chore: import upstream snapshot with attribution

vllm/distributed/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from .communication_op import *
+from .parallel_state import *
+from .utils import *

vllm/distributed/communication_op.py

@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any
+
+import torch
+import torch.distributed
+
+from .parallel_state import get_tp_group
+
+
+def tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor:
+    """All-reduce the input tensor across model parallel group."""
+    return get_tp_group().all_reduce(input_)
+
+
+def tensor_model_parallel_all_gather(
+    input_: torch.Tensor, dim: int = -1
+) -> torch.Tensor:
+    """All-gather the input tensor across model parallel group."""
+    return get_tp_group().all_gather(input_, dim)
+
+
+def tensor_model_parallel_reduce_scatter(
+    input_: torch.Tensor, dim: int = -1
+) -> torch.Tensor:
+    """Reduce-Scatter the input tensor across model parallel group."""
+    return get_tp_group().reduce_scatter(input_, dim)
+
+
+def tensor_model_parallel_gather(
+    input_: torch.Tensor, dst: int = 0, dim: int = -1
+) -> torch.Tensor | None:
+    """Gather the input tensor across model parallel group."""
+    return get_tp_group().gather(input_, dst, dim)
+
+
+def broadcast_tensor_dict(
+    tensor_dict: dict[Any, torch.Tensor | Any] | None = None, src: int = 0
+):
+    if not torch.distributed.is_initialized():
+        return tensor_dict
+    return get_tp_group().broadcast_tensor_dict(tensor_dict, src)

vllm/distributed/device_communicators/aiter_custom_all_reduce.py

@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""vLLM-owned wrapper over AITER's ``CustomAllreduce``.
+
+vLLM's ``CudaCommunicator`` stores one of these as ``aiter_ar_comm`` (when
+``VLLM_ROCM_USE_AITER_CUSTOM_AR`` is set) so the plain allreduce and
+the fused allreduce+RMSNorm path share a single AITER instance with its IPC buffers.
+
+"""
+
+import torch
+from torch.distributed import ProcessGroup
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class AiterCustomAllreduce:
+    # Default IPC buffer size for AITER's CustomAllreduce.
+    MAX_SIZE: int = 8192 * 1024 * 8 * 2
+
+    @classmethod
+    def effective_max_size(cls) -> int:
+        """
+        Max input byte size eligible for AITER custom allreduce.
+        """
+        return cls.MAX_SIZE // 2
+
+    def __init__(
+        self,
+        group: ProcessGroup,
+        device: int | str | torch.device,
+        max_size: int | None = None,
+    ):
+        from aiter.dist.device_communicators.custom_all_reduce import (
+            CustomAllreduce as _AiterCustomAllreduce,
+        )
+
+        if max_size is None:
+            max_size = self.MAX_SIZE
+
+        self._impl = _AiterCustomAllreduce(group, device, max_size=max_size)
+
+    @property
+    def aiter_ca(self):
+        return self._impl
+
+    @property
+    def disabled(self) -> bool:
+        return self._impl.disabled
+
+    def should_custom_ar(self, inp: torch.Tensor) -> bool:
+        return self._impl.should_custom_ar(inp)
+
+    def custom_all_reduce(self, inp: torch.Tensor) -> torch.Tensor | None:
+        return self._impl.custom_all_reduce(inp)
+
+    def capture(self):
+        return self._impl.capture()
+
+    def close(self) -> None:
+        self._impl.close()
+
+    @property
+    def supports_dynamic_hidden_dim(self) -> bool:
+        """Aiter's fused_allreduce_rmsnorm kernel dispatches on hidden_dim.
+        Before aiter v0.1.12 the launcher was template-specialized on HIDDEN_DIM
+        and silently no-op'd for sizes outside {512, 1024, 2048, 4096}. From v0.1.12
+        hidden_dim is a runtime argument. Older builds are detected via
+        AiterCustomAllreduce.supports_dynamic_hidden_dim; This function is used to
+        skip fusion for unsupported sizes on them.
+        Ref (old kernel): https://github.com/ROCm/aiter/blob/6a0e7b26ccf33164785531212cc2ec2cde0b9243/csrc/include/custom_all_reduce.cuh#L2590
+        """
+        return hasattr(self._impl, "_pool")
+
+    @staticmethod
+    def build_supports_per_group_quant() -> bool:
+        """True if the running AITER build exposes the per-group AR+RMS+quant
+        kernel (added in ROCm/aiter PR #2823).
+
+        The pattern registration in ``RocmAiterAllReduceFusionPass`` keys off
+        this so vLLM degrades to the AR+RMS-only fusion when run against an
+        older aiter that lacks the per-group launcher.
+        """
+        from aiter.dist.device_communicators.custom_all_reduce import (
+            CustomAllreduce as _AiterCustomAllreduce,
+        )
+
+        return hasattr(_AiterCustomAllreduce, "fused_ar_rms_per_group_quant")
+
+    # TODO(frida-andersson): drop once vLLM pins AITER >= 0.1.14 (ROCm/aiter#2823).
+    @property
+    def supports_per_group_quant(self) -> bool:
+        return self.build_supports_per_group_quant()

vllm/distributed/device_communicators/all2all.py

@@ -0,0 +1,987 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import threading
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.distributed as dist
+
+import vllm.envs as envs
+from vllm.distributed import get_dp_group, get_ep_group
+from vllm.distributed.utils import StatelessProcessGroup
+from vllm.forward_context import get_forward_context
+from vllm.logger import init_logger
+from vllm.utils.flashinfer import (
+    has_flashinfer_nvlink_one_sided,
+    has_flashinfer_nvlink_two_sided,
+)
+from vllm.utils.import_utils import has_deep_ep, has_deep_ep_v2, has_mori
+
+from .base_device_communicator import All2AllManagerBase, Cache
+
+if has_flashinfer_nvlink_two_sided():
+    from flashinfer.comm import Mapping  # type: ignore[import-not-found]
+    from flashinfer.comm.mnnvl import MnnvlConfig  # type: ignore[import-not-found]
+    from flashinfer.comm.trtllm_alltoall import (
+        MnnvlMoe,  # type: ignore[import-not-found]
+    )
+
+if has_flashinfer_nvlink_one_sided():
+    from flashinfer.comm import Mapping  # type: ignore[import-not-found]
+    from flashinfer.comm.mnnvl import MnnvlConfig  # type: ignore[import-not-found]
+    from flashinfer.comm.trtllm_moe_alltoall import (
+        MoeAlltoAll,  # type: ignore[import-not-found]
+        moe_a2a_get_workspace_size_per_rank,
+    )
+
+
+logger = init_logger(__name__)
+
+
+class AgRsAll2AllManager(All2AllManagerBase):
+    """
+    An implementation of all2all communication based on
+    all-gather (dispatch) and reduce-scatter (combine).
+    """
+
+    def __init__(self, cpu_group, tcp_store_group=None):
+        super().__init__(cpu_group, tcp_store_group)
+
+    def dispatch_router_logits(
+        self,
+        hidden_states: torch.Tensor,
+        router_logits: torch.Tensor,
+        is_sequence_parallel: bool = False,
+        extra_tensors: list[torch.Tensor] | None = None,
+    ) -> (
+        tuple[torch.Tensor, torch.Tensor]
+        | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
+    ):
+        """
+        Gather hidden_states and router_logits from all dp ranks.
+        """
+        dp_metadata = get_forward_context().dp_metadata
+        assert dp_metadata is not None
+        sizes = dp_metadata.get_chunk_sizes_across_dp_rank()
+        assert sizes is not None
+        dist_group = get_ep_group() if is_sequence_parallel else get_dp_group()
+        assert sizes[dist_group.rank_in_group] == hidden_states.shape[0]
+
+        tensors_to_gather = [hidden_states, router_logits]
+        if extra_tensors is not None:
+            tensors_to_gather.extend(extra_tensors)
+
+        gathered_tensors = dist_group.all_gatherv(
+            tensors_to_gather,
+            dim=0,
+            sizes=sizes,
+        )
+
+        if extra_tensors is not None:
+            return (gathered_tensors[0], gathered_tensors[1], gathered_tensors[2:])
+        return gathered_tensors[0], gathered_tensors[1]
+
+    def dispatch(
+        self,
+        hidden_states: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        is_sequence_parallel: bool = False,
+        extra_tensors: list[torch.Tensor] | None = None,
+    ) -> (
+        tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+        | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
+    ):
+        """
+        Gather hidden_states and router_logits from all dp ranks.
+        """
+        dp_metadata = get_forward_context().dp_metadata
+        assert dp_metadata is not None
+        sizes = dp_metadata.get_chunk_sizes_across_dp_rank()
+        assert sizes is not None
+        dist_group = get_ep_group() if is_sequence_parallel else get_dp_group()
+        assert sizes[dist_group.rank_in_group] == hidden_states.shape[0]
+
+        tensors_to_gather = [hidden_states, topk_weights, topk_ids]
+        if extra_tensors is not None:
+            tensors_to_gather.extend(extra_tensors)
+
+        gathered_tensors = dist_group.all_gatherv(
+            tensors_to_gather,
+            dim=0,
+            sizes=sizes,
+        )
+
+        hidden_states = gathered_tensors[0]
+        topk_weights = gathered_tensors[1]
+        topk_ids = gathered_tensors[2]
+
+        if extra_tensors is None:
+            return hidden_states, topk_weights, topk_ids
+
+        return hidden_states, topk_weights, topk_ids, gathered_tensors[3:]
+
+    def combine(
+        self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
+    ) -> torch.Tensor:
+        """
+        Reduce-scatter hidden_states across all dp ranks.
+        """
+        dp_metadata = get_forward_context().dp_metadata
+        assert dp_metadata is not None
+        sizes = dp_metadata.get_chunk_sizes_across_dp_rank()
+        assert sizes is not None
+
+        dist_group = get_ep_group() if is_sequence_parallel else get_dp_group()
+        hidden_states = dist_group.reduce_scatterv(hidden_states, dim=0, sizes=sizes)
+        return hidden_states
+
+    def destroy(self):
+        pass
+
+
+class DeepEPAll2AllManagerBase(All2AllManagerBase):
+    """
+    All2All communication based on DeepEP High-Throughput kernels.
+    """
+
+    def __init__(self, cpu_group, tcp_store_group=None):
+        assert has_deep_ep(), (
+            "DeepEP kernels not found. Please follow https://github.com/vllm-project/vllm/blob/main/tools/ep_kernels/README.md"
+            " to install DeepEP kernels."
+        )  # noqa
+        super().__init__(cpu_group, tcp_store_group)
+        self.handle_cache = Cache()
+
+        # This is the DeepEP default. Stick to it till we can establish
+        # reasonable defaults based on profiling.
+        self.num_sms = 20
+
+    def get_handle(self, kwargs):
+        raise NotImplementedError
+
+    def dispatch_router_logits(
+        self,
+        hidden_states: torch.Tensor,
+        router_logits: torch.Tensor,
+        is_sequence_parallel: bool = False,
+        extra_tensors: list[torch.Tensor] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        raise NotImplementedError
+
+    def dispatch(
+        self,
+        hidden_states: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        is_sequence_parallel: bool = False,
+        extra_tensors: list[torch.Tensor] | None = None,
+    ) -> (
+        tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+        | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
+    ):
+        raise NotImplementedError
+
+    def combine(
+        self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
+    ) -> torch.Tensor:
+        raise NotImplementedError
+
+    def destroy(self):
+        with self.handle_cache._lock:
+            for _, handle in self.handle_cache._cache.items():
+                handle.destroy()
+            self.handle_cache._cache.clear()
+
+
+class DeepEPHTAll2AllManager(DeepEPAll2AllManagerBase):
+    """
+    All2All communication based on DeepEP High-Throughput kernels.
+    """
+
+    def __init__(self, cpu_group, tcp_store_group=None):
+        super().__init__(cpu_group, tcp_store_group)
+
+    def _make_all2all_kwargs(self) -> dict[Any, Any]:
+        # Defaults for internode and intranode are taken from DeepEP tests.
+        num_nvl_bytes = envs.VLLM_DEEPEP_BUFFER_SIZE_MB * 1024 * 1024
+        num_rdma_bytes = None
+        num_qps_per_rank = None
+
+        if self.internode and not envs.VLLM_DEEPEP_HIGH_THROUGHPUT_FORCE_INTRA_NODE:
+            num_rdma_bytes = envs.VLLM_DEEPEP_BUFFER_SIZE_MB * 1024 * 1024
+            num_qps_per_rank = self.num_sms // 2
+        else:
+            num_rdma_bytes = 0
+            num_qps_per_rank = 1
+
+        assert num_rdma_bytes is not None
+        assert num_qps_per_rank is not None
+        # TODO: remove platform-specific logic
+        # once ROCm DeepEP is updated with the latest APIs.
+        kwargs = dict(
+            group=self.cpu_group,
+            num_nvl_bytes=num_nvl_bytes,
+            num_rdma_bytes=num_rdma_bytes,
+            low_latency_mode=False,
+            num_qps_per_rank=num_qps_per_rank,
+            explicitly_destroy=True,
+        )
+        return kwargs
+
+    def get_handle(self, kwargs):
+        assert len(kwargs) == 0, (
+            "DeepEPHTAll2AllManager expects no arguments. All the required "
+            "args are computed in the Manager itself."
+        )
+
+        import deep_ep  # type: ignore[import-not-found]
+
+        buffer_kwargs = self._make_all2all_kwargs()
+        logger.debug("DeepEP all2all args %s", buffer_kwargs)
+        handle: deep_ep.Buffer = self.handle_cache.get_or_create(
+            buffer_kwargs, deep_ep.Buffer
+        )
+        return handle
+
+    def set_num_sms(self, num_sms: int):
+        import deep_ep  # type: ignore[import-not-found]
+
+        # Right now the buffers are sized for only what the kernels were
+        # created with. So we can only reduce the number of SMS used
+        # but not increase it.
+        if num_sms > self.num_sms:
+            num_sms = self.num_sms
+        deep_ep.Buffer.set_num_sms(num_sms)
+
+
+class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):
+    """
+    All2All communication based on DeepEP Low-Latency kernels.
+    """
+
+    _buffer: Any = None
+    _mask: torch.Tensor | None = None
+    _last_mask: torch.Tensor | None = None
+
+    def __init__(self, cpu_group, tcp_store_group=None):
+        super().__init__(cpu_group, tcp_store_group)
+        self.support_fault_tolerance = False  # TODO: set to True when FT is supported.
+
+    def _make_all2all_kwargs(
+        self,
+        max_num_tokens_per_dp_rank: int,
+        token_hidden_size: int,
+        num_ep_ranks: int,
+        num_global_experts: int,
+        num_local_experts: int,
+    ) -> dict[Any, Any]:
+        """
+        max_num_tokens_per_dp_rank : the maximum number of tokens a DP rank
+          can dispatch all the ranks must hold the same value.
+        token_hidden_size: the hidden dimension of each token.
+        num_ep_ranks: the number of EP group ranks.
+        num_global_experts: Number of experts in the model.
+        num_local_experts: Number of experts in an EP rank.
+        """
+        import deep_ep  # type: ignore[import-not-found]
+
+        # Defaults for internode and intranode are taken from DeepEP tests.
+        num_nvl_bytes = envs.VLLM_DEEPEP_BUFFER_SIZE_MB * 1024 * 1024
+        num_qps_per_rank = num_local_experts
+        num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint(
+            num_max_dispatch_tokens_per_rank=max_num_tokens_per_dp_rank,
+            hidden=token_hidden_size,
+            num_ranks=num_ep_ranks,
+            num_experts=num_global_experts,
+        )
+
+        assert num_rdma_bytes is not None
+        # TODO: remove platform-specific logic
+        # once ROCm DeepEP is updated with the latest APIs.
+        kwargs = dict(
+            group=self.cpu_group,
+            num_nvl_bytes=num_nvl_bytes,
+            num_rdma_bytes=num_rdma_bytes,
+            low_latency_mode=True,
+            num_qps_per_rank=num_qps_per_rank,
+            allow_nvlink_for_low_latency_mode=True,
+            allow_mnnvl=envs.VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL,
+            explicitly_destroy=True,
+            enable_shrink=self.support_fault_tolerance,
+        )
+        return kwargs
+
+    def get_handle(self, kwargs):
+        """
+        The kwargs for DeepEPLLAll2AllManager is dictated by
+        _make_all2all_kwargs.
+        """
+        import deep_ep  # type: ignore[import-not-found]
+
+        buffer_kwargs = self._make_all2all_kwargs(**kwargs)
+        logger.debug("DeepEP all2all args %s", buffer_kwargs)
+        handle: deep_ep.Buffer = self.handle_cache.get_or_create(
+            buffer_kwargs, deep_ep.Buffer
+        )
+        DeepEPLLAll2AllManager._buffer = handle
+        return handle
+
+    # DeepEP LL uses RDMA so no SMs are used for communication
+    def max_sms_used(self) -> int | None:
+        return 0
+
+    def query_active_mask(self) -> torch.Tensor:
+        buf = DeepEPLLAll2AllManager._buffer
+        assert buf is not None
+        if DeepEPLLAll2AllManager._mask is None:
+            DeepEPLLAll2AllManager._mask = torch.zeros(
+                self.world_size, device="cuda", dtype=torch.int32
+            )
+        buf.low_latency_query_mask_buffer(DeepEPLLAll2AllManager._mask)
+        return DeepEPLLAll2AllManager._mask
+
+    def query_fault(self) -> torch.Tensor:
+        current = self.query_active_mask()
+        if DeepEPLLAll2AllManager._last_mask is None:
+            DeepEPLLAll2AllManager._last_mask = torch.zeros_like(current)
+        has_fault = (current != DeepEPLLAll2AllManager._last_mask).any()
+        return has_fault
+
+
+@dataclass
+class _NixlEPBufferState:
+    buffer: Any
+    connected_ep_size: int
+    active_ep_size: int
+
+
+class NixlEPAll2AllManager(All2AllManagerBase):
+    """
+    All2All communication based on NIXL EP kernels.
+    This backend supports elastic EP with dynamic rank connection/disconnection.
+    """
+
+    _buffer: _NixlEPBufferState | None = None
+    _lock = threading.RLock()
+    _mask: torch.Tensor | None = None
+    _last_mask: torch.Tensor | None = None
+
+    def __init__(self, cpu_group, tcp_store_group=None):
+        if tcp_store_group is None:
+            tcp_store_group = StatelessProcessGroup(
+                rank=cpu_group.rank(),
+                world_size=cpu_group.size(),
+                store=dist.PrefixStore("nixl_ep", cpu_group.get_group_store()),
+            )
+        super().__init__(cpu_group, tcp_store_group)
+        self.support_fault_tolerance = True
+
+        self.max_num_ep_ranks = envs.VLLM_NIXL_EP_MAX_NUM_RANKS
+
+    def _init_buffer(
+        self,
+        max_num_tokens_per_dp_rank: int,
+        token_hidden_size: int,
+        num_experts_per_rank: int,
+    ) -> None:
+        from nixl_ep import Buffer  # type: ignore[import-not-found]
+
+        max_num_global_experts = self.max_num_ep_ranks * num_experts_per_rank
+        num_rdma_bytes = Buffer.get_rdma_size_hint(
+            num_max_dispatch_tokens_per_rank=max_num_tokens_per_dp_rank,
+            hidden=token_hidden_size,
+            num_ranks=self.max_num_ep_ranks,
+            num_experts=max_num_global_experts,
+        )
+        assert NixlEPAll2AllManager._buffer is None, (
+            "NIXL EP buffer already initialized"
+        )
+        buffer = Buffer(
+            rank=self.rank,
+            tcp_store_group=self.tcp_store_group.store,
+        )
+        buffer.update_memory_buffers(
+            num_ranks=self.max_num_ep_ranks,
+            num_experts_per_rank=num_experts_per_rank,
+            num_rdma_bytes=num_rdma_bytes,
+        )
+        ranks_to_connect = list(range(self.world_size))
+        buffer.connect_ranks(ranks_to_connect)
+        NixlEPAll2AllManager._buffer = _NixlEPBufferState(
+            buffer=buffer,
+            connected_ep_size=self.world_size,
+            active_ep_size=self.world_size,
+        )
+
+    def _connect_to_ep_size(self, ep_size: int, *, make_active: bool) -> None:
+        assert NixlEPAll2AllManager._buffer is not None
+        state = NixlEPAll2AllManager._buffer
+        if ep_size <= state.connected_ep_size:
+            return
+
+        state.buffer.set_tcp_store_group(self.tcp_store_group.store)
+        ranks_to_connect = list(range(state.connected_ep_size, ep_size))
+        state.buffer.connect_ranks(ranks_to_connect, activate=make_active)
+        state.connected_ep_size = ep_size
+        if make_active:
+            state.active_ep_size = ep_size
+
+    def _disconnect_to_ep_size(self, ep_size: int) -> None:
+        assert NixlEPAll2AllManager._buffer is not None
+        state = NixlEPAll2AllManager._buffer
+        if ep_size >= state.connected_ep_size:
+            return
+
+        state.buffer.set_tcp_store_group(self.tcp_store_group.store)
+        ranks_to_disconnect = list(range(ep_size, state.connected_ep_size))
+        state.buffer.disconnect_ranks(ranks_to_disconnect)
+        state.connected_ep_size = ep_size
+        state.active_ep_size = min(state.active_ep_size, ep_size)
+
+    def _unmask_connected_ranks(self, target_ep_size: int) -> None:
+        assert NixlEPAll2AllManager._buffer is not None
+        state = NixlEPAll2AllManager._buffer
+        state.buffer.set_tcp_store_group(self.tcp_store_group.store)
+        if target_ep_size <= state.active_ep_size:
+            return
+        assert state.connected_ep_size >= target_ep_size
+
+        for rank in range(state.active_ep_size, target_ep_size):
+            state.buffer.update_mask_buffer(rank, mask=False)
+        state.active_ep_size = target_ep_size
+
+    def _stage_ep_size(self) -> None:
+        assert NixlEPAll2AllManager._buffer is not None
+        state = NixlEPAll2AllManager._buffer
+        target_ep_size = self.world_size
+
+        # Scale-up can safely connect standby ranks while leaving them masked.
+        # Scale-down must not disconnect active ranks until commit.
+        if target_ep_size > state.connected_ep_size:
+            self._connect_to_ep_size(target_ep_size, make_active=False)
+
+    def commit_staged_state(self) -> None:
+        """Commit staged NIXL EP state to the active communication set."""
+        with NixlEPAll2AllManager._lock:
+            assert NixlEPAll2AllManager._buffer is not None
+            state = NixlEPAll2AllManager._buffer
+            target_ep_size = self.world_size
+
+            if target_ep_size < state.connected_ep_size:
+                self._disconnect_to_ep_size(target_ep_size)
+            elif target_ep_size > state.connected_ep_size:
+                self._connect_to_ep_size(target_ep_size, make_active=True)
+
+            self._unmask_connected_ranks(target_ep_size)
+
+    def _ensure_ep_size(self, *, stage: bool) -> None:
+        if stage:
+            self._stage_ep_size()
+        else:
+            self.commit_staged_state()
+
+    def get_handle(self, kwargs):
+        with NixlEPAll2AllManager._lock:
+            stage = bool(kwargs.get("stage", False))
+            state = NixlEPAll2AllManager._buffer
+            if state is None:
+                assert not stage, (
+                    "NIXL EP staged initialization requires an existing buffer"
+                )
+                max_num_tokens_per_dp_rank = kwargs["max_num_tokens_per_dp_rank"]
+                num_experts_per_rank = (
+                    kwargs["num_global_experts"] // kwargs["num_ep_ranks"]
+                )
+                self._init_buffer(
+                    max_num_tokens_per_dp_rank=max_num_tokens_per_dp_rank,
+                    token_hidden_size=kwargs["token_hidden_size"],
+                    num_experts_per_rank=num_experts_per_rank,
+                )
+            else:
+                self._ensure_ep_size(stage=stage)
+
+            assert NixlEPAll2AllManager._buffer is not None
+            handle = NixlEPAll2AllManager._buffer.buffer
+            return handle
+
+    def dispatch(
+        self,
+        hidden_states: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        is_sequence_parallel: bool = False,
+        extra_tensors: list[torch.Tensor] | None = None,
+    ) -> (
+        tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+        | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
+    ):
+        raise NotImplementedError
+
+    def combine(
+        self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
+    ) -> torch.Tensor:
+        raise NotImplementedError
+
+    def destroy(self):
+        # NOTE(yongji): NIXLEPAll2AllManager instance is recreated during
+        # scale-up/down, so we cannot destroy the persistent buffer here.
+        assert NixlEPAll2AllManager._buffer is not None
+        buffer = NixlEPAll2AllManager._buffer.buffer
+        buffer.set_tcp_store_group(None)
+
+    # NIXL EP uses RDMA so no SMs are used for communication
+    def max_sms_used(self) -> int | None:
+        return 0
+
+    def query_active_mask(self) -> torch.Tensor:
+        state = NixlEPAll2AllManager._buffer
+        assert state is not None
+        if NixlEPAll2AllManager._mask is None:
+            NixlEPAll2AllManager._mask = torch.zeros(
+                self.max_num_ep_ranks, device="cuda", dtype=torch.int32
+            )
+        state.buffer.query_mask_buffer(NixlEPAll2AllManager._mask)
+        return NixlEPAll2AllManager._mask[: state.active_ep_size]
+
+    def query_fault(self) -> torch.Tensor:
+        current = self.query_active_mask()
+        last = NixlEPAll2AllManager._last_mask
+        if last is None or last.shape != current.shape:
+            NixlEPAll2AllManager._last_mask = torch.zeros_like(current)
+            last = NixlEPAll2AllManager._last_mask
+        has_fault = (current != last).any()
+        return has_fault
+
+
+class FlashInferNVLinkTwoSidedManager(All2AllManagerBase):
+    """
+    All2All communication based on flashinfer all2allv/two-sided NVLink kernels.
+    """
+
+    # This type lint could be removed after all of the work in
+    # https://github.com/vllm-project/vllm/issues/26533 done.
+    rank: int
+    world_size: int
+
+    def __init__(self, cpu_group, tcp_store_group=None):
+        assert has_flashinfer_nvlink_two_sided(), (
+            "flashinfer all2all module not found. Please install/check flashinfer"
+        )  # noqa
+        super().__init__(cpu_group, tcp_store_group)
+        logger.debug(
+            "Initialize for flashinfer All2All rank=%d, world size=%d",
+            self.rank,
+            self.world_size,
+        )
+        self.initialized = False
+        self.alltoall_info = None
+
+    def initialize(
+        self,
+        world_size: int,
+        rank: int,
+        gpus_per_node: int,
+    ):
+        """Initialize workspace"""
+        if self.initialized:
+            return
+
+        self.cleanup()
+        logger.debug("making map: rank=%d, world size=%d", rank, world_size)
+        self.mapping = Mapping(
+            world_size,
+            rank,
+            gpus_per_node,
+            tp_size=world_size,
+        )
+
+        from vllm.distributed.device_communicators.mnnvl_compat import (
+            CustomCommunicator,
+        )
+
+        # MNNVL workspace is allocated per rank in the comm_backend's group; the
+        # flashinfer kernel asserts workspace.size(0) == moe_ep_size, so the backend
+        # must span the EP group (= DP*PCP*TP), not the DP group.
+        ep_config = MnnvlConfig(
+            comm_backend=CustomCommunicator(self.cpu_group),
+            fabric_page_size=1 << 29,  # 512MB
+            allocation_granularity=0,  # Auto-detect
+        )
+
+        self.workspace_tensor = MnnvlMoe.get_moe_workspaces(self.mapping, ep_config)
+        self.prepare_workspace_tensor = MnnvlMoe.get_moe_prepare_workspace(
+            self.mapping, ep_config
+        )
+
+        self.world_size = world_size
+        self.rank = rank
+        self.gpus_per_node = gpus_per_node
+        self.initialized = True
+
+        logger.info(
+            "FlashInfer All2All initialized for rank %s, size %s", rank, world_size
+        )
+
+    def ensure_alltoall_workspace_initialized(self):
+        """Ensure workspace is initialized"""
+        if not has_flashinfer_nvlink_two_sided():
+            return False
+
+        if self.world_size <= 1:
+            return False
+
+        if not self.initialized:
+            self.initialize(
+                world_size=self.world_size,
+                rank=self.rank,
+                gpus_per_node=torch.accelerator.device_count,
+            )
+        return self.initialized
+
+    def get_handle(self, kwargs):
+        return self
+
+    def cleanup(self):
+        """Clean up workspace"""
+        if (
+            self.initialized
+            and self.workspace_tensor is not None
+            and self.prepare_workspace_tensor is not None
+        ):
+            try:
+                del self.workspace_tensor
+                del self.prepare_workspace_tensor
+            except Exception as e:
+                logger.warning("Failed to cleanup FlashInfer workspace: %s", e)
+            finally:
+                self.workspace_tensor = None
+                self.prepare_workspace_tensor = None
+                self.mapping = None
+                self.initialized = False
+
+
+class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
+    """
+    All2All communication based on FlashInfer's MoeAlltoAll/One-sided NVLink kernel.
+    This is a newer kernel from trtllm that should perform better than the kernel
+    used by flashinfer_nvlink_two_sided.
+    """
+
+    rank: int
+    world_size: int
+
+    def __init__(self, cpu_group):
+        assert has_flashinfer_nvlink_one_sided(), (
+            "flashinfer trtllm_moe_alltoall module not found. "
+            "Please install/check flashinfer"
+        )
+        super().__init__(cpu_group)
+        logger.debug(
+            "Initialize FlashInfer One-sided NVLink rank=%d, world size=%d",
+            self.rank,
+            self.world_size,
+        )
+        self.initialized = False
+        self.moe_alltoall: MoeAlltoAll | None = None
+        self.mapping = None
+        self.workspace_size = 0
+        self.max_num_tokens = 0
+        self.top_k = 0
+        self.num_experts = 0
+
+    def initialize(
+        self,
+        max_num_tokens: int,
+        top_k: int,
+        num_experts: int,
+        hidden_size: int,
+        dispatch_dtype_bytes_per_elem: int = 0,
+        dispatch_scale_bytes_per_token: int = 0,
+    ):
+        """Initialize (or grow) the MoeAlltoAll workspace."""
+        if dispatch_dtype_bytes_per_elem == 0:
+            hidden_bytes = hidden_size // 2
+        else:
+            hidden_bytes = hidden_size * dispatch_dtype_bytes_per_elem
+        total_dispatch_payload_size_per_token = (
+            hidden_bytes
+            + dispatch_scale_bytes_per_token
+            + top_k * 4  # int32 topks ids
+            + top_k * 4  # float32 topk weights
+        )
+        combine_payload_size_per_token = hidden_size * 2  # bf16 hidden states
+        needed_workspace_size = moe_a2a_get_workspace_size_per_rank(
+            ep_size=self.world_size,
+            max_num_tokens=max_num_tokens,
+            total_dispatch_payload_size_per_token=total_dispatch_payload_size_per_token,
+            combine_payload_size_per_token=combine_payload_size_per_token,
+        )
+        # workspace_size and max_num_tokens are kernel-side max-bounds, so
+        # heterogeneous MoE layers (e.g. NVFP4 base + bf16 MTP head) only
+        # need the shared workspace grown to the union. top_k and num_experts
+        # must match across layers: top_k is a strict-equality assert at
+        # dispatch (FlashInfer csrc/trtllm_moe_alltoall.cu), and num_experts
+        # feeds the expert-to-rank routing math, so any mismatch would crash
+        # or silently corrupt routing. All ranks see the same MoE layers in
+        # the same order with identical shapes, so the skip / rebuild
+        # branches are taken consistently across ranks.
+        if self.initialized:
+            assert top_k == self.top_k, (
+                "FlashInfer one-sided MoeAlltoAll does not support "
+                f"heterogeneous top_k across MoE layers (got {top_k}, "
+                f"was built with {self.top_k})"
+            )
+            assert num_experts == self.num_experts, (
+                "FlashInfer one-sided MoeAlltoAll does not support "
+                f"heterogeneous num_experts across MoE layers (got "
+                f"{num_experts}, was built with {self.num_experts})"
+            )
+            if (
+                needed_workspace_size <= self.workspace_size
+                and max_num_tokens <= self.max_num_tokens
+            ):
+                return
+
+        self.workspace_size = max(self.workspace_size, needed_workspace_size)
+        self.max_num_tokens = max(self.max_num_tokens, max_num_tokens)
+        self.top_k = top_k
+        self.num_experts = num_experts
+
+        self.cleanup()
+        from vllm.platforms.interface import get_assigned_physical_gpu_ids
+
+        assigned_physical_gpu_ids = get_assigned_physical_gpu_ids()
+        gpus_per_node = (
+            len(assigned_physical_gpu_ids)
+            if assigned_physical_gpu_ids is not None
+            else torch.accelerator.device_count()
+        )
+        logger.debug(
+            "Making One-sided NVLink mapping: rank=%d, world size=%d",
+            self.rank,
+            self.world_size,
+        )
+        self.mapping = Mapping(
+            self.world_size,
+            self.rank,
+            gpus_per_node,
+            tp_size=self.world_size,
+            moe_ep_size=self.world_size,
+        )
+
+        from vllm.distributed.device_communicators.mnnvl_compat import (
+            CustomCommunicator,
+        )
+
+        # MNNVL workspace is allocated per rank in the comm_backend's group; the
+        # flashinfer kernel asserts workspace.size(0) == moe_ep_size, so the backend
+        # must span the EP group (= DP*PCP*TP), not the DP group.
+        ep_config = MnnvlConfig(
+            comm_backend=CustomCommunicator(self.cpu_group),
+        )
+
+        self.moe_alltoall = MoeAlltoAll(
+            mapping=self.mapping,
+            max_num_tokens=self.max_num_tokens,
+            top_k=self.top_k,
+            num_experts=self.num_experts,
+            workspace_size_per_rank=self.workspace_size,
+            mnnvl_config=ep_config,
+        )
+
+        self.gpus_per_node = gpus_per_node
+        self.initialized = True
+
+        logger.info(
+            "FlashInfer One-sided NVLink initialized for rank %s, size %s",
+            self.rank,
+            self.world_size,
+        )
+        # Scope barrier to the EP group: with PP, different EP groups can
+        # rebuild a different number of times if their MoE layers have
+        # different shape sequences, so a world-level barrier would deadlock.
+        dist.barrier(group=self.cpu_group)
+
+    def get_handle(self, kwargs):
+        return self
+
+    def cleanup(self):
+        """Clean up resources."""
+        if self.initialized and self.moe_alltoall is not None:
+            try:
+                del self.moe_alltoall
+            except Exception as e:
+                logger.warning(
+                    "Failed to cleanup FlashInfer One-sided NVLink workspace: %s", e
+                )
+            finally:
+                self.moe_alltoall = None
+                self.mapping = None
+                self.initialized = False
+
+
+class MoriAll2AllManager(All2AllManagerBase):
+    def __init__(self, cpu_group, all2all_backend: str):
+        assert has_mori(), (
+            "MoRI kernels not found. Please follow https://github.com/ROCm/mori/blob/main/README.md"
+            " to install MoRI kernels."
+        )  # noqa
+        assert all2all_backend in (
+            "mori_high_throughput",
+            "mori_low_latency",
+        ), f"unsupported MoRI all2all backend: {all2all_backend!r}"
+        import mori
+
+        super().__init__(cpu_group)
+        self._all2all_backend = all2all_backend
+        self.handle_cache = Cache()
+
+        torch._C._distributed_c10d._register_process_group("mori", cpu_group)
+        mori.shmem.shmem_torch_process_group_init("mori")
+
+    def _make_all2all_kwargs(
+        self,
+        rank: int,
+        num_ep_ranks: int,
+        input_dtype: torch.dtype,
+        quant_dtype: torch.dtype,
+        token_hidden_size: int,
+        scale_dim: int,
+        scale_type_size: int,
+        max_num_tokens_per_dp_rank: int,
+        num_local_experts: int,
+        num_experts_per_token: int,
+    ):
+        import mori  # type: ignore[import-not-found]
+
+        from vllm.platforms.rocm import on_gfx942, on_gfx950
+
+        assert on_gfx942() or on_gfx950(), (
+            "mori currently only support arch gfx942 and gfx950"
+        )
+
+        if not self.internode:
+            # single node
+            kernel_type = mori.ops.EpDispatchCombineKernelType.IntraNode
+            rdma_block_num = 0
+            warp_num_per_block = 16
+            block_num = 80
+        else:
+            # Multi-node: kernel follows --all2all-backend (mirrors deepep_* split).
+            # mori_low_latency → InterNodeV1LL; mori_high_throughput → V1.
+            if self._all2all_backend == "mori_low_latency":
+                kernel_type = mori.ops.EpDispatchCombineKernelType.InterNodeV1LL
+            else:
+                kernel_type = mori.ops.EpDispatchCombineKernelType.InterNodeV1
+            if on_gfx942():
+                warp_num_per_block = 16
+                block_num = 32
+                rdma_block_num = 16
+            elif on_gfx950():
+                warp_num_per_block = 8
+                block_num = 64
+                rdma_block_num = 32
+            else:
+                raise NotImplementedError(
+                    "mori currently only support arch gfx942 and gfx950"
+                )
+
+        return dict(
+            rank=rank,
+            world_size=num_ep_ranks,
+            data_type=quant_dtype,
+            hidden_dim=token_hidden_size,
+            scale_dim=scale_dim,
+            scale_type_size=scale_type_size,
+            max_token_type_size=input_dtype.itemsize,
+            max_num_inp_token_per_rank=max_num_tokens_per_dp_rank,
+            num_experts_per_rank=num_local_experts,
+            num_experts_per_token=num_experts_per_token,
+            warp_num_per_block=warp_num_per_block,
+            block_num=block_num,
+            kernel_type=kernel_type,
+            rdma_block_num=rdma_block_num,
+            gpu_per_node=min(8, num_ep_ranks),
+        )
+
+    def _make_handle(self, **kwargs):
+        import mori  # type: ignore[import-not-found]
+
+        mori_config = mori.ops.EpDispatchCombineConfig(**kwargs)
+        handle = mori.ops.EpDispatchCombineOp(mori_config)
+        return handle
+
+    def get_handle(self, kwargs):
+        import mori  # type: ignore[import-not-found]
+
+        mori_kwargs = self._make_all2all_kwargs(**kwargs)
+        logger.debug("MoRI all2all args %s", mori_kwargs)
+        handle: mori.ops.EpDispatchCombineOp = self.handle_cache.get_or_create(
+            mori_kwargs, self._make_handle
+        )
+        return handle
+
+
+class DeepEPV2All2AllManager(All2AllManagerBase):
+    """
+    All2All communication based on DeepEP v2 ElasticBuffer (unified API).
+    Uses NCCL Gin backend with analytical SM calculation.
+    """
+
+    def __init__(self, cpu_group, tcp_store_group=None, device_group=None):
+        assert has_deep_ep_v2(), (
+            "DeepEP v2 (ElasticBuffer) not available. Requires DeepEP >= 2.0 "
+            "(https://github.com/deepseek-ai/DeepEP) and NCCL >= 2.30.4."
+        )
+        super().__init__(cpu_group, tcp_store_group)
+        self._device_group = device_group
+        self.handle_cache = Cache()
+        self._num_sms: int | None = None
+
+    def _make_all2all_kwargs(
+        self,
+        num_max_tokens_per_rank: int,
+        hidden: int,
+        num_topk: int,
+        use_fp8_dispatch: bool,
+    ) -> dict:
+        return dict(
+            group=self._device_group
+            if self._device_group is not None
+            else self.cpu_group,
+            num_max_tokens_per_rank=num_max_tokens_per_rank,
+            hidden=hidden,
+            num_topk=num_topk,
+            use_fp8_dispatch=use_fp8_dispatch,
+            allow_hybrid_mode=envs.VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE,
+            prefer_overlap_with_compute=envs.VLLM_DEEPEP_V2_PREFER_OVERLAP,
+            allow_multiple_reduction=(envs.VLLM_DEEPEP_V2_ALLOW_MULTIPLE_REDUCTION),
+            explicitly_destroy=True,
+        )
+
+    def get_handle(self, kwargs):
+        import deep_ep  # type: ignore[import-not-found]
+
+        num_experts = kwargs.pop("num_experts", 256)
+        buffer_kwargs = self._make_all2all_kwargs(**kwargs)
+        logger.debug("DeepEP v2 all2all args %s", buffer_kwargs)
+        handle: deep_ep.ElasticBuffer = self.handle_cache.get_or_create(
+            buffer_kwargs, deep_ep.ElasticBuffer
+        )
+        if self._num_sms is None:
+            self._num_sms = handle.get_theoretical_num_sms(
+                num_experts=num_experts,
+                num_topk=kwargs["num_topk"],
+            )
+        return handle
+
+    def max_sms_used(self) -> int | None:
+        return self._num_sms
+
+    def destroy(self):
+        with self.handle_cache._lock:
+            for _, handle in self.handle_cache._cache.items():
+                handle.destroy()
+            self.handle_cache._cache.clear()

vllm/distributed/device_communicators/all_reduce_utils.py

@@ -0,0 +1,421 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import ctypes
+import json
+import os
+import pickle
+import subprocess
+import sys
+import tempfile
+from collections.abc import Sequence
+from itertools import product
+from typing import Any
+
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+import vllm.envs as envs
+from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+from vllm.utils.system_utils import update_environment_variables
+
+logger = init_logger(__name__)
+
+KiB = 1024
+MiB = 1024 * 1024
+# Max size for each world size in case symmetric memory is available
+# For different SM architectures
+CUSTOM_ALL_REDUCE_MAX_SIZES = {
+    "9.0": {
+        2: 64 * MiB,  # 64 MB
+        4: 32 * MiB,  # 32 MB
+        6: MiB // 2,  # 512 KB
+        8: MiB // 4,  # 256 KB
+    },
+    "10.0": {
+        2: 2 * MiB,  # 2 MB
+        4: 2 * MiB,  # 2 MB
+        6: 1 * MiB,  # 1 MB
+        8: 1 * MiB,  # 1 MB
+    },
+    "10.3": {
+        2: 4 * MiB,  # 4 MB
+        4: 4 * MiB,  # 4 MB
+        6: 8 * MiB,  # 8 MB
+        8: 4 * MiB,  # 4 MB
+    },
+}
+
+SYMM_MEM_ALL_REDUCE_MAX_SIZES = {
+    "9.0": {
+        2: 64 * MiB,  # 64 MB
+        4: 32 * MiB,  # 32 MB
+        6: 64 * MiB,  # 64 MB
+        8: 64 * MiB,  # 64 MB
+    },
+    "10.0": {
+        2: 8 * MiB,  # 8 MB
+        4: 32 * MiB,  # 32 MB
+        6: 128 * MiB,  # 128 MB
+        8: 128 * MiB,  # 128 MB
+    },
+    "10.3": {
+        2: 4 * MiB,  # 4 MB
+        4: 32 * MiB,  # 32 MB
+        6: 32 * MiB,  # 32 MB
+        8: 64 * MiB,  # 64 MB
+    },
+}
+
+# NCCL symmetric memory allreduce configuration based on H100 and GB200 benchmarks.
+# PyNCCL-symm outperforms custom_AR for small and large tensor sizes,
+# while custom_AR wins for mid-range sizes.
+#
+# Benchmark results (8 GPUs):
+#   2K - 16K:   PyNCCL-symm wins (1.35x - 1.48x faster)
+#   32K - 64K:  custom_AR wins
+#   128K - 1G:  PyNCCL-symm wins (1.12x - 6.14x faster)
+#
+# Benchmark results (4 GPUs):
+#   2K - 16K:   PyNCCL-symm wins (1.21x - 1.30x faster)
+#   32K - 256K: custom_AR wins (1.07x - 1.35x faster)
+#   512K - 1G:  PyNCCL-symm wins (1.10x - 2.32x faster)
+#
+# The config defines ranges where custom_AR is preferred (symm_mem disabled).
+NCCL_SYMM_MEM_ALL_REDUCE_CONFIG: dict[str, Any] = {
+    "min_world_size": 4,
+    # Ranges where custom_AR outperforms NCCL symm_mem: (lower_bound, upper_bound)
+    # NCCL symm_mem will NOT be used for sizes in range: lower < size < upper
+    "custom_ar_preferred_ranges": {
+        4: (16 * KiB, 512 * KiB),  # custom_AR wins for 32K-256K
+        8: (16 * KiB, 128 * KiB),  # custom_AR wins for 32K-64K
+    },
+    "always_use_above_world_size": 8,  # Always use symm mem for world_size > 8
+}
+
+
+def should_nccl_symm_mem_allreduce(world_size: int, input_tensor: torch.Tensor) -> bool:
+    """
+    Determine if NCCL symmetric memory allreduce should be used.
+
+    Based on H100 and GB200 benchmarks, NCCL symm_mem is preferred for:
+    - Small tensors (≤16K): Lower latency than custom_AR
+    - Large tensors (≥128K for 8 GPUs, ≥512K for 4 GPUs): Better bandwidth
+
+    Custom_AR is preferred for mid-range sizes where its P2P approach
+    has lower overhead than the symm_mem copy-in/copy-out pattern.
+    """
+    from vllm.distributed.device_communicators.pynccl_allocator import (
+        is_symmetric_memory_enabled,
+    )
+
+    if envs.VLLM_BATCH_INVARIANT:
+        return False
+
+    if not is_symmetric_memory_enabled():
+        return False
+
+    if world_size < NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["min_world_size"]:
+        return False
+
+    tensor_size = input_tensor.nbytes
+    custom_ar_range = NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["custom_ar_preferred_ranges"].get(
+        world_size
+    )
+
+    if custom_ar_range is not None:
+        lower_bound, upper_bound = custom_ar_range
+        # Use symm_mem for small sizes (≤ lower_bound) and large sizes (≥ upper_bound)
+        # Use custom_AR (not symm_mem) for mid-range sizes
+        return tensor_size <= lower_bound or tensor_size >= upper_bound
+    return world_size > NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["always_use_above_world_size"]
+
+
+def should_nccl_symm_mem_ag_rs() -> bool:
+    """Check whether NCCL symmetric memory should be used for
+    AllGather / ReduceScatter collectives."""
+    from vllm.distributed.device_communicators.pynccl_allocator import (
+        is_symmetric_memory_enabled,
+    )
+
+    if envs.VLLM_BATCH_INVARIANT:
+        return False
+    return is_symmetric_memory_enabled()
+
+
+def producer(
+    batch_src: Sequence[int],
+    producer_queue,
+    consumer_queue,
+    result_queue,
+    cuda_visible_devices: str | None = None,
+):
+    if cuda_visible_devices is not None:
+        update_environment_variables({"CUDA_VISIBLE_DEVICES": cuda_visible_devices})
+
+    lib = CudaRTLibrary()
+    for i in batch_src:
+        lib.cudaSetDevice(i)
+        pointer = lib.cudaMalloc(1024)
+        lib.cudaMemset(pointer, 1, 1024)
+        lib.cudaDeviceSynchronize()
+        handle = lib.cudaIpcGetMemHandle(pointer)
+        producer_queue.put(handle)
+        open_success = consumer_queue.get()
+        if open_success:
+            # use two queues to simulate barrier
+            producer_queue.put(0)
+            consumer_queue.get()
+            # check if the memory is modified
+            host_data = (ctypes.c_char * 1024)()
+            lib.cudaMemcpy(host_data, pointer, 1024)  # type: ignore
+            for i in range(1024):
+                if ord(host_data[i]) != 2:
+                    open_success = False
+                    break
+        result_queue.put(open_success)
+        lib.cudaDeviceReset()
+
+
+def consumer(
+    batch_tgt: Sequence[int],
+    producer_queue,
+    consumer_queue,
+    result_queue,
+    cuda_visible_devices: str | None = None,
+):
+    if cuda_visible_devices is not None:
+        update_environment_variables({"CUDA_VISIBLE_DEVICES": cuda_visible_devices})
+
+    lib = CudaRTLibrary()
+    for j in batch_tgt:
+        lib.cudaSetDevice(j)
+        handle = producer_queue.get()
+        open_success = False
+        try:
+            pointer = lib.cudaIpcOpenMemHandle(handle)  # type: ignore
+            open_success = True
+        except RuntimeError:
+            # cannot error out here, because the producer process
+            # is still waiting for the response.
+            pass
+        consumer_queue.put(open_success)
+        if open_success:
+            # modify the memory
+            lib.cudaMemset(pointer, 2, 1024)
+            lib.cudaDeviceSynchronize()
+            # use two queues to simulate barrier
+            producer_queue.get()
+            consumer_queue.put(0)
+            # check if the memory is modified
+            host_data = (ctypes.c_char * 1024)()
+            lib.cudaMemcpy(host_data, pointer, 1024)  # type: ignore
+            for i in range(1024):
+                if ord(host_data[i]) != 2:
+                    open_success = False
+                    break
+        result_queue.put(open_success)
+        lib.cudaDeviceReset()
+
+
+def can_actually_p2p(
+    batch_src: Sequence[int],
+    batch_tgt: Sequence[int],
+) -> Sequence[bool]:
+    """
+    Usually, checking if P2P access is enabled can be done by
+    `torch.cuda.can_device_access_peer(src, tgt)`. However, sometimes
+    the driver might be broken, and `torch.cuda.can_device_access_peer(src, tgt)`
+    returns `True` even if P2P access is not actually possible.
+    See https://github.com/vllm-project/vllm/issues/2728 and
+    https://forums.developer.nvidia.com/t/direct-gpu-gpu-communication-does-not-seem-to-work-properly/283264/10
+    Therefore, we have to perform a real P2P access to check if it is actually
+    possible.
+
+    Note on p2p and cuda IPC:
+    Usually, one process uses one GPU:
+    GPU src --> cuda context src --> tensor src --> process src
+
+    We need to combine p2p and cuda IPC, so that:
+    GPU src --> cuda context src --> tensor src --> process src
+                                      |shared|
+    GPU tgt --> cuda context tgt --> tensor tgt --> process tgt
+    That is to say, process src creates a tensor in GPU src, passes IPC handle to
+    process tgt, and process tgt accesses the tensor in GPU tgt. Any operation on the
+    tensor in process tgt will be reflected in the tensor in process src, because
+    they are the same memory segment.
+    It is important to note that process tgt accesses the tensor in GPU tgt, not
+    GPU src. That's why we need p2p access.
+
+    The most time-consuming part is the process creation. To avoid creating
+    processes for every pair of GPUs, we use batched testing. We create two
+    processes for testing all pairs of GPUs in batch. The trick is to reset
+    the device after each test (which is not available in PyTorch).
+    """  # noqa
+    cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES
+    # pass the CUDA_VISIBLE_DEVICES to the child process
+    # to make sure they see the same set of GPUs
+
+    # make sure the processes are spawned
+    smp = mp.get_context("spawn")
+    producer_queue = smp.Queue()
+    consumer_queue = smp.Queue()
+    result_queue = smp.Queue()
+    p_src = smp.Process(
+        target=producer,
+        args=(
+            batch_src,
+            producer_queue,
+            consumer_queue,
+            result_queue,
+            cuda_visible_devices,
+        ),
+    )
+    p_tgt = smp.Process(
+        target=consumer,
+        args=(
+            batch_tgt,
+            producer_queue,
+            consumer_queue,
+            result_queue,
+            cuda_visible_devices,
+        ),
+    )
+    p_src.start()
+    p_tgt.start()
+    p_src.join()
+    p_tgt.join()
+    assert p_src.exitcode == 0 and p_tgt.exitcode == 0
+    result: list[bool] = []
+    for src, tgt in zip(batch_src, batch_tgt):
+        a = result_queue.get()
+        b = result_queue.get()
+        if a != b:
+            logger.warning(
+                "Two processes do not agree on the P2P access"
+                " status on %d -> %d, treat as disabled.",
+                src,
+                tgt,
+            )
+            result.append(False)
+        else:
+            result.append(a)
+    return result
+
+
+# why do we need this cache?
+# we are testing peer-to-peer (p2p) access between GPUs,across processes.
+# if we test it every time, it will be very slow, because we need to create
+#  N * N * 2 processes, where N is the world size. This is very slow.
+# to reduce the time, we use a cache file to store the p2p access status.
+# the cache file is generated by the master process if it does not exist.
+# then all the processes can read the cache file to check the p2p access status.
+# Note that the cache file is suffixed by the CUDA_VISIBLE_DEVICES, so that we
+#  can have different cache files for different CUDA_VISIBLE_DEVICES settings,
+#  e.g. used by different vllm engines. The device id in the cache file is a
+#  **local** device id, i.e. from 0 to num_dev-1, where num_dev is the number
+#  of visible devices in the vllm engine.
+_gpu_p2p_access_cache: dict[str, bool] | None = None
+
+
+def gpu_p2p_access_check(src: int, tgt: int) -> bool:
+    """Check if GPU src can access GPU tgt."""
+
+    # if the cache variable is already calculated,
+    # read from the cache instead of checking it again
+    global _gpu_p2p_access_cache
+    if _gpu_p2p_access_cache is not None:
+        return _gpu_p2p_access_cache[f"{src}->{tgt}"]
+
+    is_distributed = dist.is_initialized()
+
+    from vllm.platforms.interface import get_assigned_physical_gpu_ids
+
+    assigned_physical_gpu_ids = get_assigned_physical_gpu_ids()
+    if assigned_physical_gpu_ids is not None:
+        # Key by the ordered list: the cache stores directed local-index
+        # pairs, so permutations of the same set are distinct mappings.
+        cache_key = ",".join(str(i) for i in assigned_physical_gpu_ids)
+        num_dev = len(assigned_physical_gpu_ids)
+    else:
+        num_dev = current_platform.device_count()
+        cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES
+        cache_key = cuda_visible_devices or ",".join(str(i) for i in range(num_dev))
+
+    path = os.path.join(
+        envs.VLLM_CACHE_ROOT, f"gpu_p2p_access_cache_for_{cache_key}.json"
+    )
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    from vllm.distributed.parallel_state import get_world_group
+
+    if (not is_distributed or get_world_group().local_rank == 0) and (
+        not os.path.exists(path)
+    ):
+        # only the local master process (with local_rank == 0) can
+        #  enter this block to calculate the cache
+        logger.info("generating GPU P2P access cache in %s", path)
+        cache: dict[str, bool] = {}
+        # The probe subprocesses inherit this process's device-control env
+        # var, so they must be given visible ordinals, not physical IDs.
+        if assigned_physical_gpu_ids is not None:
+            ids = [
+                current_platform.logical_device_id_to_visible_device_id(local)
+                for local in range(num_dev)
+            ]
+        else:
+            ids = list(range(num_dev))
+        # batch of all pairs of GPUs
+        batch_src, batch_tgt = zip(*list(product(ids, ids)))
+        # NOTE: we use `subprocess` rather than `multiprocessing` here
+        # because the caller might not have `if __name__ == "__main__":`,
+        # in that case we cannot use spawn method in multiprocessing.
+        # However, `can_actually_p2p` requires spawn method.
+        # The fix is, we use `subprocess` to call the function,
+        # where we have `if __name__ == "__main__":` in this file.
+
+        # use a temporary file to store the result
+        # we don't use the output of the subprocess directly,
+        # because the subprocess might produce logging output
+        with tempfile.NamedTemporaryFile() as output_file:
+            input_bytes = pickle.dumps((batch_src, batch_tgt, output_file.name))
+            returned = subprocess.run(
+                [sys.executable, __file__], input=input_bytes, capture_output=True
+            )
+            # check if the subprocess is successful
+            try:
+                returned.check_returncode()
+            except Exception as e:
+                # wrap raised exception to provide more information
+                raise RuntimeError(
+                    f"Error happened when batch testing "
+                    f"peer-to-peer access from {batch_src} to {batch_tgt}:\n"
+                    f"{returned.stderr.decode()}"
+                ) from e
+            with open(output_file.name, "rb") as f:
+                result = pickle.load(f)
+        # Cache entries must be keyed by local indices (0..N-1) because
+        # gpu_p2p_access_check() is called with local ranks.
+        id_to_local = {device_id: local for local, device_id in enumerate(ids)}
+        for _i, _j, r in zip(batch_src, batch_tgt, result):
+            cache[f"{id_to_local[_i]}->{id_to_local[_j]}"] = r
+        with open(path, "w") as f:
+            json.dump(cache, f, indent=4)
+    if is_distributed:
+        get_world_group().barrier()
+    logger.info("reading GPU P2P access cache from %s", path)
+    with open(path) as f:
+        cache = json.load(f)
+    _gpu_p2p_access_cache = cache
+    return _gpu_p2p_access_cache[f"{src}->{tgt}"]
+
+
+__all__ = ["gpu_p2p_access_check"]
+
+if __name__ == "__main__":
+    batch_src, batch_tgt, output_file = pickle.loads(sys.stdin.buffer.read())
+    result = can_actually_p2p(batch_src, batch_tgt)
+    with open(output_file, "wb") as f:
+        f.write(pickle.dumps(result))

vllm/distributed/device_communicators/base_device_communicator.py

@@ -0,0 +1,383 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import threading
+from weakref import WeakValueDictionary
+
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+
+from vllm.utils import is_moe_layer
+
+
+class Cache:
+    def __init__(self):
+        self._cache: WeakValueDictionary = WeakValueDictionary()
+        self._lock = threading.RLock()  # Reentrant lock for thread safety
+
+    def get_or_create(self, kwargs, func):
+        # Create a hashable key from the kwargs
+        key = tuple(sorted((k, v) for k, v in kwargs.items()))
+
+        with self._lock:
+            instance = self._cache.get(key)
+            if instance is None:
+                instance = func(**kwargs)
+                self._cache[key] = instance
+            return instance
+
+
+class All2AllManagerBase:
+    rank: int
+    world_size: int
+
+    def __init__(self, cpu_group, tcp_store_group=None):
+        self.cpu_group = cpu_group
+        self.tcp_store_group = tcp_store_group
+
+        # compute some common properties
+        from vllm.distributed.parallel_state import (
+            get_dp_group,
+            get_tp_group,
+            in_the_same_node_as,
+        )
+
+        # all2all lives in ep group, which is merged from dp and tp group
+        self.dp_group = get_dp_group()
+        self.tp_group = get_tp_group()
+
+        # no self.ep_group since self.ep_group is still in construction
+        # when we create this object
+        self.dp_rank = self.dp_group.rank_in_group
+        self.dp_world_size = self.dp_group.world_size
+        self.rank = cpu_group.rank()
+        self.world_size = cpu_group.size()
+
+        # all2all communication often has separate implementations for
+        # intra-node and inter-node communication
+        if tcp_store_group is None:
+            self.internode = not all(in_the_same_node_as(cpu_group, source_rank=0))
+        else:
+            self.internode = not all(
+                in_the_same_node_as(tcp_store_group, source_rank=0)
+            )
+
+        self.support_fault_tolerance = False
+
+    def get_handle(self, kwargs):
+        # get a handle for the all2all communication,
+        # based on the kwargs.
+        # different layers can have different configs,
+        # e.g. one layer has hidden size 1024, another has 2048.
+        # usually the underlying implementation caches the handle
+        # and reuse it for the same config.
+        raise NotImplementedError
+
+    def dispatch_router_logits(
+        self,
+        hidden_states: torch.Tensor,
+        router_logits: torch.Tensor,
+        is_sequence_parallel: bool = False,
+        extra_tensors: list[torch.Tensor] | None = None,
+    ) -> (
+        tuple[torch.Tensor, torch.Tensor]
+        | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
+    ):
+        # Subclasses should either:
+        # - implement handling for extra_tensors, or
+        # - raise a clear error if extra_tensors is not supported.
+        raise NotImplementedError
+
+    def dispatch(
+        self,
+        hidden_states: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        is_sequence_parallel: bool = False,
+        extra_tensors: list[torch.Tensor] | None = None,
+    ) -> (
+        tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+        | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
+    ):
+        # Subclasses should either:
+        # - implement handling for extra_tensors, or
+        # - raise a clear error if extra_tensors is not supported.
+        raise NotImplementedError
+
+    def query_active_mask(self) -> torch.Tensor:
+        raise NotImplementedError
+
+    def query_fault(self) -> torch.Tensor:
+        """Returns has_fault scalar."""
+        raise NotImplementedError
+
+    def set_num_sms(self, num_sms: int):
+        pass
+
+    def max_sms_used(self) -> int | None:
+        return None  # None means it could use the whole GPU
+
+    def combine(self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False):
+        raise NotImplementedError
+
+    def destroy(self):
+        pass
+
+
+class DeviceCommunicatorBase:
+    """
+    Base class for device-specific communicator.
+    It can use the `cpu_group` to initialize the communicator.
+    If the device has PyTorch integration (PyTorch can recognize its
+    communication backend), the `device_group` will also be given.
+    """
+
+    def __init__(
+        self,
+        cpu_group: ProcessGroup,
+        device: torch.device | None = None,
+        device_group: ProcessGroup | None = None,
+        unique_name: str = "",
+        global_ranks: list[int] | None = None,
+        global_world_size: int | None = None,
+    ):
+        self.device = device or torch.device("cpu")
+        self.cpu_group = cpu_group
+        self.device_group = device_group
+        self.unique_name = unique_name
+
+        # Check if this is a stateless process group
+        from torch.distributed.distributed_c10d import _world
+
+        is_stateless = _world.pg_map.get(cpu_group, None) is None
+
+        if is_stateless:
+            # For stateless groups, we can't use torch.distributed methods
+            self.rank = cpu_group.rank()
+            self.world_size = cpu_group.size()
+            assert global_ranks is not None
+            assert global_world_size is not None
+            self.ranks = global_ranks
+            self.global_rank = self.ranks[self.rank]
+            self.global_world_size = global_world_size
+            self.rank_in_group = self.rank
+        else:
+            self.rank = dist.get_rank(cpu_group)
+            self.world_size = dist.get_world_size(cpu_group)
+            self.ranks = dist.get_process_group_ranks(cpu_group)
+            self.global_rank = dist.get_rank()
+            self.global_world_size = dist.get_world_size()
+            self.rank_in_group = dist.get_group_rank(self.cpu_group, self.global_rank)
+
+        use_ep = False
+        all2all_backend = None
+        from vllm.config import get_current_vllm_config_or_none
+
+        config = get_current_vllm_config_or_none()
+        if config is not None:
+            # initialize the all2all manager for DP or sequence-parallel EP.
+            use_ep = (
+                config.parallel_config.data_parallel_size > 1
+                or config.parallel_config.use_sequence_parallel_moe
+            )
+            all2all_backend = config.parallel_config.all2all_backend
+
+        self.is_ep_communicator = unique_name.split(":")[0] == "ep"
+        self.use_all2all = self.is_ep_communicator and use_ep
+        self.all2all_backend = all2all_backend
+        self.all2all_manager: All2AllManagerBase | None = None
+
+    def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
+        dist.all_reduce(input_, group=self.device_group)
+        return input_
+
+    def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
+        if dim < 0:
+            # Convert negative dim to positive.
+            dim += input_.dim()
+        input_size = input_.size()
+        # NOTE: we have to use concat-style all-gather here,
+        # stack-style all-gather has compatibility issues with
+        # torch.compile . see https://github.com/pytorch/pytorch/issues/138795
+        output_size = (input_size[0] * self.world_size,) + input_size[1:]
+        # Allocate output tensor.
+        output_tensor = torch.empty(
+            output_size, dtype=input_.dtype, device=input_.device
+        )
+        # All-gather.
+        dist.all_gather_into_tensor(output_tensor, input_, group=self.device_group)
+        # Reshape
+        output_tensor = output_tensor.reshape((self.world_size,) + input_size)
+        output_tensor = output_tensor.movedim(0, dim)
+        output_tensor = output_tensor.reshape(
+            input_size[:dim]
+            + (self.world_size * input_size[dim],)
+            + input_size[dim + 1 :]
+        )
+        return output_tensor
+
+    def all_gatherv(
+        self,
+        input_: torch.Tensor | list[torch.Tensor],
+        dim: int = 0,
+        sizes: list[int] | None = None,
+    ) -> torch.Tensor | list[torch.Tensor]:
+        raise NotImplementedError
+
+    def reduce_scatter(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
+        world_size = self.world_size
+        # Bypass the function if we are using only 1 GPU.
+        if world_size == 1:
+            return input_
+        assert -input_.dim() <= dim < input_.dim(), (
+            f"Invalid dim ({dim}) for input tensor with shape {input_.size()}"
+        )
+
+        if dim < 0:
+            # Convert negative dim to positive.
+            dim += input_.dim()
+
+        # Note: This will produce an incorrect answer if we don't make
+        # the input_tensor contiguous. Possible bug in reduce_scatter_tensor?
+        input_tensor = input_.movedim(0, dim).contiguous()
+
+        assert input_tensor.shape[0] % world_size == 0
+        chunk_size = input_tensor.shape[0] // world_size
+        output_shape = (chunk_size,) + input_tensor.shape[1:]
+
+        output_tensor = torch.empty(
+            output_shape, dtype=input_tensor.dtype, device=input_tensor.device
+        )
+
+        # Perform reduce-scatter operation
+        torch.distributed.reduce_scatter_tensor(
+            output_tensor, input_tensor, group=self.device_group
+        )
+
+        # Reshape before returning
+        return output_tensor.movedim(0, dim).contiguous()
+
+    def reduce_scatterv(
+        self, input_: torch.Tensor, dim: int = -1, sizes: list[int] | None = None
+    ) -> torch.Tensor:
+        raise NotImplementedError
+
+    def gather(
+        self, input_: torch.Tensor, dst: int = 0, dim: int = -1
+    ) -> torch.Tensor | None:
+        """
+        NOTE: We assume that the input tensor is on the same device across
+        all the ranks.
+        NOTE: `dst` is the local rank of the destination rank.
+        """
+        world_size = self.world_size
+        assert -input_.dim() <= dim < input_.dim(), (
+            f"Invalid dim ({dim}) for input tensor with shape {input_.size()}"
+        )
+        if dim < 0:
+            # Convert negative dim to positive.
+            dim += input_.dim()
+
+        # Allocate output tensor.
+        if self.rank_in_group == dst:
+            gather_list = [torch.empty_like(input_) for _ in range(world_size)]
+        else:
+            gather_list = None
+        # Gather.
+        torch.distributed.gather(
+            input_, gather_list, dst=self.ranks[dst], group=self.device_group
+        )
+        if self.rank_in_group == dst:
+            output_tensor = torch.cat(gather_list, dim=dim)
+        else:
+            output_tensor = None
+        return output_tensor
+
+    def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
+        """Sends a tensor to the destination rank in a blocking way"""
+        """NOTE: `dst` is the local rank of the destination rank."""
+        if dst is None:
+            dst = (self.rank_in_group + 1) % self.world_size
+        torch.distributed.send(tensor, self.ranks[dst], self.device_group)
+
+    def recv(
+        self, size: torch.Size, dtype: torch.dtype, src: int | None = None
+    ) -> torch.Tensor:
+        """Receives a tensor from the source rank."""
+        """NOTE: `src` is the local rank of the source rank."""
+        if src is None:
+            src = (self.rank_in_group - 1) % self.world_size
+
+        tensor = torch.empty(size, dtype=dtype, device=self.device)
+        torch.distributed.recv(tensor, self.ranks[src], self.device_group)
+        return tensor
+
+    def broadcast(self, tensor: torch.Tensor, src: int = 0) -> torch.Tensor:
+        """Broadcast a tensor from source rank to all ranks."""
+        if self.world_size == 1:
+            return tensor
+        torch.distributed.broadcast(tensor, self.ranks[src], self.device_group)
+        return tensor
+
+    def destroy(self):
+        pass
+
+    def prepare_communication_buffer_for_model(self, model: torch.nn.Module) -> None:
+        """
+        Prepare the communication buffer for the model.
+        """
+        if not self.is_ep_communicator:
+            return
+
+        moe_modules = [module for module in model.modules() if is_moe_layer(module)]
+        for module in moe_modules:
+            module.maybe_init_modular_kernel()
+
+    def dispatch_router_logits(
+        self,
+        hidden_states: torch.Tensor,
+        router_logits: torch.Tensor,
+        is_sequence_parallel: bool = False,
+        extra_tensors: list[torch.Tensor] | None = None,
+    ) -> (
+        tuple[torch.Tensor, torch.Tensor]
+        | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
+    ):
+        """
+        Dispatch the hidden states and router logits to the appropriate device.
+        This is a no-op in the base class.
+        """
+        if extra_tensors is not None:
+            return hidden_states, router_logits, extra_tensors
+        return hidden_states, router_logits
+
+    def dispatch(
+        self,
+        hidden_states: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        is_sequence_parallel: bool = False,
+        extra_tensors: list[torch.Tensor] | None = None,
+    ) -> (
+        tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+        | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
+    ):
+        """
+        Dispatch the hidden states and topk weights/ids to the appropriate device.
+        This is a no-op in the base class.
+        """
+        if extra_tensors is not None:
+            return hidden_states, topk_weights, topk_ids, extra_tensors
+        return hidden_states, topk_weights, topk_ids
+
+    def combine(
+        self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
+    ) -> torch.Tensor:
+        """
+        Combine the hidden states and router logits from the appropriate device.
+        This is a no-op in the base class.
+        """
+        return hidden_states
+
+    def batch_isend_irecv(self, p2p_ops: list):
+        raise NotImplementedError

vllm/distributed/device_communicators/cpu_communicator.py

@@ -0,0 +1,341 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+from typing import Any
+
+import torch
+from torch.distributed import ProcessGroup
+
+from vllm.distributed.utils import pickle
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+from vllm.platforms.interface import CpuArchEnum
+
+from .base_device_communicator import DeviceCommunicatorBase
+
+logger = init_logger(__name__)
+
+
+class CpuCommunicator(DeviceCommunicatorBase):
+    def __init__(
+        self,
+        cpu_group: ProcessGroup,
+        device: torch.device | None = None,
+        device_group: ProcessGroup | None = None,
+        unique_name: str = "",
+    ):
+        super().__init__(cpu_group, device, device_group, unique_name)
+        self.dist_module = torch.distributed
+
+        if (
+            (
+                current_platform.get_cpu_architecture() == CpuArchEnum.X86
+                or current_platform.get_cpu_architecture() == CpuArchEnum.ARM
+                or current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC
+            )
+            and hasattr(torch.ops._C, "init_shm_manager")
+            and (unique_name.startswith("tp") or unique_name.startswith("pp"))
+            and self._all_group_ranks_share_shm_group_name()
+        ):
+            self.dist_module = _CPUSHMDistributed(self)
+        elif unique_name.startswith("tp") or unique_name.startswith("pp"):
+            logger.info(
+                "CPU SHM communicator disabled for group %s: ranks do not share "
+                "the same SHM group name, falling back to torch.distributed.",
+                unique_name,
+            )
+
+        # send/recv tensor_dict is only supported through the SHM communicator backend
+        self.supports_tensor_dict = isinstance(self.dist_module, _CPUSHMDistributed)
+
+        if self.use_all2all:
+            if self.all2all_backend not in (
+                "naive",
+                "allgather_reducescatter",
+            ):  # type: ignore[has-type]
+                logger.warning(
+                    "`%s` all2all manager is not supported on CPU. "
+                    "Falling back to `allgather_reducescatter` manager.",
+                    self.all2all_backend,  # type: ignore[has-type]
+                )
+            from .all2all import AgRsAll2AllManager
+
+            self.all2all_manager = AgRsAll2AllManager(self.cpu_group)
+            logger.info("Using allgather_reducescatter all2all manager.")
+
+    def _all_group_ranks_share_shm_group_name(self) -> bool:
+        """
+        CPUSHM requires all ranks in this group to agree on one SHM group name.
+        This is a lightweight consistency check for VLLM_DIST_IDENT/name inputs.
+        """
+        local_name = _CPUSHMDistributed.make_group_name(self)
+        names: list[str] = [""] * self.world_size
+        torch.distributed.all_gather_object(
+            names,
+            local_name,
+            group=self.device_group,
+        )
+        return len(set(names)) == 1
+
+    def all_reduce(self, input_):
+        self.dist_module.all_reduce(input_, group=self.device_group)
+        return input_
+
+    def gather(
+        self, input_: torch.Tensor, dst: int = 0, dim: int = -1
+    ) -> torch.Tensor | None:
+        """
+        NOTE: We assume that the input tensor is on the same device across
+        all the ranks.
+        NOTE: `dst` is the local rank of the destination rank.
+        """
+        world_size = self.world_size
+        assert -input_.dim() <= dim < input_.dim(), (
+            f"Invalid dim ({dim}) for input tensor with shape {input_.size()}"
+        )
+        if dim < 0:
+            # Convert negative dim to positive.
+            dim += input_.dim()
+
+        # Allocate output tensor.
+        if self.rank_in_group == dst:
+            gather_list = [torch.empty_like(input_) for _ in range(world_size)]
+        else:
+            gather_list = None
+
+        # Gather.
+        self.dist_module.gather(
+            input_, gather_list, dst=self.ranks[dst], group=self.device_group
+        )
+
+        if self.rank_in_group == dst:
+            output_tensor = torch.cat(gather_list, dim=dim)
+        else:
+            output_tensor = None
+        return output_tensor
+
+    def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
+        if dim < 0:
+            # Convert negative dim to positive.
+            dim += input_.dim()
+        input_size = input_.size()
+        # NOTE: we have to use concat-style all-gather here,
+        # stack-style all-gather has compatibility issues with
+        # torch.compile . see https://github.com/pytorch/pytorch/issues/138795
+        output_size = (input_size[0] * self.world_size,) + input_size[1:]
+        # Allocate output tensor.
+        output_tensor = torch.empty(
+            output_size, dtype=input_.dtype, device=input_.device
+        )
+        # All-gather.
+        self.dist_module.all_gather_into_tensor(
+            output_tensor, input_, group=self.device_group
+        )
+
+        # Reshape
+        output_tensor = output_tensor.reshape((self.world_size,) + input_size)
+        output_tensor = output_tensor.movedim(0, dim)
+        output_tensor = output_tensor.reshape(
+            input_size[:dim]
+            + (self.world_size * input_size[dim],)
+            + input_size[dim + 1 :]
+        )
+        return output_tensor
+
+    def send_tensor_dict(
+        self,
+        tensor_dict: dict[str, torch.Tensor | Any],
+        dst: int,
+    ) -> None:
+        if not self.supports_tensor_dict:
+            raise NotImplementedError(
+                "CpuCommunicator does not support tensor dict fastpath with "
+                "torch.distributed backend."
+            )
+        return self.dist_module.send_tensor_dict(tensor_dict, dst)
+
+    def recv_tensor_dict(
+        self,
+        src: int,
+    ) -> dict[str, torch.Tensor | Any]:
+        if not self.supports_tensor_dict:
+            raise NotImplementedError(
+                "CpuCommunicator does not support tensor dict fastpath with "
+                "torch.distributed backend."
+            )
+        return self.dist_module.recv_tensor_dict(src)
+
+    def dispatch_router_logits(
+        self,
+        hidden_states: torch.Tensor,
+        router_logits: torch.Tensor,
+        is_sequence_parallel: bool = False,
+        extra_tensors: list[torch.Tensor] | None = None,
+    ) -> (
+        tuple[torch.Tensor, torch.Tensor]
+        | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
+    ):
+        """
+        Dispatch the hidden states and router logits to the appropriate device.
+        This is a no-op in the base class.
+        """
+
+        assert self.all2all_manager is not None
+        return self.all2all_manager.dispatch_router_logits(
+            hidden_states,
+            router_logits,
+            is_sequence_parallel,
+            extra_tensors,
+        )
+
+    def dispatch(
+        self,
+        hidden_states: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        is_sequence_parallel: bool = False,
+        extra_tensors: list[torch.Tensor] | None = None,
+    ) -> (
+        tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+        | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
+    ):
+        """
+        Dispatch the hidden states and topk weights/ids to the appropriate device.
+        This is a no-op in the base class.
+        """
+        assert self.all2all_manager is not None
+        return self.all2all_manager.dispatch(
+            hidden_states,
+            topk_weights,
+            topk_ids,
+            is_sequence_parallel,
+            extra_tensors=extra_tensors,
+        )
+
+    def combine(
+        self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
+    ) -> torch.Tensor:
+        """
+        Combine the hidden states and router logits from the appropriate device.
+        This is a no-op in the base class.
+        """
+        assert self.all2all_manager is not None
+        return self.all2all_manager.combine(
+            hidden_states,
+            is_sequence_parallel,
+        )
+
+
+class _CPUSHMDistributed:
+    def __init__(self, communicator: CpuCommunicator):
+        self.communicator = communicator
+
+        self.group_name = self.make_group_name(communicator)
+
+        self.handle = self._init_cpu_shm()
+
+    @staticmethod
+    def make_group_name(communicator: CpuCommunicator) -> str:
+        instance_identifier = os.environ["VLLM_DIST_IDENT"]
+        unique_name = communicator.unique_name
+        instance_identifier = f"{instance_identifier}-{unique_name}"
+        group_ranks = [str(rank) for rank in communicator.ranks]
+        shm_group_identifier = f"[{'-'.join(group_ranks)}]"
+        return f"{instance_identifier}-{shm_group_identifier}-cpushm"
+
+    def _init_cpu_shm(self) -> int:
+        thread_num_tensor = torch.tensor(
+            [torch.get_num_threads()],
+            dtype=torch.int64,
+        )
+        torch.distributed.all_reduce(
+            thread_num_tensor,
+            op=torch.distributed.ReduceOp.MIN,
+            group=self.communicator.device_group,
+        )
+        thread_num = thread_num_tensor.item()
+
+        handle = torch.ops._C.init_shm_manager(
+            self.group_name,
+            self.communicator.world_size,
+            self.communicator.rank,
+            thread_num,
+        )
+        torch.distributed.barrier(self.communicator.device_group)
+        torch.ops._C.join_shm_manager(
+            handle,
+            self.group_name,
+        )
+        torch.distributed.barrier(self.communicator.device_group)
+
+        return handle
+
+    def all_reduce(
+        self, input: torch.Tensor, group: ProcessGroup | None = None
+    ) -> None:
+        torch.ops._C.shm_allreduce(self.handle, input)
+
+    def gather(
+        self,
+        input: torch.Tensor,
+        gather_list: list[torch.Tensor] | None,
+        dst: int = -1,
+        group: ProcessGroup | None = None,
+    ) -> None:
+        # Note: different from the torch gather, here we use local dst rank.
+        torch.ops._C.shm_gather(
+            self.handle,
+            input,
+            gather_list,
+            torch.distributed.get_group_rank(group, dst),
+        )
+
+    def all_gather_into_tensor(
+        self,
+        output: torch.Tensor,
+        input: torch.Tensor,
+        group: ProcessGroup | None = None,
+    ) -> None:
+        torch.ops._C.shm_all_gather(self.handle, input, output)
+
+    def send_tensor_dict(
+        self,
+        tensor_dict: dict[str, torch.Tensor | Any],
+        dst: int,
+    ) -> None:
+        key_list = list(tensor_dict.keys())
+        value_list = list(tensor_dict.values())
+        size_list = []
+        for v in value_list:
+            if not isinstance(v, torch.Tensor):
+                raise RuntimeError("CpuCommunicator only supports sending tensors.")
+            size_list.append(v.size())
+        key_size_tensor = torch.frombuffer(
+            pickle.dumps([key_list, size_list]), dtype=torch.uint8
+        )
+        value_list.append(key_size_tensor)
+
+        torch.ops._C.shm_send_tensor_list(self.handle, value_list, dst)
+
+        return None
+
+    def recv_tensor_dict(
+        self,
+        src: int,
+    ) -> dict[str, torch.Tensor | Any]:
+        tensor_list = torch.ops._C.shm_recv_tensor_list(self.handle, src)
+
+        value_list: list[torch.Tensor] = tensor_list[:-1]
+        key_size_tensor = tensor_list[-1]
+
+        key_size = pickle.loads(key_size_tensor.numpy().tobytes())
+        key_list = key_size[0]
+        size_list = key_size[1]
+        assert len(key_list) == len(size_list)
+        assert len(key_list) == len(value_list)
+
+        tensor_dict: dict[str, torch.Tensor] = {}
+        for key, size, t in zip(key_list, size_list, value_list):
+            tensor_dict[key] = t.view(size)
+        return tensor_dict

vllm/distributed/device_communicators/cuda_communicator.py

@@ -0,0 +1,745 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import torch
+from torch.distributed import ProcessGroup
+
+import vllm.envs as envs
+from vllm._aiter_ops import rocm_aiter_ops
+from vllm.distributed.device_communicators.all_reduce_utils import (
+    NCCL_SYMM_MEM_ALL_REDUCE_CONFIG,
+    should_nccl_symm_mem_ag_rs,
+    should_nccl_symm_mem_allreduce,
+)
+from vllm.distributed.device_communicators.pynccl import register_nccl_symmetric_ops
+from vllm.distributed.device_communicators.pynccl_allocator import (
+    is_symmetric_memory_enabled,
+)
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+from ..utils import StatelessProcessGroup
+from .aiter_custom_all_reduce import AiterCustomAllreduce
+from .base_device_communicator import DeviceCommunicatorBase
+
+logger = init_logger(__name__)
+
+
+class CudaCommunicator(DeviceCommunicatorBase):
+    def __init__(
+        self,
+        cpu_group: ProcessGroup,
+        device: torch.device | None = None,
+        device_group: ProcessGroup | None = None,
+        unique_name: str = "",
+        global_ranks: list[int] | None = None,
+        global_world_size: int | None = None,
+        tcp_store_group: StatelessProcessGroup | None = None,
+    ):
+        super().__init__(
+            cpu_group,
+            device,
+            device_group,
+            unique_name,
+            global_ranks,
+            global_world_size,
+        )
+        if "tp" not in unique_name:
+            # custom allreduce or torch symm mem can be used only by tp
+            use_custom_allreduce = False
+            use_torch_symm_mem = False
+            use_flashinfer_allreduce = False
+            use_aiter_allreduce = False
+        else:
+            from vllm.distributed.parallel_state import _ENABLE_CUSTOM_ALL_REDUCE
+
+            use_custom_allreduce = _ENABLE_CUSTOM_ALL_REDUCE
+            use_torch_symm_mem = envs.VLLM_ALLREDUCE_USE_SYMM_MEM
+            use_flashinfer_allreduce = envs.VLLM_ALLREDUCE_USE_FLASHINFER
+            use_aiter_allreduce = use_custom_allreduce and bool(
+                rocm_aiter_ops.is_custom_all_reduce_enabled()
+            )
+
+        self.use_custom_allreduce = use_custom_allreduce
+        self.use_torch_symm_mem = use_torch_symm_mem
+        self.use_flashinfer_allreduce = use_flashinfer_allreduce
+        self.use_aiter_allreduce = use_aiter_allreduce
+
+        # lazy import to avoid documentation build error
+        from vllm.distributed.device_communicators.custom_all_reduce import (
+            CustomAllreduce,
+        )
+        from vllm.distributed.device_communicators.flashinfer_all_reduce import (
+            FlashInferAllReduce,
+        )
+        from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
+        from vllm.distributed.device_communicators.quick_all_reduce import (
+            QuickAllReduce,
+        )
+        from vllm.distributed.device_communicators.symm_mem import SymmMemCommunicator
+
+        self.pynccl_comm: PyNcclCommunicator | None = None
+        if self.world_size > 1:
+            self.pynccl_comm = PyNcclCommunicator(
+                group=self.cpu_group if tcp_store_group is None else tcp_store_group,
+                device=self.device,
+            )
+            if is_symmetric_memory_enabled():
+                register_nccl_symmetric_ops(self.pynccl_comm)
+
+        self.ca_comm: CustomAllreduce | None = None
+        self.qr_comm: QuickAllReduce | None = None
+        self.symm_mem_comm: SymmMemCommunicator | None = None
+        self.fi_ar_comm: FlashInferAllReduce | None = None
+        self.aiter_ar_comm: AiterCustomAllreduce | None = None
+
+        if use_torch_symm_mem and current_platform.is_cuda():
+            self.symm_mem_comm = SymmMemCommunicator(
+                group=self.cpu_group,
+                device=self.device,
+            )
+
+        if self.use_flashinfer_allreduce and self.world_size > 1:
+            self.fi_ar_comm = FlashInferAllReduce(
+                group=self.cpu_group,
+                device=self.device,
+            )
+
+        if self.use_aiter_allreduce and self.world_size > 1:
+            self.aiter_ar_comm = AiterCustomAllreduce(
+                group=self.cpu_group,
+                device=self.device,
+            )
+
+        if use_custom_allreduce and self.aiter_ar_comm is None and self.world_size > 1:
+            # Initialize a custom fast all-reduce implementation.
+            self.ca_comm = CustomAllreduce(
+                group=self.cpu_group,
+                device=self.device,
+                symm_mem_enabled=(
+                    self.symm_mem_comm is not None and not self.symm_mem_comm.disabled
+                ),
+            )
+
+        if use_custom_allreduce and self.world_size > 1 and current_platform.is_rocm():
+            # Initialize a custom quick all-reduce implementation for AMD.
+            # Quick reduce is designed as a complement to custom allreduce
+            # (vLLM's or AITER's), so it is initialized for either backend.
+            # Based on quickreduce (https://github.com/mk1-project/quickreduce).
+            # On ROCm, 'use_custom_allreduce==True' means it must currently be
+            # an MI300 series.
+            self.qr_comm = QuickAllReduce(group=self.cpu_group, device=self.device)
+
+        if self.world_size > 1:
+            self._log_all_reduce_backend_selection()
+
+        if self.use_all2all:
+            if self.all2all_backend in ("naive", "allgather_reducescatter"):
+                from .all2all import AgRsAll2AllManager
+
+                self.all2all_manager = AgRsAll2AllManager(
+                    self.cpu_group, tcp_store_group
+                )
+            elif self.all2all_backend == "deepep_high_throughput":
+                from .all2all import DeepEPHTAll2AllManager
+
+                self.all2all_manager = DeepEPHTAll2AllManager(
+                    self.cpu_group, tcp_store_group
+                )
+            elif self.all2all_backend == "deepep_low_latency":
+                from .all2all import DeepEPLLAll2AllManager
+
+                self.all2all_manager = DeepEPLLAll2AllManager(
+                    self.cpu_group, tcp_store_group
+                )
+            elif self.all2all_backend in (
+                "mori_high_throughput",
+                "mori_low_latency",
+            ):
+                from .all2all import MoriAll2AllManager
+
+                self.all2all_manager = MoriAll2AllManager(
+                    self.cpu_group, self.all2all_backend
+                )
+            elif self.all2all_backend == "deepep_v2":
+                from .all2all import DeepEPV2All2AllManager
+
+                self.all2all_manager = DeepEPV2All2AllManager(
+                    self.cpu_group,
+                    tcp_store_group,
+                    device_group=self.device_group,
+                )
+            elif self.all2all_backend == "nixl_ep":
+                from .all2all import NixlEPAll2AllManager
+
+                self.all2all_manager = NixlEPAll2AllManager(
+                    self.cpu_group, tcp_store_group
+                )
+            elif (
+                self.all2all_backend == "flashinfer_all2allv"
+                or self.all2all_backend == "flashinfer_nvlink_two_sided"
+            ):
+                if self.all2all_backend == "flashinfer_all2allv":
+                    logger.warning_once(
+                        "'flashinfer_all2allv' is deprecated and has been renamed to"
+                        "'flashinfer_nvlink_two_sided'. It will be removed in a future"
+                        "release."
+                    )
+                from .all2all import FlashInferNVLinkTwoSidedManager
+
+                self.all2all_manager = FlashInferNVLinkTwoSidedManager(
+                    self.cpu_group, tcp_store_group
+                )
+            elif self.all2all_backend == "flashinfer_nvlink_one_sided":
+                from .all2all import FlashInferNVLinkOneSidedManager
+
+                self.all2all_manager = FlashInferNVLinkOneSidedManager(self.cpu_group)
+            else:
+                raise ValueError(f"Unknown all2all backend: {self.all2all_backend}")
+
+            logger.info_once(
+                "Using %s all2all manager.",
+                self.all2all_manager.__class__.__name__,
+                scope="global",
+            )
+
+    def _log_all_reduce_backend_selection(self) -> None:
+        """Log the all-reduce backends that are active for this group.
+
+        The dispatch chain in ``all_reduce`` tries backends in this order and
+        falls through to the next one if the current backend rejects the
+        input (size/dtype gates) or is disabled. The list of "enabled"
+        backends below is the subset of potential backends that may be
+        chosen at dispatch time for this group; the actual per-call choice
+        depends on the input tensor.
+        """
+        all_potential_ar_backends = [
+            "NCCL_SYMM_MEM",
+            "QUICK_REDUCE",
+            "FLASHINFER",
+            "AITER_CUSTOM",
+            "CUSTOM",
+            "SYMM_MEM",
+            "PYNCCL",
+        ]
+        enabled_ar_backends: list[str] = []
+        # Mirror the static preconditions of `should_nccl_symm_mem_allreduce`:
+        # VLLM_BATCH_INVARIANT off, NCCL symm mem enabled, world_size meets
+        # min_world_size, and world_size either has a tuned entry in
+        # `custom_ar_preferred_ranges` or is greater than
+        # `always_use_above_world_size`. World sizes that fail the latter (e.g.
+        # 5/6/7 with the default config) never dispatch NCCL symm mem
+        # regardless of input. The per-tensor-size check inside the function
+        # stays as a runtime decision.
+        nccl_symm_ws_ok = self.world_size >= NCCL_SYMM_MEM_ALL_REDUCE_CONFIG[
+            "min_world_size"
+        ] and (
+            self.world_size
+            in NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["custom_ar_preferred_ranges"]
+            or self.world_size
+            > NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["always_use_above_world_size"]
+        )
+        if (
+            self.pynccl_comm is not None
+            and not self.pynccl_comm.disabled
+            and is_symmetric_memory_enabled()
+            and not envs.VLLM_BATCH_INVARIANT
+            and nccl_symm_ws_ok
+        ):
+            enabled_ar_backends.append("NCCL_SYMM_MEM")
+        if self.qr_comm is not None and not self.qr_comm.disabled:
+            enabled_ar_backends.append("QUICK_REDUCE")
+        if self.fi_ar_comm is not None and not self.fi_ar_comm.disabled:
+            enabled_ar_backends.append("FLASHINFER")
+        if self.aiter_ar_comm is not None and not self.aiter_ar_comm.disabled:
+            enabled_ar_backends.append("AITER_CUSTOM")
+        if self.ca_comm is not None and not self.ca_comm.disabled:
+            enabled_ar_backends.append("CUSTOM")
+        if self.symm_mem_comm is not None and not self.symm_mem_comm.disabled:
+            enabled_ar_backends.append("SYMM_MEM")
+        if self.pynccl_comm is not None and not self.pynccl_comm.disabled:
+            enabled_ar_backends.append("PYNCCL")
+
+        logger.info_once(
+            "Using %s all-reduce backends (in dispatch order) for group "
+            "'%s' out of potential backends: %s.",
+            "[" + ", ".join(f"'{b}'" for b in enabled_ar_backends) + "]",
+            self.unique_name or "<unnamed>",
+            "[" + ", ".join(f"'{b}'" for b in all_potential_ar_backends) + "]",
+            scope="global",
+        )
+
+    def all_reduce(self, input_):
+        # since currently we perform copy input -> symm_input -> out-of-place AR
+        # return symm_output, we don't need to check if input is symmetric
+        if self.pynccl_comm is not None and should_nccl_symm_mem_allreduce(
+            self.pynccl_comm.world_size, input_
+        ):
+            out = torch.ops.vllm.all_reduce_symmetric_with_copy(input_)
+            if out is not None:
+                return out
+        # always try quick reduce first, then flashinfer, then the AITER or vLLM
+        # custom allreduce, and then pynccl. (quick reduce just for ROCM MI3*)
+        qr_comm = self.qr_comm
+        if (
+            qr_comm is not None
+            and not qr_comm.disabled
+            and qr_comm.should_quick_allreduce(input_)
+        ):
+            out = qr_comm.quick_all_reduce(input_)
+            assert out is not None
+            return out
+        fi_ar_comm = self.fi_ar_comm
+        if (
+            fi_ar_comm is not None
+            and not fi_ar_comm.disabled
+            and fi_ar_comm.should_use_fi_ar(input_)
+        ):
+            out = fi_ar_comm.all_reduce(input_)
+            assert out is not None
+            return out
+        aiter_ar_comm = self.aiter_ar_comm
+        if (
+            aiter_ar_comm is not None
+            and not aiter_ar_comm.disabled
+            and aiter_ar_comm.should_custom_ar(input_)
+        ):
+            out = aiter_ar_comm.custom_all_reduce(input_)
+            assert out is not None
+            return out
+        ca_comm = self.ca_comm
+        if (
+            ca_comm is not None
+            and not ca_comm.disabled
+            and ca_comm.should_custom_ar(input_)
+        ):
+            out = ca_comm.custom_all_reduce(input_)
+            assert out is not None
+            return out
+        symm_mem_comm = self.symm_mem_comm
+        if symm_mem_comm is not None and symm_mem_comm.should_use_symm_mem(input_):
+            out = symm_mem_comm.all_reduce(input_)
+            assert out is not None
+            return out
+        pynccl_comm = self.pynccl_comm
+        if pynccl_comm is None or pynccl_comm.disabled:
+            out = input_.clone()
+            torch.distributed.all_reduce(out, group=self.device_group)
+            return out
+        assert pynccl_comm is not None
+        out = pynccl_comm.all_reduce(input_)
+        if out is None:
+            # fall back to the default all-reduce using PyTorch.
+            # this usually happens during testing.
+            # when we run the model, allreduce only happens for the TP
+            # group, where we always have either custom allreduce or pynccl.
+            out = input_.clone()
+            torch.distributed.all_reduce(out, group=self.device_group)
+        return out
+
+    def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
+        # Route uniform dim-0 all-gathers through NVLS symmetric memory when
+        # enabled (mirrors reduce_scatter); otherwise fall back to the
+        # PyNccl/base-class all-gather. Sequence parallelism's
+        # gather-before-GEMM uses dim=0 with tp-aligned (uniform) shards.
+        if dim < 0:
+            dim += input_.dim()
+        if dim == 0 and should_nccl_symm_mem_ag_rs():
+            return self._all_gather_symm_mem(input_.contiguous())
+
+        pynccl_comm = self.pynccl_comm
+        if pynccl_comm is None or pynccl_comm.disabled:
+            return super().all_gather(input_, dim)
+
+        # On ROCm, the base-class all_gather (all_gather_into_tensor) is faster
+        # than the manual pynccl + torch.empty + movedim + reshape path below,
+        # which adds a per-call output allocation and (for dim != 0) an extra
+        # copy on every step. This is on the hot path for TP forward passes, so
+        # keep ROCm on the base-class collective to avoid a decode regression.
+        if current_platform.is_rocm():
+            return super().all_gather(input_, dim)
+
+        input_size = input_.size()
+        output_size = (input_size[0] * self.world_size,) + input_size[1:]
+        output_tensor = torch.empty(
+            output_size, dtype=input_.dtype, device=input_.device
+        )
+        pynccl_comm.all_gather(output_tensor, input_.contiguous())
+        output_tensor = output_tensor.reshape((self.world_size,) + input_size)
+        output_tensor = output_tensor.movedim(0, dim)
+        return output_tensor.reshape(
+            input_size[:dim]
+            + (self.world_size * input_size[dim],)
+            + input_size[dim + 1 :]
+        )
+
+    def reduce_scatter(self, input_: torch.Tensor, dim: int = -1):
+        world_size = self.world_size
+        pynccl_comm = self.pynccl_comm
+        assert pynccl_comm is not None
+        if dim < 0:
+            # Convert negative dim to positive.
+            dim += input_.dim()
+
+        # Note: This will produce an incorrect answer if we don't make
+        # the input_tensor contiguous. Possible bug in reduce_scatter_tensor?
+        input_tensor = input_.movedim(0, dim).contiguous()
+
+        assert input_tensor.shape[0] % world_size == 0
+        chunk_size = input_tensor.shape[0] // world_size
+        output_shape = (chunk_size,) + input_tensor.shape[1:]
+
+        if should_nccl_symm_mem_ag_rs():
+            output = self._reduce_scatter_symm_mem(input_tensor)
+        else:
+            output = torch.empty(
+                output_shape, dtype=input_tensor.dtype, device=input_tensor.device
+            )
+            pynccl_comm.reduce_scatter(output, input_tensor)
+
+        # Reshape before returning
+        return output.movedim(0, dim).contiguous()
+
+    def reduce_scatterv(
+        self, input_: torch.Tensor, dim: int = -1, sizes: list[int] | None = None
+    ):
+        world_size = self.world_size
+        pynccl_comm = self.pynccl_comm
+        assert pynccl_comm is not None
+        if dim < 0:
+            # Convert negative dim to positive.
+            dim += input_.dim()
+
+        # Note: This will produce an incorrect answer if we don't make
+        # the input_tensor contiguous. Possible bug in reduce_scatter_tensor?
+        input_tensor = input_.movedim(0, dim).contiguous()
+
+        if sizes is not None:
+            assert len(sizes) == world_size, f"{len(sizes)} == {world_size}"
+            assert input_tensor.shape[0] == sum(sizes)
+            chunk_size = sizes[self.rank_in_group]
+        else:
+            assert input_tensor.shape[0] % world_size == 0
+            chunk_size = input_tensor.shape[0] // world_size
+        output_shape = (chunk_size,) + input_tensor.shape[1:]
+
+        # Symmetric memory is only used when all ranks have uniform sizes.
+        # ncclCommWindowRegister is collective: asymmetric pool allocations
+        # from variable per-rank sizes cause deadlocks.
+        use_symm_mem = sizes is None and should_nccl_symm_mem_ag_rs()
+        if use_symm_mem:
+            output = self._reduce_scatter_symm_mem(input_tensor)
+        else:
+            output = torch.empty(
+                output_shape, dtype=input_tensor.dtype, device=input_tensor.device
+            )
+            if sizes is not None and sizes.count(sizes[0]) != len(sizes):
+                pynccl_comm.reduce_scatterv(output, input_tensor, sizes=sizes)
+            else:
+                pynccl_comm.reduce_scatter(output, input_tensor)
+
+        # Reshape before returning
+        return output.movedim(0, dim).contiguous()
+
+    def _get_symm_scratch(
+        self,
+        role: str,
+        shape: tuple[int, ...],
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> torch.Tensor:
+        """Persistent, pre-registered NCCL symmetric-memory scratch buffer.
+
+        Allocating a fresh symm tensor per collective pays the
+        ``nccl_symm_mem_context`` snapshot + window-registration scan on every
+        call (~0.5 ms/RS+AG pair, dwarfing the NVLS transfer itself). Instead we
+        allocate once per ``(role, shape, dtype)``, register once, and reuse.
+
+        Safe for serial (eager) sequence parallelism: each collective's result
+        is consumed on the same stream before the next same-role collective
+        reuses the buffer. Distinct roles (e.g. ``rs_in`` vs ``ag_out``, both
+        full-size) get distinct buffers so a reduce-scatter input copy never
+        clobbers a still-live all-gather output.
+        """
+        from vllm.distributed.device_communicators.pynccl_allocator import (
+            nccl_symm_mem_context,
+        )
+
+        pynccl_comm = self.pynccl_comm
+        assert pynccl_comm is not None
+        cache = self.__dict__.setdefault("_symm_scratch_bufs", {})
+        key = (role, tuple(shape), dtype)
+        buf = cache.get(key)
+        if buf is None:
+            with nccl_symm_mem_context(pynccl_comm):
+                buf = torch.empty(shape, dtype=dtype, device=device)
+            cache[key] = buf
+        return buf
+
+    def _reduce_scatter_symm_mem(
+        self,
+        input_tensor: torch.Tensor,
+    ) -> torch.Tensor:
+        """ReduceScatter using NCCL symmetric memory (NVLS).
+
+        Only called for uniform-size reduce_scatter (variable sizes are
+        guarded out by the caller to avoid asymmetric ncclCommWindowRegister).
+        Uses persistent pre-registered scratch (see _get_symm_scratch).
+        """
+        from vllm.distributed.device_communicators.pynccl_allocator import (
+            is_symmetric_memory_tensor,
+        )
+
+        pynccl_comm = self.pynccl_comm
+        assert pynccl_comm is not None
+
+        chunk = input_tensor.shape[0] // self.world_size
+        output_shape = (chunk,) + tuple(input_tensor.shape[1:])
+
+        symm_output = self._get_symm_scratch(
+            "rs_out", output_shape, input_tensor.dtype, input_tensor.device
+        )
+        # NVLS reduce-scatter (LDMC) requires the input in symmetric memory.
+        if is_symmetric_memory_tensor(input_tensor):
+            symm_input = input_tensor
+        else:
+            symm_input = self._get_symm_scratch(
+                "rs_in",
+                tuple(input_tensor.shape),
+                input_tensor.dtype,
+                input_tensor.device,
+            )
+            symm_input.copy_(input_tensor)
+
+        pynccl_comm.reduce_scatter(symm_output, symm_input)
+        return symm_output
+
+    def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
+        """Sends a tensor to the destination rank in a blocking way"""
+        """NOTE: `dst` is the local rank of the destination rank."""
+        if dst is None:
+            dst = (self.rank_in_group + 1) % self.world_size
+
+        pynccl_comm = self.pynccl_comm
+        if pynccl_comm is not None and not pynccl_comm.disabled:
+            pynccl_comm.send(tensor, dst)
+        else:
+            torch.distributed.send(tensor, self.ranks[dst], self.device_group)
+
+    def recv(
+        self, size: torch.Size, dtype: torch.dtype, src: int | None = None
+    ) -> torch.Tensor:
+        """Receives a tensor from the source rank."""
+        """NOTE: `src` is the local rank of the source rank."""
+        if src is None:
+            src = (self.rank_in_group - 1) % self.world_size
+
+        tensor = torch.empty(size, dtype=dtype, device=self.device)
+        pynccl_comm = self.pynccl_comm
+        if pynccl_comm is not None and not pynccl_comm.disabled:
+            pynccl_comm.recv(tensor, src)
+        else:
+            torch.distributed.recv(tensor, self.ranks[src], self.device_group)
+        return tensor
+
+    def broadcast(self, tensor: torch.Tensor, src: int = 0) -> torch.Tensor:
+        """Broadcast a tensor from source rank to all ranks."""
+        if self.world_size == 1:
+            return tensor
+
+        pynccl_comm = self.pynccl_comm
+        if pynccl_comm is not None and not pynccl_comm.disabled:
+            pynccl_comm.broadcast(tensor, src)
+            return tensor
+        else:
+            raise ValueError("No PyNCCL communicator found")
+
+    def destroy(self):
+        if self.pynccl_comm is not None:
+            self.pynccl_comm.destroy()
+            self.pynccl_comm = None
+        if self.ca_comm is not None:
+            self.ca_comm = None
+        if self.aiter_ar_comm is not None:
+            self.aiter_ar_comm.close()
+            self.aiter_ar_comm = None
+        if self.fi_ar_comm is not None:
+            self.fi_ar_comm.destroy()
+            self.fi_ar_comm = None
+        if self.all2all_manager is not None:
+            self.all2all_manager.destroy()
+            self.all2all_manager = None  # type: ignore[assignment]
+
+    def all_gatherv(
+        self,
+        input_: torch.Tensor | list[torch.Tensor],
+        dim: int = 0,
+        sizes: list[int] | None = None,
+    ):
+        if dim != 0:
+            raise NotImplementedError("only dim 0 all-gatherv is supported")
+        world_size = self.world_size
+        pynccl_comm = self.pynccl_comm
+        assert pynccl_comm is not None and not pynccl_comm.disabled
+
+        # 'sizes' is not needed if all inputs in the same group have the same
+        # shape
+        if sizes is not None and all(s == sizes[0] for s in sizes):
+            sizes = None
+
+        # Symmetric memory is only used when all ranks have uniform sizes.
+        # ncclCommWindowRegister is collective: asymmetric pool allocations
+        # from variable per-rank sizes cause deadlocks.
+        if sizes is None and should_nccl_symm_mem_ag_rs():
+            if isinstance(input_, torch.Tensor):
+                return self._all_gather_symm_mem(input_)
+            return self._all_gather_batched_symm_mem(input_)
+
+        def _all_gather_single(input_: torch.Tensor, sizes: list[int] | None = None):
+            input_size = input_.size()
+            if sizes is not None:
+                assert len(sizes) == world_size
+                assert input_.shape[dim] == sizes[self.rank_in_group], (
+                    f"{input_.shape[dim]} != {sizes[self.rank_in_group]}"
+                )
+                output_size = (sum(sizes),) + input_size[1:]
+            else:
+                output_size = (input_size[0] * world_size,) + input_size[1:]
+            # Allocate output tensor.
+            output_tensor = torch.empty(
+                output_size, dtype=input_.dtype, device=input_.device
+            )
+            if sizes is not None:
+                pynccl_comm.all_gatherv(output_tensor, input_, sizes=sizes)
+            else:
+                pynccl_comm.all_gather(output_tensor, input_)
+            return output_tensor
+
+        if isinstance(input_, torch.Tensor):
+            return _all_gather_single(input_, sizes)
+
+        output_list = []
+        pynccl_comm.group_start()
+        for inp in input_:
+            output_list.append(_all_gather_single(inp, sizes=sizes))
+        pynccl_comm.group_end()
+
+        return output_list
+
+    def _all_gather_symm_mem(self, input_: torch.Tensor) -> torch.Tensor:
+        """AllGather a single tensor using NCCL symmetric memory (NVLS).
+
+        Only the output needs to be in symmetric memory; NCCL does not
+        require the AG input to be symmetrically allocated.
+        """
+        pynccl_comm = self.pynccl_comm
+        assert pynccl_comm is not None
+
+        out_size = (input_.size(0) * self.world_size,) + tuple(input_.size()[1:])
+        # Persistent pre-registered scratch avoids the per-call symm-mem context
+        # snapshot/registration overhead (see _get_symm_scratch).
+        symm_output = self._get_symm_scratch(
+            "ag_out", out_size, input_.dtype, input_.device
+        )
+        pynccl_comm.all_gather(symm_output, input_)
+        return symm_output
+
+    def _all_gather_batched_symm_mem(
+        self, inputs: list[torch.Tensor]
+    ) -> list[torch.Tensor]:
+        """AllGather a list of tensors using NCCL symmetric memory (NVLS).
+
+        Uses group_start/group_end to batch the collectives.
+        Only the output needs to be in symmetric memory (see
+        _all_gather_symm_mem).
+        """
+        from vllm.distributed.device_communicators.pynccl_allocator import (
+            nccl_symm_mem_context,
+        )
+
+        pynccl_comm = self.pynccl_comm
+        assert pynccl_comm is not None
+        world_size = self.world_size
+
+        symm_outputs = []
+        with nccl_symm_mem_context(pynccl_comm):
+            for inp in inputs:
+                out_size = (inp.size(0) * world_size,) + inp.size()[1:]
+                symm_outputs.append(
+                    torch.empty(out_size, dtype=inp.dtype, device=inp.device)
+                )
+
+        pynccl_comm.group_start()
+        for symm_out, inp in zip(symm_outputs, inputs):
+            pynccl_comm.all_gather(symm_out, inp)
+        pynccl_comm.group_end()
+
+        return symm_outputs
+
+    def dispatch_router_logits(
+        self,
+        hidden_states: torch.Tensor,
+        router_logits: torch.Tensor,
+        is_sequence_parallel: bool = False,
+        extra_tensors: list[torch.Tensor] | None = None,
+    ) -> (
+        tuple[torch.Tensor, torch.Tensor]
+        | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
+    ):
+        """
+        Dispatch the hidden states and router logits to the appropriate device.
+        This is a no-op in the base class.
+        """
+
+        assert self.all2all_manager is not None
+        return self.all2all_manager.dispatch_router_logits(
+            hidden_states,
+            router_logits,
+            is_sequence_parallel,
+            extra_tensors,
+        )
+
+    def dispatch(
+        self,
+        hidden_states: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        is_sequence_parallel: bool = False,
+        extra_tensors: list[torch.Tensor] | None = None,
+    ) -> (
+        tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+        | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
+    ):
+        """
+        Dispatch the hidden states and topk weights/ids to the appropriate device.
+        This is a no-op in the base class.
+        """
+        assert self.all2all_manager is not None
+        return self.all2all_manager.dispatch(
+            hidden_states,
+            topk_weights,
+            topk_ids,
+            is_sequence_parallel,
+            extra_tensors=extra_tensors,
+        )
+
+    def combine(
+        self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
+    ) -> torch.Tensor:
+        """
+        Combine the hidden states and router logits from the appropriate device.
+        This is a no-op in the base class.
+        """
+        assert self.all2all_manager is not None
+        return self.all2all_manager.combine(
+            hidden_states,
+            is_sequence_parallel,
+        )
+
+    def batch_isend_irecv(self, p2p_ops: list):
+        pynccl_comm = self.pynccl_comm
+        if pynccl_comm is not None and not pynccl_comm.disabled:
+            pynccl_comm.batch_isend_irecv(p2p_ops)
+        else:
+            raise ValueError("No PyNCCL communicator found")

vllm/distributed/device_communicators/cuda_wrapper.py

@@ -0,0 +1,187 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""This file is a pure Python wrapper for the cudart library.
+It avoids the need to compile a separate shared library, and is
+convenient for use when we just need to call a few functions.
+"""
+
+import ctypes
+from dataclasses import dataclass
+from typing import Any
+
+# this line makes it possible to directly load `libcudart.so` using `ctypes`
+import torch  # noqa
+
+import vllm.envs as envs
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+from vllm.utils.system_utils import find_loaded_library
+
+logger = init_logger(__name__)
+
+# === export types and functions from cudart to Python ===
+# for the original cudart definition, please check
+# https://docs.nvidia.com/cuda/cuda-runtime-api/index.html
+
+cudaError_t = ctypes.c_int
+cudaMemcpyKind = ctypes.c_int
+
+
+class cudaIpcMemHandle_t(ctypes.Structure):
+    _fields_ = [("internal", ctypes.c_byte * 128)]
+
+
+@dataclass
+class Function:
+    name: str
+    restype: Any
+    argtypes: list[Any]
+
+
+class CudaRTLibrary:
+    exported_functions = [
+        # ​cudaError_t cudaSetDevice ( int  device )
+        Function("cudaSetDevice", cudaError_t, [ctypes.c_int]),
+        # cudaError_t 	cudaDeviceSynchronize ( void )
+        Function("cudaDeviceSynchronize", cudaError_t, []),
+        # ​cudaError_t cudaDeviceReset ( void )
+        Function("cudaDeviceReset", cudaError_t, []),
+        # const char* 	cudaGetErrorString ( cudaError_t error )
+        Function("cudaGetErrorString", ctypes.c_char_p, [cudaError_t]),
+        # ​cudaError_t 	cudaMalloc ( void** devPtr, size_t size )
+        Function(
+            "cudaMalloc",
+            cudaError_t,
+            [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t],
+        ),
+        # ​cudaError_t 	cudaFree ( void* devPtr )
+        Function("cudaFree", cudaError_t, [ctypes.c_void_p]),
+        # ​cudaError_t cudaMemset ( void* devPtr, int  value, size_t count )
+        Function(
+            "cudaMemset", cudaError_t, [ctypes.c_void_p, ctypes.c_int, ctypes.c_size_t]
+        ),
+        # ​cudaError_t cudaMemcpy ( void* dst, const void* src, size_t count, cudaMemcpyKind kind ) # noqa
+        Function(
+            "cudaMemcpy",
+            cudaError_t,
+            [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, cudaMemcpyKind],
+        ),
+        # cudaError_t cudaIpcGetMemHandle ( cudaIpcMemHandle_t* handle, void* devPtr ) # noqa
+        Function(
+            "cudaIpcGetMemHandle",
+            cudaError_t,
+            [ctypes.POINTER(cudaIpcMemHandle_t), ctypes.c_void_p],
+        ),
+        # ​cudaError_t cudaIpcOpenMemHandle ( void** devPtr, cudaIpcMemHandle_t handle, unsigned int  flags ) # noqa
+        Function(
+            "cudaIpcOpenMemHandle",
+            cudaError_t,
+            [ctypes.POINTER(ctypes.c_void_p), cudaIpcMemHandle_t, ctypes.c_uint],
+        ),
+    ]
+
+    # https://rocm.docs.amd.com/projects/HIPIFY/en/latest/tables/CUDA_Runtime_API_functions_supported_by_HIP.html # noqa
+    cuda_to_hip_mapping = {
+        "cudaSetDevice": "hipSetDevice",
+        "cudaDeviceSynchronize": "hipDeviceSynchronize",
+        "cudaDeviceReset": "hipDeviceReset",
+        "cudaGetErrorString": "hipGetErrorString",
+        "cudaMalloc": "hipMalloc",
+        "cudaFree": "hipFree",
+        "cudaMemset": "hipMemset",
+        "cudaMemcpy": "hipMemcpy",
+        "cudaIpcGetMemHandle": "hipIpcGetMemHandle",
+        "cudaIpcOpenMemHandle": "hipIpcOpenMemHandle",
+    }
+
+    # class attribute to store the mapping from the path to the library
+    # to avoid loading the same library multiple times
+    path_to_library_cache: dict[str, Any] = {}
+
+    # class attribute to store the mapping from library path
+    #  to the corresponding dictionary
+    path_to_dict_mapping: dict[str, dict[str, Any]] = {}
+
+    def __init__(self, so_file: str | None = None):
+        if so_file is None:
+            so_file = (
+                find_loaded_library(
+                    "libamdhip64" if current_platform.is_rocm() else "libcudart"
+                )
+                or envs.VLLM_CUDART_SO_PATH  # fallback to env var
+            )
+            assert so_file is not None, (
+                "libcudart is not loaded in the current process, "
+                "try setting VLLM_CUDART_SO_PATH"
+            )
+        if so_file not in CudaRTLibrary.path_to_library_cache:
+            lib = ctypes.CDLL(so_file)
+            CudaRTLibrary.path_to_library_cache[so_file] = lib
+        self.lib = CudaRTLibrary.path_to_library_cache[so_file]
+
+        if so_file not in CudaRTLibrary.path_to_dict_mapping:
+            _funcs = {}
+            for func in CudaRTLibrary.exported_functions:
+                f = getattr(
+                    self.lib,
+                    CudaRTLibrary.cuda_to_hip_mapping[func.name]
+                    if current_platform.is_rocm()
+                    else func.name,
+                )
+                f.restype = func.restype
+                f.argtypes = func.argtypes
+                _funcs[func.name] = f
+            CudaRTLibrary.path_to_dict_mapping[so_file] = _funcs
+        self.funcs = CudaRTLibrary.path_to_dict_mapping[so_file]
+
+    def CUDART_CHECK(self, result: cudaError_t) -> None:
+        if result != 0:
+            error_str = self.cudaGetErrorString(result)
+            raise RuntimeError(f"CUDART error: {error_str}")
+
+    def cudaGetErrorString(self, error: cudaError_t) -> str:
+        return self.funcs["cudaGetErrorString"](error).decode("utf-8")
+
+    def cudaSetDevice(self, device: int) -> None:
+        self.CUDART_CHECK(self.funcs["cudaSetDevice"](device))
+
+    def cudaDeviceSynchronize(self) -> None:
+        self.CUDART_CHECK(self.funcs["cudaDeviceSynchronize"]())
+
+    def cudaDeviceReset(self) -> None:
+        self.CUDART_CHECK(self.funcs["cudaDeviceReset"]())
+
+    def cudaMalloc(self, size: int) -> ctypes.c_void_p:
+        devPtr = ctypes.c_void_p()
+        self.CUDART_CHECK(self.funcs["cudaMalloc"](ctypes.byref(devPtr), size))
+        return devPtr
+
+    def cudaFree(self, devPtr: ctypes.c_void_p) -> None:
+        self.CUDART_CHECK(self.funcs["cudaFree"](devPtr))
+
+    def cudaMemset(self, devPtr: ctypes.c_void_p, value: int, count: int) -> None:
+        self.CUDART_CHECK(self.funcs["cudaMemset"](devPtr, value, count))
+
+    def cudaMemcpy(
+        self, dst: ctypes.c_void_p, src: ctypes.c_void_p, count: int
+    ) -> None:
+        cudaMemcpyDefault = 4
+        kind = cudaMemcpyDefault
+        self.CUDART_CHECK(self.funcs["cudaMemcpy"](dst, src, count, kind))
+
+    def cudaIpcGetMemHandle(self, devPtr: ctypes.c_void_p) -> cudaIpcMemHandle_t:
+        handle = cudaIpcMemHandle_t()
+        self.CUDART_CHECK(
+            self.funcs["cudaIpcGetMemHandle"](ctypes.byref(handle), devPtr)
+        )
+        return handle
+
+    def cudaIpcOpenMemHandle(self, handle: cudaIpcMemHandle_t) -> ctypes.c_void_p:
+        cudaIpcMemLazyEnablePeerAccess = 1
+        devPtr = ctypes.c_void_p()
+        self.CUDART_CHECK(
+            self.funcs["cudaIpcOpenMemHandle"](
+                ctypes.byref(devPtr), handle, cudaIpcMemLazyEnablePeerAccess
+            )
+        )
+        return devPtr

vllm/distributed/device_communicators/custom_all_reduce.py

@@ -0,0 +1,323 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from contextlib import contextmanager
+from typing import cast
+
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+
+import vllm.envs as envs
+from vllm import _custom_ops as ops
+from vllm.distributed.device_communicators.all_reduce_utils import (
+    CUSTOM_ALL_REDUCE_MAX_SIZES,
+    gpu_p2p_access_check,
+)
+from vllm.distributed.parallel_state import in_the_same_node_as
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+try:
+    ops.meta_size()
+    custom_ar = True
+except Exception:
+    # For CPUs
+    custom_ar = False
+
+logger = init_logger(__name__)
+
+
+def _can_p2p(rank: int, world_size: int) -> bool:
+    for i in range(world_size):
+        if i == rank:
+            continue
+        if envs.VLLM_SKIP_P2P_CHECK:
+            logger.debug("Skipping P2P check and trusting the driver's P2P report.")
+            # can_device_access_peer takes visible device ordinals, while
+            # rank and i are logical local IDs.
+            return torch.cuda.can_device_access_peer(
+                current_platform.logical_device_id_to_visible_device_id(rank),
+                current_platform.logical_device_id_to_visible_device_id(i),
+            )
+        if not gpu_p2p_access_check(rank, i):
+            return False
+    return True
+
+
+from vllm.distributed.utils import is_weak_contiguous  # noqa: E402
+
+
+class CustomAllreduce:
+    _SUPPORTED_WORLD_SIZES = [2, 4, 6, 8]
+
+    # max_size: max supported allreduce size
+    def __init__(
+        self,
+        group: ProcessGroup,
+        device: int | str | torch.device,
+        max_size=8192 * 1024,
+        symm_mem_enabled=False,
+    ) -> None:
+        """
+        Args:
+            group: the process group to work on. If None, it will use the
+                default process group.
+            device: the device to bind the CustomAllreduce to. If None,
+                it will be bound to f"cuda:{local_rank}".
+        It is the caller's responsibility to make sure each communicator
+        is bind to a unique device, and all communicators in this group
+        are in the same node.
+        """
+        self._IS_CAPTURING = False
+        self.disabled = True
+
+        if not custom_ar:
+            # disable because of missing custom allreduce library
+            # e.g. in a non-GPU environment
+            logger.info(
+                "Custom allreduce is disabled because "
+                "of missing custom allreduce library"
+            )
+            return
+
+        self.group = group
+
+        assert dist.get_backend(group) != dist.Backend.NCCL, (
+            "CustomAllreduce should be attached to a non-NCCL group."
+        )
+
+        if not all(in_the_same_node_as(group, source_rank=0)):
+            # No need to initialize custom allreduce for multi-node case.
+            logger.warning(
+                "Custom allreduce is disabled because this process group"
+                " spans across nodes."
+            )
+            return
+
+        rank = dist.get_rank(group=self.group)
+        self.rank = rank
+        world_size = dist.get_world_size(group=self.group)
+        if world_size == 1:
+            # No need to initialize custom allreduce for single GPU case.
+            return
+
+        if world_size not in CustomAllreduce._SUPPORTED_WORLD_SIZES:
+            logger.warning(
+                "Custom allreduce is disabled due to an unsupported world"
+                " size: %d. Supported world sizes: %s. To silence this "
+                "warning, specify disable_custom_all_reduce=True explicitly.",
+                world_size,
+                str(CustomAllreduce._SUPPORTED_WORLD_SIZES),
+            )
+            return
+
+        if isinstance(device, int):
+            device = torch.device(f"cuda:{device}")
+        elif isinstance(device, str):
+            device = torch.device(device)
+        # now `device` is a `torch.device` object
+        assert isinstance(device, torch.device)
+        self.device = device
+        device_capability = current_platform.get_device_capability()
+        if (
+            current_platform.is_cuda()
+            and symm_mem_enabled
+            and device_capability is not None
+        ):
+            device_capability_str = device_capability.as_version_str()
+            if device_capability_str in CUSTOM_ALL_REDUCE_MAX_SIZES:
+                max_size = min(
+                    CUSTOM_ALL_REDUCE_MAX_SIZES[device_capability_str][world_size],
+                    max_size,
+                )
+        # device.index is a visible ordinal, not a logical local ID.
+        physical_device_id = current_platform.visible_device_id_to_physical_device_id(
+            device.index
+        )
+        tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")
+        gather_list = [
+            torch.tensor([0], dtype=torch.int, device="cpu") for _ in range(world_size)
+        ]
+        dist.all_gather(gather_list, tensor, group=self.group)
+        physical_device_ids = [t.item() for t in gather_list]
+
+        # test nvlink first, this will filter out most of the cases
+        # where custom allreduce is not supported
+        # this checks hardware and driver support for NVLink
+        assert current_platform.is_cuda_alike()
+        fully_connected = current_platform.is_fully_connected(physical_device_ids)
+        if world_size > 2 and not fully_connected:
+            logger.warning(
+                "Custom allreduce is disabled because it's not supported on"
+                " more than two PCIe-only GPUs. To silence this warning, "
+                "specify disable_custom_all_reduce=True explicitly."
+            )
+            return
+        # test P2P capability, this checks software/cudaruntime support
+        # this is expensive to compute at the first time
+        # then we cache the result
+        # On AMD GPU, p2p is always enabled between XGMI connected GPUs
+        if not current_platform.is_rocm() and not _can_p2p(rank, world_size):
+            logger.warning(
+                "Custom allreduce is disabled because your platform lacks "
+                "GPU P2P capability or P2P test failed. To silence this "
+                "warning, specify disable_custom_all_reduce=True explicitly."
+            )
+            return
+
+        self.disabled = False
+        # Buffers memory are owned by this Python class and passed to C++.
+        # Metadata composes of two parts: metadata for synchronization and a
+        # temporary buffer for storing intermediate allreduce results.
+        self.meta_ptrs = self.create_shared_buffer(
+            ops.meta_size() + max_size, group=group, uncached=True
+        )
+        # This is a pre-registered IPC buffer. In eager mode, input tensors
+        # are first copied into this buffer before allreduce is performed
+        self.buffer_ptrs = self.create_shared_buffer(max_size, group=group)
+        # This is a buffer for storing the tuples of pointers pointing to
+        # IPC buffers from all ranks. Each registered tuple has size of
+        # 8*world_size bytes where world_size is at most 8. Allocating 8MB
+        # is enough for 131072 such tuples. The largest model I've seen only
+        # needs less than 10000 of registered tuples.
+        self.rank_data = torch.empty(
+            8 * 1024 * 1024, dtype=torch.uint8, device=self.device
+        )
+        self.max_size = max_size
+        self.rank = rank
+        self.world_size = world_size
+        self.fully_connected = fully_connected
+        self._ptr = ops.init_custom_ar(
+            self.meta_ptrs, self.rank_data, rank, self.fully_connected
+        )
+        ops.register_buffer(self._ptr, self.buffer_ptrs)
+
+    @contextmanager
+    def capture(self):
+        """
+        The main responsibility of this context manager is the
+        `register_graph_buffers` call at the end of the context.
+        It records all the buffer addresses used in the CUDA graph.
+        """
+        try:
+            self._IS_CAPTURING = True
+            yield
+        finally:
+            self._IS_CAPTURING = False
+            if not self.disabled:
+                self.register_graph_buffers()
+
+    def register_graph_buffers(self):
+        handle, offset = ops.get_graph_buffer_ipc_meta(self._ptr)
+        logger.info("Registering %d cuda graph addresses", len(offset))
+        # We cannot directly use `dist.all_gather_object` here
+        # because it is incompatible with `gloo` backend under inference mode.
+        # see https://github.com/pytorch/pytorch/issues/126032 for details.
+        all_data: list[list[list[int] | None]]
+        all_data = [[None, None] for _ in range(dist.get_world_size(group=self.group))]
+        all_data[self.rank] = [handle, offset]
+        ranks = sorted(dist.get_process_group_ranks(group=self.group))
+        for i, rank in enumerate(ranks):
+            dist.broadcast_object_list(
+                all_data[i], src=rank, group=self.group, device="cpu"
+            )
+        # Unpack list of tuples to tuple of lists.
+        handles = cast(list[list[int]], [d[0] for d in all_data])
+        offsets = cast(list[list[int]], [d[1] for d in all_data])
+        ops.register_graph_buffers(self._ptr, handles, offsets)
+
+    def should_custom_ar(self, inp: torch.Tensor):
+        if self.disabled:
+            return False
+        inp_size = inp.numel() * inp.element_size()
+        # custom allreduce requires input byte size to be multiples of 16
+        if inp_size % 16 != 0:
+            return False
+        if not is_weak_contiguous(inp):
+            return False
+        # for 4 or more non NVLink-capable GPUs, custom allreduce provides
+        # little performance improvement over NCCL.
+        if self.world_size == 2 or self.fully_connected:
+            return inp_size < self.max_size
+        return False
+
+    def all_reduce(
+        self, inp: torch.Tensor, *, out: torch.Tensor = None, registered: bool = False
+    ):
+        """Performs an out-of-place all reduce.
+
+        If registered is True, this assumes inp's pointer is already
+        IPC-registered. Otherwise, inp is first copied into a pre-registered
+        buffer.
+        """
+        if out is None:
+            out = torch.empty_like(inp)
+        if registered:
+            ops.all_reduce(self._ptr, inp, out, 0, 0)
+        else:
+            ops.all_reduce(
+                self._ptr, inp, out, self.buffer_ptrs[self.rank], self.max_size
+            )
+        return out
+
+    def custom_all_reduce(self, input: torch.Tensor) -> torch.Tensor | None:
+        """The main allreduce API that provides support for cuda graph."""
+        # When custom allreduce is disabled, this will be None.
+        if self.disabled or not self.should_custom_ar(input):
+            return None
+        if self._IS_CAPTURING:
+            if torch.cuda.is_current_stream_capturing():
+                return self.all_reduce(input, registered=True)
+            else:
+                # If warm up, mimic the allocation pattern since custom
+                # allreduce is out-of-place.
+                return torch.empty_like(input)
+        else:
+            # Note: outside of cuda graph context, custom allreduce incurs a
+            # cost of cudaMemcpy, which should be small (<=1% of overall
+            # latency) compared to the performance gain of using custom kernels
+            return self.all_reduce(input, registered=False)
+
+    def close(self):
+        if not self.disabled and self._ptr:
+            if ops is not None:
+                ops.dispose(self._ptr)
+            self._ptr = 0
+            self.free_shared_buffer(self.meta_ptrs, rank=self.rank)
+            self.free_shared_buffer(self.buffer_ptrs, rank=self.rank)
+
+    def __del__(self):
+        self.close()
+
+    @staticmethod
+    def create_shared_buffer(
+        size_in_bytes: int,
+        group: ProcessGroup | None = None,
+        uncached: bool | None = False,
+    ) -> list[int]:
+        pointer, handle = ops.allocate_shared_buffer_and_handle(size_in_bytes)
+
+        world_size = dist.get_world_size(group=group)
+        rank = dist.get_rank(group=group)
+        handles = [None] * world_size
+        dist.all_gather_object(handles, handle, group=group)
+
+        pointers: list[int] = []
+        for i, h in enumerate(handles):
+            if i == rank:
+                pointers.append(pointer)  # type: ignore
+            else:
+                pointers.append(ops.open_mem_handle(h))
+        return pointers
+
+    @staticmethod
+    def free_shared_buffer(
+        pointers: list[int],
+        group: ProcessGroup | None = None,
+        rank: int | None = None,
+    ) -> None:
+        if rank is None:
+            rank = dist.get_rank(group=group)
+        if ops is not None:
+            ops.free_shared_buffer(pointers[rank])

vllm/distributed/device_communicators/flashinfer_all_reduce.py

@@ -0,0 +1,353 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import atexit
+import os
+import random
+import threading
+
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+
+import vllm.envs as envs
+from vllm.config.compilation import PassConfig
+from vllm.distributed.parallel_state import get_node_count
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+logger = init_logger(__name__)
+
+# The empirical value for small batch
+PDL_ADVANCE_LAUNCH_TOKENS = 16
+
+fi_ar_available = False
+try:
+    import flashinfer.comm as flashinfer_comm  # type: ignore[no-redef]
+    from flashinfer.comm.mnnvl import (
+        TorchDistBackend,  # type: ignore[import-not-found, no-redef]
+    )
+
+    fi_ar_available = hasattr(flashinfer_comm, "allreduce_fusion")
+except ImportError:
+    pass
+
+# Workspace for standalone allreduce and non-quant ar+rms fusion
+_fi_ar_workspace = None
+# Extra workspace for quant fusion patterns (only supported by trtllm backend)
+_fi_ar_quant_workspace = None
+
+
+def _create_workspace(
+    backend: str,
+    world_size: int,
+    rank: int,
+    max_token_num: int,
+    hidden_dim: int,
+    dtype: torch.dtype,
+    group: ProcessGroup,
+):
+    """Create a flashinfer allreduce workspace, returning None on failure."""
+    comm_backend = TorchDistBackend(group=group)
+    rng_state = random.getstate()
+    try:
+        random.seed(int.from_bytes(os.urandom(16), byteorder="big"))
+        workspace = flashinfer_comm.create_allreduce_fusion_workspace(
+            backend=backend,
+            world_size=world_size,
+            rank=rank,
+            max_token_num=max_token_num,
+            hidden_dim=hidden_dim,
+            dtype=dtype,
+            comm_backend=comm_backend,
+            group=group,
+        )
+    except Exception as e:
+        if "multicast" in str(e).lower():
+            logger.warning_once(
+                "Failed to initialize FlashInfer All Reduce workspace: %s. "
+                "This is expected on GPUs without NVSwitch (e.g., NVLink "
+                "bridge-only or PCIe topologies).",
+                e,
+            )
+        else:
+            logger.warning_once(
+                "Failed to initialize FlashInfer All Reduce workspace: %s.",
+                e,
+            )
+        return None
+    finally:
+        random.setstate(rng_state)
+    logger.debug(
+        "Initialized FlashInfer All Reduce workspace: backend=%s, "
+        "world_size=%d, rank=%d, max_token_num=%d, hidden_dim=%d, dtype=%s",
+        backend,
+        world_size,
+        rank,
+        max_token_num,
+        hidden_dim,
+        dtype,
+    )
+    return workspace
+
+
+def _resolve_fi_ar_backend() -> tuple[str, bool]:
+    """Resolve the flashinfer allreduce backend for the current setup.
+
+    Returns:
+        A ``(backend, allow_trtllm_fallback)`` tuple. ``allow_trtllm_fallback``
+        is True only when ``auto`` selects mnnvl for a single node, so that
+        workspace creation can fall back to trtllm on single-node topologies
+        without NVSwitch multicast support (where mnnvl is unavailable).
+    """
+    backend = envs.VLLM_FLASHINFER_ALLREDUCE_BACKEND
+    if backend != "auto":
+        logger.info_once(f"Using flashinfer allreduce backend: {backend}")
+        return backend, False
+
+    # Default to mnnvl for both single- and multi-node setups. The mnnvl
+    # cudagraph hang that previously forced single-node to trtllm
+    # (https://github.com/vllm-project/vllm/issues/35772) was fixed upstream in
+    # FlashInfer (>= 0.6.12, vLLM pins 0.6.13), so mnnvl is safe here. trtllm
+    # does not support multi-node allreduce, so mnnvl is required there anyway.
+    # mnnvl needs NVSwitch multicast; on single-node topologies without it,
+    # fall back to trtllm so fused allreduce stays enabled.
+    backend = "mnnvl"
+    allow_trtllm_fallback = get_node_count() == 1
+
+    logger.info_once(f"Auto-selected flashinfer allreduce backend: {backend}")
+    return backend, allow_trtllm_fallback
+
+
+def get_fi_ar_workspace(
+    world_size: int,
+    rank: int,
+    max_token_num: int,
+    hidden_dim: int,
+    dtype: torch.dtype,
+    group: ProcessGroup,
+):
+    """
+    Return the allreduce workspace for non-quant patterns, initializing if needed.
+
+    Used by AllReduceFusionPass (non-quant patterns) and FlashInferAllReduce
+    for standalone allreduce. Backend is controlled by
+    VLLM_FLASHINFER_ALLREDUCE_BACKEND env var.
+    """
+    global _fi_ar_workspace
+    if _fi_ar_workspace is not None:
+        return _fi_ar_workspace
+
+    backend, allow_trtllm_fallback = _resolve_fi_ar_backend()
+
+    if get_node_count() > 1 and backend == "trtllm":
+        raise ValueError(
+            "Flashinfer allreduce is not supported for multi-node allreduce with "
+            "'trtllm' backend. Please use 'mnnvl' backend instead."
+        )
+
+    def _get_or_create(be: str):
+        # Reuse the quant workspace if it was already created with the same backend
+        if _fi_ar_quant_workspace is not None and _fi_ar_quant_workspace.backend == be:
+            return _fi_ar_quant_workspace
+        return _create_workspace(
+            be, world_size, rank, max_token_num, hidden_dim, dtype, group
+        )
+
+    _fi_ar_workspace = _get_or_create(backend)
+    if _fi_ar_workspace is None and allow_trtllm_fallback and backend != "trtllm":
+        logger.warning_once(
+            "FlashInfer mnnvl allreduce workspace unavailable (likely no NVSwitch "
+            "multicast support); falling back to trtllm backend for single node."
+        )
+        backend = "trtllm"
+        _fi_ar_workspace = _get_or_create(backend)
+
+    if _fi_ar_workspace is not None:
+        logger.info_once(
+            "Initialized FlashInfer Allreduce norm fusion workspace "
+            f"with backend={backend}"
+        )
+    else:
+        logger.warning_once(
+            "Failed to initialize FlashInfer Allreduce norm fusion workspace "
+            f"with backend={backend}"
+        )
+
+    return _fi_ar_workspace
+
+
+def get_fi_ar_quant_workspace(
+    world_size: int,
+    rank: int,
+    max_token_num: int,
+    hidden_dim: int,
+    dtype: torch.dtype,
+    group: ProcessGroup,
+):
+    """
+    Return the allreduce workspace for quant patterns, initializing if needed.
+
+    Always uses trtllm backend as it is the only one supporting quantization
+    fusion (FP8/FP4). Returns None for multi-node setups since not supported
+    by trtllm backend.
+    """
+    global _fi_ar_quant_workspace
+    if _fi_ar_quant_workspace is not None:
+        return _fi_ar_quant_workspace
+
+    if get_node_count() > 1:
+        logger.warning_once(
+            "Flashinfer allreduce quantization fusion is not supported for "
+            "multi-node allreduce. Disabling quant fusion."
+        )
+        return None
+
+    # Reuse the non-quant workspace if it was already created with trtllm
+    if _fi_ar_workspace is not None and _fi_ar_workspace.backend == "trtllm":
+        _fi_ar_quant_workspace = _fi_ar_workspace
+        return _fi_ar_quant_workspace
+
+    _fi_ar_quant_workspace = _create_workspace(
+        "trtllm", world_size, rank, max_token_num, hidden_dim, dtype, group
+    )
+    if _fi_ar_quant_workspace is not None:
+        logger.info_once(
+            "Initialized FlashInfer Allreduce norm quantization "
+            "fusion workspace with backend=trtllm"
+        )
+    else:
+        logger.warning_once(
+            "Failed to initialize FlashInfer Allreduce norm quantization "
+            "fusion workspace with backend=trtllm"
+        )
+
+    return _fi_ar_quant_workspace
+
+
+_fi_ar_workspace_lock = threading.Lock()
+
+
+def destroy_fi_ar_workspace():
+    global _fi_ar_workspace, _fi_ar_quant_workspace
+    with _fi_ar_workspace_lock:
+        is_alias = _fi_ar_workspace is _fi_ar_quant_workspace
+
+        if _fi_ar_workspace is not None:
+            _fi_ar_workspace.destroy()
+        if _fi_ar_quant_workspace is not None and not is_alias:
+            _fi_ar_quant_workspace.destroy()
+
+        _fi_ar_workspace = _fi_ar_quant_workspace = None
+
+
+atexit.register(destroy_fi_ar_workspace)
+
+
+class FlashInferAllReduce:
+    def __init__(
+        self,
+        group: ProcessGroup,
+        device: int | str | torch.device,
+    ):
+        self.disabled = True
+
+        if not fi_ar_available:
+            logger.info(
+                "FlashInfer All Reduce is disabled because flashinfer is not available"
+            )
+            return
+
+        if not current_platform.is_cuda():
+            logger.info(
+                "FlashInfer All Reduce is disabled because it requires CUDA platform"
+            )
+            return
+
+        self.group = group
+        self.world_size = dist.get_world_size(self.group)
+        self.rank = dist.get_rank(self.group)
+        self.device = device
+        if self.world_size == 1:
+            return
+
+        # Use the same threshold as the allreduce-rms fusion pass
+        # TODO: tune the threshold
+        MiB = 1024 * 1024
+        max_workspace_size = PassConfig.default_fi_allreduce_fusion_max_size_mb().get(
+            self.world_size, None
+        )
+        if not max_workspace_size:
+            logger.warning(
+                "FlashInfer All Reduce is disabled because it "
+                "is not supported for world_size=%d.",
+                self.world_size,
+            )
+            return
+        self.max_workspace_size = max_workspace_size * MiB
+        self.max_num_tokens = 0
+        self.disabled = False
+
+    def _ensure_workspace(self, hidden_dim: int, dtype: torch.dtype) -> bool:
+        """Ensure the all reduce workspace is initialized."""
+        if self.max_num_tokens == 0:
+            element_size = torch.tensor([], dtype=dtype, device="cpu").element_size()
+            self.max_num_tokens = self.max_workspace_size // (hidden_dim * element_size)
+        workspace = get_fi_ar_workspace(
+            world_size=self.world_size,
+            rank=self.rank,
+            max_token_num=self.max_num_tokens,
+            hidden_dim=hidden_dim,
+            dtype=dtype,
+            group=self.group,
+        )
+        if workspace is None:
+            self.disabled = True
+            return False
+        return True
+
+    def should_use_fi_ar(self, input_tensor: torch.Tensor) -> bool:
+        if self.disabled:
+            return False
+
+        if not input_tensor.is_cuda:
+            return False
+
+        if not input_tensor.is_contiguous():
+            return False
+
+        if len(input_tensor.shape) != 2:
+            return False
+
+        num_tokens, hidden_dim = input_tensor.shape
+        if not self.max_num_tokens:
+            element_size = torch.tensor([], dtype=input_tensor.dtype).element_size()
+            self.max_num_tokens = self.max_workspace_size // (hidden_dim * element_size)
+
+        if num_tokens > self.max_num_tokens:
+            return False
+
+        return self._ensure_workspace(hidden_dim, input_tensor.dtype)
+
+    def all_reduce(self, input_tensor: torch.Tensor) -> torch.Tensor:
+        num_tokens, hidden_dim = input_tensor.shape
+        workspace = get_fi_ar_workspace(
+            world_size=self.world_size,
+            rank=self.rank,
+            max_token_num=self.max_num_tokens,
+            hidden_dim=hidden_dim,
+            dtype=input_tensor.dtype,
+            group=self.group,
+        )
+        return flashinfer_comm.allreduce_fusion(
+            input=input_tensor,
+            workspace=workspace,
+            pattern=flashinfer_comm.AllReduceFusionPattern.kAllReduce,
+            launch_with_pdl=True,
+            trigger_completion_at_end=num_tokens > PDL_ADVANCE_LAUNCH_TOKENS,
+        )
+
+    def destroy(self):
+        if not self.disabled:
+            destroy_fi_ar_workspace()

vllm/distributed/device_communicators/mnnvl_compat.py

@@ -0,0 +1,38 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import Any
+
+import torch.distributed as dist
+from flashinfer.comm.mnnvl import CommBackend as CommBackend
+
+from vllm.utils.flashinfer import has_flashinfer_nvlink_two_sided
+
+assert has_flashinfer_nvlink_two_sided(), "Flashinfer alltoallv module cannot be found"
+
+
+class CustomCommunicator(CommBackend):
+    def __init__(self, group):
+        self._group = group
+
+    def Get_rank(self) -> int:
+        return self._group.rank()
+
+    def Get_size(self) -> int:
+        return self._group.size()
+
+    def allgather(self, data: int):
+        gathered = [None] * self.Get_size()
+        dist.all_gather_object(gathered, data, group=self._group)
+        return gathered
+
+    def bcast(self, data: Any, root: int) -> Any:
+        obj_list = [data]
+        # broadcast_object_list mutates obj_list in-place
+        dist.broadcast_object_list(obj_list, src=root, group=self._group)
+        return obj_list[0]
+
+    def barrier(self) -> None:
+        dist.barrier(group=self._group)
+
+    def Split(self, color: int, key: int) -> "CustomCommunicator":
+        return self

vllm/distributed/device_communicators/pynccl.py

@@ -0,0 +1,434 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+# ===================== import region =====================
+import threading
+
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup, ReduceOp
+
+import vllm.envs as envs
+from vllm.distributed.device_communicators.pynccl_wrapper import (
+    NCCLLibrary,
+    buffer_type,
+    cudaStream_t,
+    ncclComm_t,
+    ncclDataTypeEnum,
+    ncclRedOpTypeEnum,
+    ncclUniqueId,
+)
+from vllm.distributed.utils import StatelessProcessGroup
+from vllm.logger import init_logger
+from vllm.utils.torch_utils import current_stream
+
+logger = init_logger(__name__)
+
+_NCCL_SYMM_OPS_REGISTERED = False
+
+
+def register_nccl_symmetric_ops(pynccl_comm):
+    from vllm.distributed.device_communicators.pynccl_allocator import (
+        nccl_symm_mem_context,
+    )
+    from vllm.utils.torch_utils import direct_register_custom_op
+
+    global _NCCL_SYMM_OPS_REGISTERED
+    if _NCCL_SYMM_OPS_REGISTERED:
+        return
+    _NCCL_SYMM_OPS_REGISTERED = True
+
+    def all_reduce_symmetric_with_copy_impl(input_tensor: torch.Tensor) -> torch.Tensor:
+        with nccl_symm_mem_context(pynccl_comm):
+            symm_input = torch.empty_like(input_tensor)
+            symm_output = torch.empty_like(input_tensor)
+        symm_input.copy_(input_tensor)
+        symm_output = pynccl_comm.all_reduce(symm_input, symm_output)
+        return symm_output
+
+    def all_reduce_symmetric_with_copy_fake(input_tensor: torch.Tensor) -> torch.Tensor:
+        return torch.empty_like(input_tensor)
+
+    direct_register_custom_op(
+        op_name="all_reduce_symmetric_with_copy",
+        op_func=all_reduce_symmetric_with_copy_impl,
+        fake_impl=all_reduce_symmetric_with_copy_fake,
+    )
+
+
+class PyNcclCommunicator:
+    def __init__(
+        self,
+        group: ProcessGroup | StatelessProcessGroup,
+        device: int | str | torch.device,
+        library_path: str | None = None,
+    ):
+        """
+        Args:
+            group: the process group to work on. If None, it will use the
+                default process group.
+            device: the device to bind the PyNcclCommunicator to. If None,
+                it will be bound to f"cuda:{local_rank}".
+            library_path: the path to the NCCL library. If None, it will
+                use the default library path.
+        It is the caller's responsibility to make sure each communicator
+        is bind to a unique device.
+        """
+        if not isinstance(group, StatelessProcessGroup):
+            assert dist.is_initialized()
+            assert dist.get_backend(group) != dist.Backend.NCCL, (
+                "PyNcclCommunicator should be attached to a non-NCCL group."
+            )
+            # note: this rank is the rank in the group
+            self.rank = dist.get_rank(group)
+            self.world_size = dist.get_world_size(group)
+        else:
+            self.rank = group.rank
+            self.world_size = group.world_size
+
+        self.group = group
+
+        # if world_size == 1, no need to create communicator
+        if self.world_size == 1 or envs.VLLM_DISABLE_PYNCCL:
+            self.available = False
+            self.disabled = True
+            return
+        try:
+            self.nccl = NCCLLibrary(library_path)
+        except Exception:
+            # disable because of missing NCCL library
+            # e.g. in a non-GPU environment
+            self.available = False
+            self.disabled = True
+            return
+
+        self.available = True
+        self.disabled = False
+
+        self.nccl_version = self.nccl.ncclGetRawVersion()
+        if self.rank == 0:
+            # get the unique id from NCCL
+            self.unique_id = self.nccl.ncclGetUniqueId()
+            logger.info_once("vLLM is using nccl==%s", self.nccl.ncclGetVersion())
+        else:
+            # construct an empty unique id
+            self.unique_id = ncclUniqueId()
+
+        if not isinstance(group, StatelessProcessGroup):
+            tensor = torch.ByteTensor(list(self.unique_id.internal))
+            ranks = dist.get_process_group_ranks(group)
+            # arg `src` in `broadcast` is the global rank
+            dist.broadcast(tensor, src=ranks[0], group=group)
+            byte_list = tensor.tolist()
+            for i, byte in enumerate(byte_list):
+                self.unique_id.internal[i] = byte
+        else:
+            self.unique_id = group.broadcast_obj(self.unique_id, src=0)
+        if isinstance(device, int):
+            device = torch.device(f"cuda:{device}")
+        elif isinstance(device, str):
+            device = torch.device(device)
+        # now `device` is a `torch.device` object
+        assert isinstance(device, torch.device)
+        self.device = device
+        # nccl communicator and stream will use this device
+        with torch.accelerator.device_index(device.index):
+            self.comm: ncclComm_t = self.nccl.ncclCommInitRank(
+                self.world_size, self.unique_id, self.rank
+            )
+
+            stream = current_stream()
+            # A small all_reduce for warmup.
+            data = torch.zeros(1, device=device)
+            self.all_reduce(data)
+            stream.synchronize()
+            del data
+
+    def destroy(self):
+        if self.available and not self.disabled:
+            # ncclCommAbort can block until all CUDA graphs that
+            # captured NCCL ops on this comm are destroyed — and
+            # those graphs are released later in this same main-
+            # thread teardown, so a direct call here self-deadlocks.
+            # Run it in a daemon thread with a timeout: the main
+            # thread proceeds, the graphs drop, and the abort returns.
+            def _abort():
+                with torch.accelerator.device_index(self.device.index):
+                    self.nccl.ncclCommAbort(self.comm)
+
+            abort_thread = threading.Thread(target=_abort, daemon=True)
+            abort_thread.start()
+            abort_thread.join(timeout=5.0)
+            self.available = False
+            self.disabled = True
+
+    def all_reduce(
+        self,
+        in_tensor: torch.Tensor,
+        out_tensor: torch.Tensor = None,
+        op: ReduceOp = ReduceOp.SUM,
+        stream=None,
+    ) -> torch.Tensor:
+        if self.disabled:
+            return None
+        # nccl communicator created on a specific device
+        # will only work on tensors on the same device
+        # otherwise it will cause "illegal memory access"
+        assert in_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {in_tensor.device}"
+        )
+
+        if out_tensor is None:
+            out_tensor = torch.empty_like(in_tensor)
+
+        if stream is None:
+            stream = current_stream()
+        self.nccl.ncclAllReduce(
+            buffer_type(in_tensor.data_ptr()),
+            buffer_type(out_tensor.data_ptr()),
+            in_tensor.numel(),
+            ncclDataTypeEnum.from_torch(in_tensor.dtype),
+            ncclRedOpTypeEnum.from_torch(op),
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )
+        return out_tensor
+
+    def all_gather(
+        self, output_tensor: torch.Tensor, input_tensor: torch.Tensor, stream=None
+    ):
+        if self.disabled:
+            return
+        # nccl communicator created on a specific device
+        # will only work on tensors on the same device
+        # otherwise it will cause "illegal memory access"
+        assert input_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {input_tensor.device}"
+        )
+        if stream is None:
+            stream = current_stream()
+        self.nccl.ncclAllGather(
+            buffer_type(input_tensor.data_ptr()),
+            buffer_type(output_tensor.data_ptr()),
+            input_tensor.numel(),
+            ncclDataTypeEnum.from_torch(input_tensor.dtype),
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )
+
+    def all_gatherv(
+        self,
+        output_tensor: torch.Tensor,
+        input_tensor: torch.Tensor,
+        sizes: list[int],
+        stream=None,
+    ):
+        if self.disabled:
+            return
+        # nccl communicator created on a specific device
+        # will only work on tensors on the same device
+        # otherwise it will cause "illegal memory access"
+        assert input_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {input_tensor.device}"
+        )
+        if stream is None:
+            stream = current_stream()
+        assert output_tensor.shape[0] == sum(sizes)
+        split_offset = 0
+        self.nccl.ncclGroupStart()
+        for root, split_size in enumerate(sizes):
+            dst_slice = output_tensor[split_offset : split_offset + split_size]
+            self.nccl.ncclBroadcast(
+                buffer_type(input_tensor.data_ptr()),
+                buffer_type(dst_slice.data_ptr()),
+                dst_slice.numel(),
+                ncclDataTypeEnum.from_torch(input_tensor.dtype),
+                root,
+                self.comm,
+                cudaStream_t(stream.cuda_stream),
+            )
+            split_offset += split_size
+        self.nccl.ncclGroupEnd()
+
+    def reduce_scatter(
+        self,
+        output_tensor: torch.Tensor,
+        input_tensor: torch.Tensor,
+        op: ReduceOp = ReduceOp.SUM,
+        stream=None,
+    ):
+        if self.disabled:
+            return
+        # nccl communicator created on a specific device
+        # will only work on tensors on the same device
+        # otherwise it will cause "illegal memory access"
+        assert input_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {input_tensor.device}"
+        )
+        if stream is None:
+            stream = current_stream()
+        self.nccl.ncclReduceScatter(
+            buffer_type(input_tensor.data_ptr()),
+            buffer_type(output_tensor.data_ptr()),
+            output_tensor.numel(),
+            ncclDataTypeEnum.from_torch(input_tensor.dtype),
+            ncclRedOpTypeEnum.from_torch(op),
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )
+
+    def reduce_scatterv(
+        self,
+        output_tensor: torch.Tensor,
+        input_tensor: torch.Tensor,
+        sizes: list[int],
+        op: ReduceOp = ReduceOp.SUM,
+        stream=None,
+    ):
+        if self.disabled:
+            return
+        # nccl communicator created on a specific device
+        # will only work on tensors on the same device
+        # otherwise it will cause "illegal memory access"
+        assert input_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {input_tensor.device}"
+        )
+        if stream is None:
+            stream = current_stream()
+
+        split_offset = 0
+        self.nccl.ncclGroupStart()
+        for root, split_size in enumerate(sizes):
+            chunk = input_tensor[split_offset : split_offset + split_size, ...]
+            self.nccl.ncclReduce(
+                buffer_type(chunk.data_ptr()),
+                buffer_type(output_tensor.data_ptr()),
+                chunk.numel(),
+                ncclDataTypeEnum.from_torch(input_tensor.dtype),
+                ncclRedOpTypeEnum.from_torch(op),
+                root,
+                self.comm,
+                cudaStream_t(stream.cuda_stream),
+            )
+            split_offset += split_size
+        self.nccl.ncclGroupEnd()
+
+    def send(self, tensor: torch.Tensor, dst: int, stream=None):
+        if self.disabled:
+            return
+        assert tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {tensor.device}"
+        )
+        if stream is None:
+            stream = current_stream()
+        if tensor.dtype in [
+            torch.float8_e5m2,
+            torch.float8_e4m3fn,
+            torch.float8_e4m3fnuz,
+            torch.float8_e5m2fnuz,
+        ]:
+            nccl_dtype = ncclDataTypeEnum.from_torch(torch.uint8)
+        else:
+            nccl_dtype = ncclDataTypeEnum.from_torch(tensor.dtype)
+        self.nccl.ncclSend(
+            buffer_type(tensor.data_ptr()),
+            tensor.numel(),
+            nccl_dtype,
+            dst,
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )
+
+    def recv(self, tensor: torch.Tensor, src: int, stream=None):
+        if self.disabled:
+            return
+        assert tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {tensor.device}"
+        )
+        if stream is None:
+            stream = current_stream()
+        if tensor.dtype in [
+            torch.float8_e5m2,
+            torch.float8_e4m3fn,
+            torch.float8_e4m3fnuz,
+            torch.float8_e5m2fnuz,
+        ]:
+            nccl_dtype = ncclDataTypeEnum.from_torch(torch.uint8)
+        else:
+            nccl_dtype = ncclDataTypeEnum.from_torch(tensor.dtype)
+        self.nccl.ncclRecv(
+            buffer_type(tensor.data_ptr()),
+            tensor.numel(),
+            nccl_dtype,
+            src,
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )
+
+    def broadcast(self, tensor: torch.Tensor, src: int, stream=None):
+        if self.disabled:
+            return
+        assert tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {tensor.device}"
+        )
+        if stream is None:
+            stream = current_stream()
+        if src == self.rank:
+            sendbuff = buffer_type(tensor.data_ptr())
+            # NCCL requires the sender also to have a receive buffer
+            recvbuff = buffer_type(tensor.data_ptr())
+        else:
+            sendbuff = buffer_type()
+            recvbuff = buffer_type(tensor.data_ptr())
+        self.nccl.ncclBroadcast(
+            sendbuff,
+            recvbuff,
+            tensor.numel(),
+            ncclDataTypeEnum.from_torch(tensor.dtype),
+            src,
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )
+
+    def group_start(self):
+        self.nccl.ncclGroupStart()
+
+    def group_end(self):
+        self.nccl.ncclGroupEnd()
+
+    def register_comm_window(self, tensor: torch.Tensor):
+        return self.nccl.ncclCommWindowRegister(
+            self.comm,
+            buffer_type(tensor.data_ptr()),
+            tensor.numel() * tensor.element_size(),
+            1,
+        )
+
+    def register_comm_window_raw(self, ptr: int, size: int):
+        return self.nccl.ncclCommWindowRegister(self.comm, buffer_type(ptr), size, 1)
+
+    def deregister_comm_window(self, window):
+        return self.nccl.ncclCommWindowDeregister(self.comm, window)
+
+    def batch_isend_irecv(self, p2p_ops: list, stream=None):
+        if self.disabled:
+            return
+        if stream is None:
+            stream = current_stream()
+        self.group_start()
+        for op in p2p_ops:
+            if op.op is torch.distributed.isend:
+                self.send(op.tensor, op.group_peer, stream)
+            elif op.op is torch.distributed.irecv:
+                self.recv(op.tensor, op.group_peer, stream)
+
+        self.group_end()

vllm/distributed/device_communicators/pynccl_allocator.py

@@ -0,0 +1,194 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import atexit
+import contextlib
+import tempfile
+from typing import Any
+
+import torch
+from packaging import version
+from torch.cuda.memory import CUDAPluggableAllocator
+from torch.utils.cpp_extension import load_inline
+
+from vllm import envs
+from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+from vllm.utils.nccl import find_nccl_include_paths
+
+logger = init_logger(__name__)
+
+nccl_allocator_source = """
+#include <nccl.h>
+extern "C" {
+
+void* nccl_alloc_plug(size_t size, int device, void* stream) {
+  void* ptr;
+  ncclResult_t err = ncclMemAlloc(&ptr, size);
+  return ptr;
+
+}
+
+void nccl_free_plug(void* ptr, size_t size, int device, void* stream) {
+  ncclResult_t err = ncclMemFree(ptr);
+}
+
+}
+"""
+
+_allocator = None
+_allocator_wrapper = None
+_mem_pool = None
+_registered_base_addrs: dict[bytes, set] = {}
+_graph_pool_id = None
+_nccl_allocator_failed_to_compile = False
+_cached_pool_snapshot = None
+
+
+def is_symmetric_memory_enabled():
+    global _nccl_allocator_failed_to_compile
+    return envs.VLLM_USE_NCCL_SYMM_MEM and not _nccl_allocator_failed_to_compile
+
+
+def is_symmetric_memory_tensor(tensor: torch.Tensor):
+    if not is_symmetric_memory_enabled() or _cached_pool_snapshot is None:
+        return False
+    for segment in _cached_pool_snapshot:
+        for block in segment["blocks"]:
+            if block["address"] == tensor.untyped_storage().data_ptr():
+                return True
+    return False
+
+
+def set_graph_pool_id(graph_pool_id: Any) -> None:
+    global _graph_pool_id
+    _graph_pool_id = graph_pool_id
+
+
+def compile_nccl_allocator():
+    global _allocator, _allocator_wrapper, _nccl_allocator_failed_to_compile
+    if not current_platform.is_cuda():
+        _nccl_allocator_failed_to_compile = True
+        return
+    try:
+        out_dir = tempfile.gettempdir()
+        nccl_allocator_libname = "nccl_allocator"
+        nccl_include_paths = find_nccl_include_paths()
+        load_inline(
+            name=nccl_allocator_libname,
+            cpp_sources=nccl_allocator_source,
+            with_cuda=True,
+            extra_ldflags=["-lnccl"],
+            verbose=envs.VLLM_LOGGING_LEVEL == "DEBUG",
+            is_python_module=False,
+            build_directory=out_dir,
+            extra_include_paths=nccl_include_paths,
+        )
+        _allocator_wrapper = CUDAPluggableAllocator(
+            f"{out_dir}/{nccl_allocator_libname}.so",
+            "nccl_alloc_plug",
+            "nccl_free_plug",
+        )
+        _allocator = _allocator_wrapper.allocator()
+    except Exception as e:
+        _nccl_allocator_failed_to_compile = True
+        logger.warning(
+            "Failed to compile NCCL memory allocator. "
+            "Symmetric memory will be disabled. "
+            "This is expected if NCCL headers are not available. "
+            "optionally set VLLM_NCCL_INCLUDE_PATH to point to a directory "
+            "containing the NCCL header. "
+            "Error: %s",
+            str(e),
+        )
+
+
+def get_nccl_mem_pool():
+    global _mem_pool, _nccl_allocator_failed_to_compile
+    if _mem_pool is None and not _nccl_allocator_failed_to_compile:
+        compile_nccl_allocator()
+        if _allocator is not None:
+            _mem_pool = torch.cuda.MemPool(_allocator)
+    return _mem_pool
+
+
+def _cleanup_nccl_mem_pool():
+    global _mem_pool
+    _mem_pool = None
+
+
+def _cleanup_nccl_allocator_wrapper():
+    global _allocator_wrapper
+    _allocator_wrapper = None
+
+
+atexit.register(_cleanup_nccl_mem_pool)
+atexit.register(_cleanup_nccl_allocator_wrapper)
+
+
+class nccl_symm_mem_context:
+    def __init__(
+        self,
+        pynccl_comm: PyNcclCommunicator,
+        disabled: bool = False,
+    ):
+        self.disabled = (
+            disabled
+            or not is_symmetric_memory_enabled()
+            or pynccl_comm.world_size == 1
+            or not current_platform.is_cuda()
+            or get_nccl_mem_pool() is None
+            or version.parse(torch.__version__) < version.parse("2.8.0.a0")
+        )
+        if self.disabled:
+            self.pynccl_comm: PyNcclCommunicator | None = None
+            self._mem_pool_ctx: contextlib.AbstractContextManager[Any] = (
+                contextlib.nullcontext()
+            )
+            self.is_graph_capture = None
+            self.device = None
+        else:
+            self.pynccl_comm = pynccl_comm
+            self._mem_pool_ctx = torch.cuda.use_mem_pool(get_nccl_mem_pool())
+            self.is_graph_capture = torch.cuda.is_current_stream_capturing()
+            self.device = torch.accelerator.current_device_index()
+
+    def __enter__(self):
+        if self.disabled:
+            return self
+        assert self.pynccl_comm is not None, (
+            "Symmetric memory requires pynccl to be initialized"
+        )
+        assert self.pynccl_comm.nccl_version >= 22703, (
+            "NCCL version 2.27.3 or higher is required for NCCL symmetric memory"
+        )
+        if self.is_graph_capture:
+            assert _graph_pool_id is not None, (
+                "graph_pool_id is not set under graph capture"
+            )
+            # Pause graph memory pool to use symmetric memory with cuda graph
+            torch._C._cuda_endAllocateToPool(self.device, _graph_pool_id)
+        self._mem_pool_ctx.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.disabled:
+            return
+        global _cached_pool_snapshot
+        global _registered_base_addrs
+        self._mem_pool_ctx.__exit__(exc_type, exc_val, exc_tb)
+        _pool = get_nccl_mem_pool()
+        assert _pool is not None
+        _cached_pool_snapshot = _pool.snapshot()
+        assert self.pynccl_comm is not None
+        comm_key = bytes(self.pynccl_comm.unique_id.internal)
+        if comm_key not in _registered_base_addrs:
+            _registered_base_addrs[comm_key] = set()
+        for segment in _cached_pool_snapshot:
+            if segment["address"] not in _registered_base_addrs[comm_key]:
+                self.pynccl_comm.register_comm_window_raw(
+                    segment["address"], segment["total_size"]
+                )
+                _registered_base_addrs[comm_key].add(segment["address"])
+        if self.is_graph_capture:
+            torch._C._cuda_beginAllocateCurrentThreadToPool(self.device, _graph_pool_id)

vllm/distributed/device_communicators/pynccl_wrapper.py

@@ -0,0 +1,589 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# This file is a pure Python wrapper for the NCCL library.
+# The main purpose is to use NCCL combined with CUDA graph.
+# Before writing this script, we tried the following approach:
+# 1. We tried to use `cupy`, it calls NCCL correctly, but `cupy` itself
+#  often gets stuck when initializing the NCCL communicator.
+# 2. We tried to use `torch.distributed`, but `torch.distributed.all_reduce`
+#  contains many other potential cuda APIs, that are not allowed during
+#  capturing the CUDA graph. For further details, please check
+# https://discuss.pytorch.org/t/pytorch-cudagraph-with-nccl-operation-failed/ .
+#
+# Another rejected idea is to write a C/C++ binding for NCCL. It is usually
+# doable, but we often encounter issues related with nccl versions, and need
+# to switch between different versions of NCCL. See
+# https://github.com/NVIDIA/nccl/issues/1234 for more details.
+# A C/C++ binding is not flexible enough to handle this. It requires
+# recompilation of the code every time we want to switch between different
+# versions. This current implementation, with a **pure** Python wrapper, is
+# more flexible. We can easily switch between different versions of NCCL by
+# changing the environment variable `VLLM_NCCL_SO_PATH`, or the `so_file`
+# variable in the code.
+
+import ctypes
+import functools
+import platform
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from torch.distributed import ReduceOp
+
+from vllm import envs
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+from vllm.utils.nccl import find_nccl_library
+
+logger = init_logger(__name__)
+
+# === export types and functions from nccl to Python ===
+# for the original nccl definition, please check
+# https://github.com/NVIDIA/nccl/blob/master/src/nccl.h.in
+
+ncclResult_t = ctypes.c_int
+ncclComm_t = ctypes.c_void_p
+ncclWindow_t = ctypes.c_void_p
+
+
+class ncclUniqueId(ctypes.Structure):
+    _fields_ = [("internal", ctypes.c_byte * 128)]
+
+
+cudaStream_t = ctypes.c_void_p
+buffer_type = ctypes.c_void_p
+
+ncclDataType_t = ctypes.c_int
+
+
+class ncclDataTypeEnum:
+    ncclInt8 = 0
+    ncclChar = 0
+    ncclUint8 = 1
+    ncclInt32 = 2
+    ncclInt = 2
+    ncclUint32 = 3
+    ncclInt64 = 4
+    ncclUint64 = 5
+    ncclFloat16 = 6
+    ncclHalf = 6
+    ncclFloat32 = 7
+    ncclFloat = 7
+    ncclFloat64 = 8
+    ncclDouble = 8
+    ncclBfloat16 = 9
+    ncclFloat8e4m3 = 10
+    ncclNumTypes = 11
+
+    @classmethod
+    @functools.lru_cache(maxsize=1)
+    def _torch_to_nccl_map(cls) -> dict[torch.dtype, int]:
+        return {
+            torch.int8: cls.ncclInt8,
+            torch.uint8: cls.ncclUint8,
+            torch.int32: cls.ncclInt32,
+            torch.int64: cls.ncclInt64,
+            torch.float16: cls.ncclFloat16,
+            torch.float32: cls.ncclFloat32,
+            torch.float64: cls.ncclFloat64,
+            torch.bfloat16: cls.ncclBfloat16,
+            current_platform.fp8_dtype(): cls.ncclFloat8e4m3,
+        }
+
+    @classmethod
+    def supports_torch_dtype(cls, dtype: torch.dtype) -> bool:
+        return dtype in cls._torch_to_nccl_map()
+
+    @classmethod
+    def try_from_torch(cls, dtype: torch.dtype) -> int | None:
+        return cls._torch_to_nccl_map().get(dtype)
+
+    @classmethod
+    def from_torch(cls, dtype: torch.dtype) -> int:
+        nccl_dtype = cls.try_from_torch(dtype)
+        if nccl_dtype is not None:
+            return nccl_dtype
+        raise ValueError(
+            f"Unsupported dtype {dtype}: should be one of "
+            f"int8, uint8, int32, int64, float16, float32, float64, bfloat16,"
+            " float8e4m3."
+        )
+
+
+ncclRedOp_t = ctypes.c_int
+
+
+class ncclRedOpTypeEnum:
+    ncclSum = 0
+    ncclProd = 1
+    ncclMax = 2
+    ncclMin = 3
+    ncclAvg = 4
+    ncclNumOps = 5
+
+    @classmethod
+    def from_torch(cls, op: ReduceOp) -> int:
+        if op == ReduceOp.SUM:
+            return cls.ncclSum
+        if op == ReduceOp.PRODUCT:
+            return cls.ncclProd
+        if op == ReduceOp.MAX:
+            return cls.ncclMax
+        if op == ReduceOp.MIN:
+            return cls.ncclMin
+        if op == ReduceOp.AVG:
+            return cls.ncclAvg
+        raise ValueError(f"Unsupported op: {op}")
+
+
+@dataclass
+class Function:
+    name: str
+    restype: Any
+    argtypes: list[Any]
+
+
+class NCCLLibrary:
+    exported_functions = [
+        # const char* ncclGetErrorString(ncclResult_t result)
+        Function("ncclGetErrorString", ctypes.c_char_p, [ncclResult_t]),
+        # ncclResult_t  ncclGetVersion(int *version);
+        Function("ncclGetVersion", ncclResult_t, [ctypes.POINTER(ctypes.c_int)]),
+        # ncclResult_t ncclGetUniqueId(ncclUniqueId* uniqueId);
+        Function("ncclGetUniqueId", ncclResult_t, [ctypes.POINTER(ncclUniqueId)]),
+        # ncclResult_t  ncclCommInitRank(
+        #   ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank);
+        # note that ncclComm_t is a pointer type, so the first argument
+        # is a pointer to a pointer
+        Function(
+            "ncclCommInitRank",
+            ncclResult_t,
+            [ctypes.POINTER(ncclComm_t), ctypes.c_int, ncclUniqueId, ctypes.c_int],
+        ),
+        # ncclResult_t  ncclAllReduce(
+        #   const void* sendbuff, void* recvbuff, size_t count,
+        #   ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm,
+        #   cudaStream_t stream);
+        # note that cudaStream_t is a pointer type, so the last argument
+        # is a pointer
+        Function(
+            "ncclAllReduce",
+            ncclResult_t,
+            [
+                buffer_type,
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ncclRedOp_t,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # ncclResult_t  ncclReduce(
+        #   const void* sendbuff, void* recvbuff, size_t count,
+        #   ncclDataType_t datatype, ncclRedOp_t op, int root,
+        #   ncclComm_t comm,  cudaStream_t stream);
+        # note that cudaStream_t is a pointer type, so the last argument
+        # is a pointer
+        Function(
+            "ncclReduce",
+            ncclResult_t,
+            [
+                buffer_type,
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ncclRedOp_t,
+                ctypes.c_int,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # ncclResult_t  ncclAllGather(
+        #   const void* sendbuff, void* recvbuff, size_t count,
+        #   ncclDataType_t datatype, ncclComm_t comm,
+        #   cudaStream_t stream);
+        # note that cudaStream_t is a pointer type, so the last argument
+        # is a pointer
+        Function(
+            "ncclAllGather",
+            ncclResult_t,
+            [
+                buffer_type,
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # ncclResult_t  ncclReduceScatter(
+        #   const void* sendbuff, void* recvbuff, size_t count,
+        #   ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm,
+        #   cudaStream_t stream);
+        # note that cudaStream_t is a pointer type, so the last argument
+        # is a pointer
+        Function(
+            "ncclReduceScatter",
+            ncclResult_t,
+            [
+                buffer_type,
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ncclRedOp_t,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # ncclResult_t  ncclSend(
+        #   const void* sendbuff, size_t count, ncclDataType_t datatype,
+        #   int dest, ncclComm_t comm, cudaStream_t stream);
+        Function(
+            "ncclSend",
+            ncclResult_t,
+            [
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ctypes.c_int,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # ncclResult_t  ncclRecv(
+        #   void* recvbuff, size_t count, ncclDataType_t datatype,
+        #   int src, ncclComm_t comm, cudaStream_t stream);
+        Function(
+            "ncclRecv",
+            ncclResult_t,
+            [
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ctypes.c_int,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # ncclResult_t ncclBroadcast(
+        #   const void* sendbuff, void* recvbuff, size_t count,
+        #   ncclDataType_t datatype, int root, ncclComm_t comm,
+        #   cudaStream_t stream);
+        Function(
+            "ncclBroadcast",
+            ncclResult_t,
+            [
+                buffer_type,
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ctypes.c_int,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # be cautious! this is a collective call, it will block until all
+        # processes in the communicator have called this function.
+        # because Python object destruction can happen in random order,
+        # it is better not to call it at all.
+        # ncclResult_t  ncclCommDestroy(ncclComm_t comm);
+        Function("ncclCommDestroy", ncclResult_t, [ncclComm_t]),
+        # ncclCommAbort frees resources associated with the communicator
+        # without requiring a collective synchronization. Unlike
+        # ncclCommDestroy, it is safe to call during an uncoordinated
+        # shutdown when peer ranks may already be gone.
+        # ncclResult_t  ncclCommAbort(ncclComm_t comm);
+        Function("ncclCommAbort", ncclResult_t, [ncclComm_t]),
+        # ncclResult_t ncclGroupStart();
+        Function("ncclGroupStart", ncclResult_t, []),
+        # ncclResult_t ncclGroupEnd();
+        Function("ncclGroupEnd", ncclResult_t, []),
+        # ncclResult_t ncclCommWindowRegister(
+        #   ncclComm_t comm, void* buff, size_t size,
+        #   ncclWindow_t* win, int winFlags);
+        Function(
+            "ncclCommWindowRegister",
+            ncclResult_t,
+            [
+                ncclComm_t,
+                buffer_type,
+                ctypes.c_size_t,
+                ctypes.POINTER(ncclWindow_t),
+                ctypes.c_int,
+            ],
+        ),
+        # ncclResult_t ncclCommWindowDeregister(
+        #   ncclComm_t comm, ncclWindow_t win);
+        Function("ncclCommWindowDeregister", ncclResult_t, [ncclComm_t, ncclWindow_t]),
+    ]
+
+    # class attribute to store the mapping from the path to the library
+    # to avoid loading the same library multiple times
+    path_to_library_cache: dict[str, Any] = {}
+
+    # class attribute to store the mapping from library path
+    #  to the corresponding dictionary
+    path_to_dict_mapping: dict[str, dict[str, Any]] = {}
+
+    def __init__(self, so_file: str | None = None):
+        so_file = so_file or find_nccl_library()
+
+        try:
+            if so_file not in NCCLLibrary.path_to_dict_mapping:
+                lib = ctypes.CDLL(so_file)
+                NCCLLibrary.path_to_library_cache[so_file] = lib
+            self.lib = NCCLLibrary.path_to_library_cache[so_file]
+        except Exception as e:
+            logger.error(
+                "Failed to load NCCL library from %s. "
+                "It is expected if you are not running on NVIDIA/AMD GPUs."
+                "Otherwise, the nccl library might not exist, be corrupted "
+                "or it does not support the current platform %s. "
+                "If you already have the library, please set the "
+                "environment variable VLLM_NCCL_SO_PATH"
+                " to point to the correct nccl library path.",
+                so_file,
+                platform.platform(),
+            )
+            raise e
+
+        if so_file not in NCCLLibrary.path_to_dict_mapping:
+            _funcs: dict[str, Any] = {}
+            for func in NCCLLibrary.exported_functions:
+                try:
+                    f = getattr(self.lib, func.name)
+                    f.restype = func.restype
+                    f.argtypes = func.argtypes
+                    _funcs[func.name] = f
+                except AttributeError:
+                    if func.name in [
+                        "ncclCommWindowRegister",
+                        "ncclCommWindowDeregister",
+                    ]:
+                        if envs.VLLM_USE_NCCL_SYMM_MEM:
+                            logger.warning_once(
+                                "The symbol %s is not found in the NCCL "
+                                "library %s. To enable VLLM_USE_NCCL_SYMM_MEM "
+                                " please update your NCCL version to >= "
+                                "2.27.03.",
+                                func.name,
+                                so_file,
+                            )
+                        if current_platform.is_rocm():
+                            # Having an exception here on ROCm platform is
+                            # not allowed during graph capturing
+                            continue
+                    raise
+            NCCLLibrary.path_to_dict_mapping[so_file] = _funcs
+        self._funcs = NCCLLibrary.path_to_dict_mapping[so_file]
+
+    def ncclGetErrorString(self, result: ncclResult_t) -> str:
+        return self._funcs["ncclGetErrorString"](result).decode("utf-8")
+
+    def NCCL_CHECK(self, result: ncclResult_t) -> None:
+        if result != 0:
+            error_str = self.ncclGetErrorString(result)
+            raise RuntimeError(f"NCCL error: {error_str}")
+
+    def ncclGetRawVersion(self) -> int:
+        version = ctypes.c_int()
+        self.NCCL_CHECK(self._funcs["ncclGetVersion"](ctypes.byref(version)))
+        # something like 21903
+        return version.value
+
+    def ncclGetVersion(self) -> str:
+        version_str = str(self.ncclGetRawVersion())
+        # something like 21903 --> "2.19.3"
+        major = version_str[0].lstrip("0")
+        minor = version_str[1:3].lstrip("0")
+        patch = version_str[3:].lstrip("0")
+        return f"{major}.{minor}.{patch}"
+
+    def ncclGetUniqueId(self) -> ncclUniqueId:
+        unique_id = ncclUniqueId()
+        self.NCCL_CHECK(self._funcs["ncclGetUniqueId"](ctypes.byref(unique_id)))
+        return unique_id
+
+    def unique_id_from_bytes(self, data: bytes) -> ncclUniqueId:
+        if len(data) != 128:
+            raise ValueError(
+                f"Expected 128 bytes for ncclUniqueId, got {len(data)} bytes"
+            )
+        unique_id = ncclUniqueId()
+        ctypes.memmove(ctypes.addressof(unique_id.internal), data, 128)
+        return unique_id
+
+    def ncclCommInitRank(
+        self, world_size: int, unique_id: ncclUniqueId, rank: int
+    ) -> ncclComm_t:
+        comm = ncclComm_t()
+        self.NCCL_CHECK(
+            self._funcs["ncclCommInitRank"](
+                ctypes.byref(comm), world_size, unique_id, rank
+            )
+        )
+        return comm
+
+    def ncclAllReduce(
+        self,
+        sendbuff: buffer_type,
+        recvbuff: buffer_type,
+        count: int,
+        datatype: int,
+        op: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        # `datatype` actually should be `ncclDataType_t`
+        # and `op` should be `ncclRedOp_t`
+        # both are aliases of `ctypes.c_int`
+        # when we pass int to a function, it will be converted to `ctypes.c_int`
+        # by ctypes automatically
+        self.NCCL_CHECK(
+            self._funcs["ncclAllReduce"](
+                sendbuff, recvbuff, count, datatype, op, comm, stream
+            )
+        )
+
+    def ncclReduce(
+        self,
+        sendbuff: buffer_type,
+        recvbuff: buffer_type,
+        count: int,
+        datatype: int,
+        op: int,
+        root: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        # `datatype` actually should be `ncclDataType_t`
+        # and `op` should be `ncclRedOp_t`
+        # both are aliases of `ctypes.c_int`
+        # when we pass int to a function, it will be converted to `ctypes.c_int`
+        # by ctypes automatically
+        self.NCCL_CHECK(
+            self._funcs["ncclReduce"](
+                sendbuff, recvbuff, count, datatype, op, root, comm, stream
+            )
+        )
+
+    def ncclReduceScatter(
+        self,
+        sendbuff: buffer_type,
+        recvbuff: buffer_type,
+        count: int,
+        datatype: int,
+        op: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        # `datatype` actually should be `ncclDataType_t`
+        # and `op` should be `ncclRedOp_t`
+        # both are aliases of `ctypes.c_int`
+        # when we pass int to a function, it will be converted to `ctypes.c_int`
+        # by ctypes automatically
+        self.NCCL_CHECK(
+            self._funcs["ncclReduceScatter"](
+                sendbuff, recvbuff, count, datatype, op, comm, stream
+            )
+        )
+
+    def ncclAllGather(
+        self,
+        sendbuff: buffer_type,
+        recvbuff: buffer_type,
+        count: int,
+        datatype: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        # `datatype` actually should be `ncclDataType_t`
+        # which is an aliases of `ctypes.c_int`
+        # when we pass int to a function, it will be converted to `ctypes.c_int`
+        # by ctypes automatically
+        self.NCCL_CHECK(
+            self._funcs["ncclAllGather"](
+                sendbuff, recvbuff, count, datatype, comm, stream
+            )
+        )
+
+    def ncclSend(
+        self,
+        sendbuff: buffer_type,
+        count: int,
+        datatype: int,
+        dest: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        self.NCCL_CHECK(
+            self._funcs["ncclSend"](sendbuff, count, datatype, dest, comm, stream)
+        )
+
+    def ncclRecv(
+        self,
+        recvbuff: buffer_type,
+        count: int,
+        datatype: int,
+        src: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        self.NCCL_CHECK(
+            self._funcs["ncclRecv"](recvbuff, count, datatype, src, comm, stream)
+        )
+
+    def ncclBroadcast(
+        self,
+        sendbuff: buffer_type,
+        recvbuff: buffer_type,
+        count: int,
+        datatype: int,
+        root: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        self.NCCL_CHECK(
+            self._funcs["ncclBroadcast"](
+                sendbuff, recvbuff, count, datatype, root, comm, stream
+            )
+        )
+
+    def ncclCommDestroy(self, comm: ncclComm_t) -> None:
+        self.NCCL_CHECK(self._funcs["ncclCommDestroy"](comm))
+
+    def ncclCommAbort(self, comm: ncclComm_t) -> None:
+        self.NCCL_CHECK(self._funcs["ncclCommAbort"](comm))
+
+    def ncclGroupStart(self) -> None:
+        self.NCCL_CHECK(self._funcs["ncclGroupStart"]())
+
+    def ncclGroupEnd(self) -> None:
+        self.NCCL_CHECK(self._funcs["ncclGroupEnd"]())
+
+    def ncclCommWindowRegister(
+        self, comm: ncclComm_t, buff: buffer_type, size: int, win_flags: int
+    ) -> ncclWindow_t:
+        window = ncclWindow_t()
+        self.NCCL_CHECK(
+            self._funcs["ncclCommWindowRegister"](
+                comm, buff, size, ctypes.byref(window), win_flags
+            )
+        )
+        return window
+
+    def ncclCommWindowDeregister(self, comm: ncclComm_t, window: ncclWindow_t) -> None:
+        self.NCCL_CHECK(self._funcs["ncclCommWindowDeregister"](comm, window))
+
+
+__all__ = [
+    "NCCLLibrary",
+    "ncclDataTypeEnum",
+    "ncclRedOpTypeEnum",
+    "ncclUniqueId",
+    "ncclComm_t",
+    "cudaStream_t",
+    "buffer_type",
+]

vllm/distributed/device_communicators/quick_all_reduce.py

@@ -0,0 +1,367 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from enum import Enum
+
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+
+import vllm.envs as envs
+from vllm import _custom_ops as ops
+from vllm.config import get_current_vllm_config_or_none
+from vllm.distributed.parallel_state import in_the_same_node_as
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+logger = init_logger(__name__)
+
+try:
+    ops.qr_max_size()
+    quick_ar = True
+except Exception:
+    # For CPUs and CUDA
+    quick_ar = False
+
+
+from vllm.distributed.utils import is_weak_contiguous  # noqa: E402, F401
+
+
+class QuickReduceRegime(Enum):
+    # Keep integer ids aligned with csrc/quickreduce/quick_reduce.h
+    FP = 0
+    INT8 = 1
+    INT6 = 2
+    INT4 = 3
+    INT3 = 4
+    NONE = 5
+
+
+KB = 1024
+MB = 1024 * KB
+
+
+class QuickAllReduce:
+    _SUPPORTED_WORLD_SIZES = [2, 4, 8]
+    _SUPPORTED_DTYPES = [torch.float16, torch.bfloat16]
+    # The following data is based on kernel tests.
+    # In this order [FP, INT8, INT6, INT4, INT3].
+    _QR_MIN_SIZE = {
+        (torch.float16, 2): [1 * MB, 2 * MB, 2 * MB, 1 * MB, 1 * MB],
+        (torch.float16, 4): [1 * MB, 16 * MB, 4 * MB, 2 * MB, 2 * MB],
+        (torch.float16, 8): [16 * MB, 4 * MB, 4 * MB, 2 * MB, 2 * MB],
+        (torch.bfloat16, 2): [2 * MB, 8 * MB, 8 * MB, 8 * MB, 8 * MB],
+        (torch.bfloat16, 4): [8 * MB, 64 * MB, 64 * MB, 16 * MB, 16 * MB],
+        (torch.bfloat16, 8): [
+            16 * MB,
+            2048 * MB,
+            2048 * MB,
+            2048 * MB,
+            2048 * MB,
+        ],
+    }
+
+    def __init__(self, group: ProcessGroup, device: int | str | torch.device) -> None:
+        """
+        Custom allreduce provides non-destructive acceleration and is
+        available for CUDA and ROCm MI300 series.
+
+        Custom quick allreduce leverages quantization for further
+        acceleration on ROCm. It currently supports Q8, Q6, Q4, and Q3
+        quantization formats and FP(float16, bfloat16). Q3 (INT3) is
+        restricted to TP2 (world_size == 2) due to poor performance on
+        larger world sizes.
+
+        Quick allreduce is designed as a complement to custom allreduce.
+        Its initialization requires even stricter conditions.
+
+        Only the ROCm MI300 series is supported for quick allreduce at
+        this time.
+
+        Args:
+            group: the process group to work on. If None, it will use the
+                default process group.
+            device: the device to bind the CustomAllreduce to. If None,
+                it will be bound to f"cuda:{local_rank}".
+        It is the caller's responsibility to make sure each communicator
+        is bind to a unique device, and all communicators in this group
+        are in the same node.
+        """
+        self.disabled = True
+        if not self._rocm_arch_available():
+            logger.debug(
+                "Custom quick allreduce is only supported on ROCm MI300 series."
+            )
+            return
+
+        if not quick_ar:
+            # disable because of missing quick reduce library
+            # e.g. in a cuda environment
+            logger.info(
+                "Custom quick allreduce is disabled because "
+                "of missing custom quick allreduce library"
+            )
+            return
+
+        self.group = group
+        assert dist.get_backend(group) != dist.Backend.NCCL, (
+            "Custom quick allreduce should be attached to a non-NCCL group."
+        )
+        if not all(in_the_same_node_as(group, source_rank=0)):
+            # No need to initialize custom quick allreduce for
+            # multi-node case.
+            logger.warning(
+                "Custom quick allreduce is disabled because this "
+                "process group spans across nodes."
+            )
+            return
+        rank = dist.get_rank(group=self.group)
+        world_size = dist.get_world_size(group=self.group)
+        self.rank = rank
+        self.world_size = world_size
+        if world_size == 1:
+            # No need to initialize QuickReduce for single GPU case.
+            return
+
+        if world_size not in QuickAllReduce._SUPPORTED_WORLD_SIZES:
+            logger.warning(
+                "Custom quick allreduce is disabled due to an "
+                "unsupported world size: %d. Supported world sizes: %s.",
+                world_size,
+                str(QuickAllReduce._SUPPORTED_WORLD_SIZES),
+            )
+            return
+
+        if isinstance(device, int):
+            device = torch.device(f"cuda:{device}")
+        elif isinstance(device, str):
+            device = torch.device(device)
+        assert isinstance(device, torch.device)
+        self.device = device
+
+        # device.index is a visible ordinal, not a logical local ID.
+        physical_device_id = current_platform.visible_device_id_to_physical_device_id(
+            device.index
+        )
+        tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")
+        gather_list = [
+            torch.tensor([0], dtype=torch.int, device="cpu")
+            for _ in range(self.world_size)
+        ]
+        dist.all_gather(gather_list, tensor, group=self.group)
+        physical_device_ids = [t.item() for t in gather_list]
+
+        # test nvlink first, this will filter out most of the cases
+        # where custom quick allreduce is not supported
+        # this checks hardware and driver support for NVLink
+        assert current_platform.is_cuda_alike()
+        self.fully_connected = current_platform.is_fully_connected(physical_device_ids)
+        if self.world_size > 2 and not self.fully_connected:
+            logger.debug(
+                "Custom quick allreduce is disabled because it's not supported "
+                "on more than two PCIe-only GPUs. "
+            )
+            return
+
+        self.init_quick_all_reduce()
+
+    def init_quick_all_reduce(self):
+        # On RocM, bfloat16 kernels are slower than fp16
+        # due to slower match operations
+        # If environment variable is set to 1, we convert input to fp16
+        self.use_fp16_kernels = envs.VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16
+        regime_str = envs.VLLM_ROCM_QUICK_REDUCE_QUANTIZATION
+        if regime_str not in QuickReduceRegime.__members__:
+            logger.warning(
+                "Custom quick allreduce:",
+                f"Invalid quantization level: {regime_str}. "
+                "Supported levels: "
+                f"{list(QuickReduceRegime.__members__.keys())}",
+            )
+            return
+
+        if regime_str == "NONE":
+            logger.debug(
+                "Custom quick allreduce is disabled based "
+                "on env variable "
+                "VLLM_ROCM_QUICK_REDUCE_QUANTIZATION='NONE'"
+            )
+            return
+        self.qr_quant_level = QuickReduceRegime[regime_str]
+
+        # INT3 is only enabled for TP2 (world_size == 2).
+        # Kernel benchmarks show INT3 all-reduce on TP4/TP8 has poor
+        # performance (the extra ranks make the 3-bit codec's pack/unpack
+        # overhead outweigh the reduced communication volume), so INT3 is
+        # restricted to 2-GPU tensor parallelism. For TP4/TP8 use a wider
+        # codec (e.g. INT4) or NONE instead.
+        if self.qr_quant_level == QuickReduceRegime.INT3 and self.world_size != 2:
+            logger.warning(
+                "Custom quick allreduce is disabled: INT3 quantization is "
+                "only supported for TP2 (world_size == 2), but world_size "
+                "is %d. INT3 on TP4/TP8 is disabled due to poor kernel "
+                "performance. Use INT4/NONE for this world size.",
+                self.world_size,
+            )
+            return
+
+        self.qr_quantization_min_size = self._get_qr_quantization_min_size()
+        vllm_config = get_current_vllm_config_or_none()
+        if (
+            vllm_config is not None
+            and hasattr(vllm_config, "model_config")
+            and hasattr(vllm_config.model_config, "dtype")
+        ):
+            dtype = vllm_config.model_config.dtype
+            if dtype not in [torch.float16, torch.bfloat16]:
+                logger.debug(
+                    "Custom quick allreduce disabled: only supports "
+                    "float16 and float16, but get %s.",
+                    dtype,
+                )
+                return
+
+            if dtype == torch.bfloat16 and self.use_fp16_kernels:
+                logger.info(
+                    "Custom quick allreduce: BF16 inputs will be converted "
+                    "to FP16 to improve performance. set "
+                    "envs.VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16=0 "
+                    "to turn off."
+                )
+
+        # VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB is specified in MB
+        qr_max_size = envs.VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB
+        if qr_max_size is not None:
+            if qr_max_size < 1:
+                logger.info(
+                    "You should not set a max_size smaller than 1MB, which can "
+                    "lead to error or degradation to custom allreduce or rccl."
+                )
+            qr_max_size = qr_max_size * MB
+        effective_qr_max_size = (
+            qr_max_size if qr_max_size is not None else ops.qr_max_size()
+        )
+        qr_min_size = self._get_qr_min_size(effective_qr_max_size)
+        self._ptr = ops.init_custom_qr(self.rank, self.world_size, qr_max_size)
+        self.qr_max_size = effective_qr_max_size
+        self.qr_min_size = qr_min_size
+        if qr_min_size is not None:
+            logger.info(
+                "Custom quick allreduce: min size override = %d MB",
+                qr_min_size // MB,
+            )
+        if self.qr_quantization_min_size is not None:
+            logger.info(
+                "Custom quick allreduce: quantization codec threshold = %d KB",
+                self.qr_quantization_min_size // KB,
+            )
+        self.create_shared_buffer()
+        self.disabled = False
+
+    @staticmethod
+    def _get_qr_min_size(qr_max_size: int | None) -> int | None:
+        qr_min_size = envs.VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB
+        if qr_min_size is None:
+            return None
+        if qr_min_size < 0:
+            raise ValueError(
+                "VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB must be non-negative, "
+                f"got {qr_min_size}"
+            )
+        qr_min_size *= MB
+        if qr_max_size is not None and qr_min_size > qr_max_size:
+            raise ValueError(
+                "VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB must be less than or "
+                "equal to the effective QuickReduce max size"
+            )
+        return qr_min_size
+
+    @staticmethod
+    def _get_qr_quantization_min_size() -> int | None:
+        quantization_min_size = envs.VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB
+        if quantization_min_size is None:
+            return None
+        if quantization_min_size < 0:
+            raise ValueError(
+                "VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB must be "
+                f"non-negative, got {quantization_min_size}"
+            )
+        return quantization_min_size * KB
+
+    def _rocm_arch_available(self):
+        if not current_platform.is_rocm():
+            return False
+        try:
+            props = torch.cuda.get_device_properties(0)
+            gcn_arch = getattr(props, "gcnArchName", "")
+            supported_archs = ["gfx94", "gfx95"]
+            return any(gfx in gcn_arch for gfx in supported_archs)
+        except Exception as e:
+            logger.warning("Failed to determine ROCm for quick allreduce: %s", e)
+            return False
+
+    def create_shared_buffer(self):
+        """
+        Creates a shared buffer for quickreduce.
+        Has to be called after init_custom_qr
+        """
+        handle = ops.qr_get_handle(self._ptr)
+        world_size = dist.get_world_size(group=self.group)
+        handles = [None] * world_size
+        dist.all_gather_object(handles, handle, group=self.group)
+        ops.qr_open_handles(self._ptr, handles)
+
+    def should_quick_allreduce(self, inp: torch.Tensor):
+        """
+        Check if quickreduce is available
+        """
+        if self.disabled:
+            return False
+        if inp.dtype not in self._SUPPORTED_DTYPES:
+            return False
+        inp_size = inp.numel() * inp.element_size()
+        # custom quick allreduce requires input byte size to be
+        # multiples of 16
+        if inp_size % 16 != 0:
+            return False
+        if not is_weak_contiguous(inp):
+            return False
+        dtype = inp.dtype
+        if self.use_fp16_kernels:
+            dtype = torch.float16
+        min_size = self.qr_min_size
+        if min_size is None:
+            min_size = self._QR_MIN_SIZE[(dtype, self.world_size)][
+                self.qr_quant_level.value
+            ]
+        return inp_size <= self.qr_max_size and inp_size >= min_size
+
+    def quick_all_reduce(self, inp: torch.Tensor, *, out: torch.Tensor = None):
+        """Performs an out-of-place custom quick all reduce."""
+        # quick allreduce doesn't require a separate graph mode,
+        # as QR uses static IPC buffer.
+        if out is None:
+            out = torch.empty_like(inp)
+        ops.qr_all_reduce(
+            self._ptr, inp, out, self._get_qr_quant_level(inp), self.use_fp16_kernels
+        )
+        return out
+
+    def _get_qr_quant_level(self, inp: torch.Tensor) -> int:
+        quantization_min_size = self.qr_quantization_min_size
+        if (
+            quantization_min_size is not None
+            and inp.numel() * inp.element_size() < quantization_min_size
+        ):
+            return QuickReduceRegime.FP.value
+        return self.qr_quant_level.value
+
+    def close(self):
+        if not self.disabled and getattr(self, "_ptr", None):
+            if ops is not None:
+                ops.qr_destroy(self._ptr)
+            self._ptr = 0
+            self.disabled = True
+
+    def __del__(self):
+        self.close()

vllm/distributed/device_communicators/ray_communicator.py

@@ -0,0 +1,259 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import uuid
+from typing import Any
+
+import ray
+import torch
+from ray.exceptions import RayChannelError
+from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator
+from torch.distributed import ReduceOp
+
+from vllm.distributed.device_communicators.base_device_communicator import (
+    DeviceCommunicatorBase,
+)
+from vllm.distributed.parallel_state import get_pp_group
+from vllm.logger import init_logger
+from vllm.utils.torch_utils import current_stream
+
+logger = init_logger(__name__)
+
+
+class RayPPCommunicator(Communicator):
+    """
+    Communicator to be used for pipeline parallelism in Ray Compiled Graph.
+    This is wraps around the vLLM _PP GroupCoordinator.
+
+    This class is not thread-safe.
+    """
+
+    _comm: DeviceCommunicatorBase | None
+
+    def __init__(
+        self,
+        world_size: int,
+        comm_id: Any,
+        rank: int | None,
+        actor_handles: list["ray.actor.ActorHandle"],
+        cuda_stream: torch.cuda.Stream | None,
+        use_communication_streams: bool = False,
+    ):
+        """
+        Initialize a RayPPCommunicator that can be used to communicate with
+        other Ray Compiled Graph actors for pipeline parallelism.
+
+        Args:
+            world_size: The number of participating actors.
+            comm_id: A unique communicator ID. This is just to conform with
+                the Ray Communicator API and is not used.
+            rank: The rank of this actor. If None, then the caller is not a
+                participant of the RayPPCommunicator group (e.g., the Ray
+                driver).
+            actor_handles: A list of actor handles.
+            cuda_stream: A CUDA stream to dispatch communication ops to. This
+                is not supported.
+            use_communication_streams: Whether to use communication streams.
+                This is not supported.
+        """
+        self._world_size = world_size
+        self._rank: int | None = None
+        self._actor_handles = actor_handles
+        if use_communication_streams:
+            raise NotImplementedError("use_communication_streams is not supported")
+        if cuda_stream is not None and cuda_stream != current_stream():
+            raise ValueError(
+                "cuda_stream other than the current stream is not supported"
+            )
+
+        if rank is not None:
+            # Rank is not None, this is Ray worker
+            assert ray.get_gpu_ids(), "RayPPCommunicator has no GPUs assigned"
+
+            self._comm = get_pp_group().device_communicator
+            assert self._comm is not None
+
+            # Since we wrap around the vLLM _PP communicator, we use
+            # the rank from the vLLM communicator, and ignore the rank
+            # passed in from Ray.
+            # TODO(rui): refactor the Ray Communicator API so that
+            # it also supports no rank passed in.
+            self._rank = self._comm.rank_in_group
+
+            self._build_actor_rank_mapping()
+        else:
+            # Rank is None, this is Ray driver
+            self._comm = None
+
+        self._closed = False
+
+    def _build_actor_rank_mapping(self):
+        """
+        Use collective communication to build a mapping from actor IDs to ranks.
+        This should be called once during initialization.
+        """
+        if self._comm is None:
+            return {}
+
+        current_actor = ray.get_runtime_context().current_actor
+        actor_id_str = current_actor._actor_id.hex()
+
+        # Ray actor IDs are 32-character hex strings (128 bits)
+        ACTOR_ID_LEN = 32
+        actor_id_bytes = bytearray(actor_id_str.encode("utf-8"))
+        assert len(actor_id_bytes) == ACTOR_ID_LEN, (
+            f"Unexpected actor ID length: {len(actor_id_bytes)}"
+        )
+
+        actor_id_tensor = torch.frombuffer(actor_id_bytes, dtype=torch.uint8).to(
+            self._comm.device
+        )
+
+        # All-gather full actor IDs from all actors
+        gathered_ids = self._comm.all_gather(actor_id_tensor, dim=0)
+
+        # Build mapping: actor_id -> device_comm_rank
+        self._actor_id_to_rank = {}
+        for rank in range(self._world_size):
+            start_idx = rank * ACTOR_ID_LEN
+            end_idx = (rank + 1) * ACTOR_ID_LEN
+            actor_bytes = gathered_ids[start_idx:end_idx].cpu().numpy().tobytes()
+            actor_id = actor_bytes.decode("utf-8")
+            self._actor_id_to_rank[actor_id] = rank
+
+    def initialize(self, rank: int) -> None:
+        # No additional initialization is needed.
+        pass
+
+    def get_actor_handles(self) -> list["ray.actor.ActorHandle"]:
+        return self._actor_handles
+
+    def get_rank(self, actor: ray.actor.ActorHandle) -> int:
+        """
+        Return the given actor's rank using device communicator collective ops.
+        """
+        assert hasattr(self, "_actor_id_to_rank"), (
+            "Actor rank mapping not built. "
+            "This should have been done during initialization."
+        )
+
+        actor_id_str = actor._actor_id.hex()
+
+        if actor_id_str in self._actor_id_to_rank:
+            return self._actor_id_to_rank[actor_id_str]  # type: ignore
+        else:
+            raise ValueError(f"Actor {actor} not found in communicator group")
+
+    def get_self_rank(self) -> int | None:
+        """
+        Return this actor's rank.
+        """
+        return self._rank
+
+    def get_world_size(self) -> int:
+        """
+        Return the number of ranks in the RayPPCommunicator group.
+        """
+        return self._world_size
+
+    def send(self, buf: "torch.Tensor", peer_rank: int) -> None:
+        """
+        Send a torch.Tensor to a peer.
+
+        This returns when the send kernel has been queued, but the kernel may
+        not have completed. Therefore, the caller should ensure that there are
+        no concurrent writes to the sent `buf` until the send has finished.
+        That is, either all writes should be submitted on the current stream
+        (self._cuda_stream) or, if on a different stream, that stream should
+        synchronize with the current stream.
+
+        Args:
+            buf: The torch.Tensor to send. It should already be on this
+                actor's default device.
+            peer_rank: The rank of the actor to send to.
+        """
+        if self._closed:
+            raise RayChannelError("RayPPCommunicator has been destroyed.")
+
+        assert self._comm is not None
+        self._comm.send(buf, peer_rank)
+
+    def recv(
+        self,
+        shape: tuple[int, ...],
+        dtype: "torch.dtype",
+        peer_rank: int,
+        allocator: TorchTensorAllocator,
+    ) -> "torch.Tensor":
+        """
+        Receive a torch.Tensor from a peer and synchronize the current stream.
+
+        After this call returns, the receive buffer is safe to read from
+        any stream. An RayChannelError will be raised if an error occurred
+        (e.g., remote actor died), and the buffer is not safe to read.
+
+        Args:
+            shape: The shape of the tensor to receive.
+            dtype: The dtype of the tensor to receive.
+            peer_rank: The rank of the actor to receive from.
+            allocator: The allocator to use to create the received tensor.
+                This is ignored for this implementation.
+        """
+        if self._closed:
+            raise RayChannelError("RayPPCommunicator has been destroyed.")
+
+        assert self._comm is not None
+        size = torch.Size(shape)
+        buf = self._comm.recv(size, dtype, src=peer_rank)
+
+        # Buffer values are undefined if NCCL ops are aborted. Therefore, we
+        # need to synchronize here and check that the channel is still
+        # open to ensure that the receive buffer is valid.
+        # TODO(swang): Avoid CUDA synchronization.
+        current_stream().synchronize()
+
+        if self._closed:
+            raise RayChannelError("RayPPCommunicator has been destroyed.")
+        return buf
+
+    def allgather(
+        self,
+        send_buf: "torch.Tensor",
+        recv_buf: "torch.Tensor",
+    ):
+        raise NotImplementedError("allgather is not supported")
+
+    def allreduce(
+        self,
+        send_buf: "torch.Tensor",
+        recv_buf: "torch.Tensor",
+        op: ReduceOp = ReduceOp.SUM,
+    ):
+        raise NotImplementedError("allreduce is not supported")
+
+    def reducescatter(
+        self,
+        send_buf: "torch.Tensor",
+        recv_buf: "torch.Tensor",
+        op: ReduceOp = ReduceOp.SUM,
+    ):
+        raise NotImplementedError("reducescatter is not supported")
+
+    @property
+    def recv_stream(self):
+        return torch.cuda.StreamContext(current_stream())
+
+    @property
+    def send_stream(self):
+        return torch.cuda.StreamContext(current_stream())
+
+    def destroy(self) -> None:
+        # Just sets a flag, vLLM manages the lifecycle of the underlying
+        # _PP GroupCoordinator.
+        self._closed = True
+
+    def get_transport_name(self) -> str:
+        return "nccl"
+
+    @classmethod
+    def generate_communicator_id(cls) -> Any:
+        return uuid.uuid4()

vllm/distributed/device_communicators/shm_broadcast.py

@@ -0,0 +1,947 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import functools
+import pickle
+import sys
+import threading
+import time
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from multiprocessing import shared_memory
+from pickle import PickleBuffer
+from typing import TYPE_CHECKING, Any, cast
+from unittest.mock import patch
+
+import torch
+import torch.distributed as dist
+import zmq
+from torch.distributed import ProcessGroup
+from zmq import (  # type: ignore
+    IPV6,  # type: ignore
+    PUB,
+    SUB,
+    SUBSCRIBE,
+    XPUB,
+    XPUB_VERBOSE,
+    Context,
+)
+
+import vllm.envs as envs
+from vllm.distributed.utils import StatelessProcessGroup, sched_yield
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+from vllm.utils.network_utils import (
+    get_ip,
+    get_open_port,
+    get_open_zmq_inproc_path,
+    get_open_zmq_ipc_path,
+    is_valid_ipv6_address,
+)
+
+logger = init_logger(__name__)
+
+
+SPINLOOP_EXT_ENABLED = False
+if envs.VLLM_USE_SPINLOOP_EXT:
+    try:
+        from vllm.spinloop import spinloop
+
+        SPINLOOP_EXT_ENABLED = True
+    except ImportError:
+        logger.warning(
+            "spinloop extension could not be loaded, disabling VLLM_USE_SPINLOOP_EXT!"
+        )
+SPINLOOP_TIMEOUT_SECONDS = 0.1
+
+if TYPE_CHECKING:
+    from _typeshed import SizedBuffer
+
+VLLM_RINGBUFFER_WARNING_INTERVAL = envs.VLLM_RINGBUFFER_WARNING_INTERVAL
+
+from_bytes_big = functools.partial(int.from_bytes, byteorder="big")
+
+
+# Memory fence for cross-process shared memory visibility.
+# Required for correct producer-consumer synchronization when using
+# shared memory without locks.
+_memory_fence_lock = threading.Lock()
+
+
+def memory_fence():
+    """
+    Full memory barrier for shared memory synchronization.
+
+    Ensures all prior memory writes are visible to other processes before
+    any subsequent reads. This is critical for lock-free producer-consumer
+    patterns using shared memory.
+
+    Implementation acquires and immediately releases a lock. Python's
+    threading.Lock provides sequentially consistent memory barrier semantics
+    across all major platforms (POSIX, Windows). This is a lightweight
+    operation (~20ns) that guarantees:
+    - All stores before the barrier are visible to other threads/processes
+    - All loads after the barrier see the latest values
+    """
+    # Lock acquire/release provides full memory barrier semantics.
+    # Using context manager ensures lock release even on exceptions.
+    with _memory_fence_lock:
+        pass
+
+
+def to_bytes_big(value: int, size: int) -> bytes:
+    return value.to_bytes(size, byteorder="big")
+
+
+LONG_WAIT_TIME_LOG_MSG = (
+    "No available shared memory broadcast block found "
+    "in %d seconds. This typically happens "
+    "when some processes are hanging or doing some "
+    "time-consuming work (e.g. compilation, "
+    "weight/kv cache quantization)."
+)
+
+
+class SpinCondition:
+    """
+    This class implements an interface similar to a threading.Condition. It
+    allows a writer to notify readers to wake up and read from the shared memory
+    buffer. This notification is done over a zmq socket.
+
+    For optimal performance under load we don't want the readers to need to poll
+    the zmq socket for every read. So the `wait` method here will return
+    immediately when reads are frequent, and will only enter "idle mode" and
+    await a notification on the zmq socket after a period of inactivity. This
+    allows the readers to spin quickly, hence "SpinCondition".
+
+    To support clean shutdown, a separate thread in the reader's process must be
+    able to wake the reader so that it can exit. A separate cancel() method is
+    implemented with an in-process socket to allow this interruption.
+    """
+
+    def __init__(
+        self,
+        is_reader: bool,
+        context: zmq.Context,
+        notify_address: str,
+        busy_loop_s: float = 1,
+    ):
+        self.is_reader = is_reader
+
+        if is_reader:
+            # Time of last shm buffer read
+            self.last_read = time.monotonic()
+
+            # Time to keep busy-looping on the shm buffer before going idle
+            self.busy_loop_s = busy_loop_s
+
+            # Readers subscribe to write notifications
+            self.local_notify_socket: zmq.Socket = context.socket(SUB)
+            # Set zmq.CONFLATE to only keep the last message that the socket
+            # receives. This prevents us from piling up notification messages
+            # under high load when we aren't polling the socket.
+            self.local_notify_socket.setsockopt(zmq.CONFLATE, 1)
+            # Subscribe to all messages on the socket
+            self.local_notify_socket.setsockopt_string(SUBSCRIBE, "")
+            self.local_notify_socket.connect(notify_address)
+
+            # Readers require a process-local socket to poll for cancellation
+            cancel_path = get_open_zmq_inproc_path()
+            self.write_cancel_socket: zmq.Socket = context.socket(zmq.PAIR)
+            self.write_cancel_socket.bind(cancel_path)
+            self.read_cancel_socket: zmq.Socket = context.socket(zmq.PAIR)
+            self.read_cancel_socket.connect(cancel_path)
+
+            # Poller allows waiting on either `.notify()` or `.cancel()`
+            self.poller = zmq.Poller()
+            self.poller.register(self.read_cancel_socket, zmq.POLLIN)
+            self.poller.register(self.local_notify_socket, zmq.POLLIN)
+        else:
+            # Writer side publishes write notifications
+            self.local_notify_socket: zmq.Socket = context.socket(PUB)  # type: ignore
+            # Set high water mark to 1 - we don't need to send a massive amount of
+            # pings during busy operation. PUB sockets will silently drop subsequent
+            # messages after the high water mark is reached.
+            self.local_notify_socket.setsockopt(zmq.SNDHWM, 1)
+            self.local_notify_socket.bind(notify_address)
+
+            self.last_read = 0
+            self.busy_loop_s = 0
+            self.read_cancel_socket = None
+            self.write_cancel_socket = None
+            self.poller = None
+
+    def record_read(self):
+        self.last_read = time.monotonic()
+
+    def cancel(self):
+        # Sends cancellation ping that will cause the reader to wake up.
+        # This is done from a monitor thread in the same process as the reader.
+        if self.is_reader:
+            logger.debug("Canceling waiting reads on SHM Buffer")
+            self.write_cancel_socket.send(b"\x00")
+
+    def wait(self, timeout_ms: int | None = None) -> None:
+        """Wait for data on the shared memory buffer.
+
+        Yields the scheduler then returns immediately if it has been less than
+        self.busy_loop_s since the last read.
+
+        Otherwise, enters idle mode and awaits a socket ping for at most
+        `timeout_ms` milliseconds, or indefinitely if timeout_ms is None.
+        """
+        assert self.is_reader, "Only readers can wait"
+
+        current_time = time.monotonic()
+        if current_time <= self.last_read + self.busy_loop_s:
+            sched_yield()
+        else:
+            events = dict(self.poller.poll(timeout=timeout_ms))
+
+            if self.read_cancel_socket in events:
+                logger.debug("Poller received cancel event")
+            elif self.local_notify_socket in events:
+                logger.debug("Poller received notify event")
+                # Since zmq.CONFLATE is set, there will only be one notification
+                # to read from the socket
+                self.local_notify_socket.recv(flags=zmq.NOBLOCK, copy=False)
+            else:
+                logger.debug("Poller timed out")
+
+    def notify(self):
+        """Notifies all readers to wake up"""
+        assert not self.is_reader, "Only writers can notify"
+        self.local_notify_socket.send(b"\x00")
+
+
+class ShmRingBuffer:
+    def __init__(
+        self,
+        n_reader: int,
+        max_chunk_bytes: int,
+        max_chunks: int,
+        name: str | None = None,
+    ):
+        """
+        A shared memory ring buffer implementation for broadcast communication.
+        Essentially, it is a queue where only one will `enqueue` and multiple
+        will `dequeue`. The max size of each item, together with the max number
+        of items that can be stored in the buffer are known in advance.
+        In this case, we don't need to synchronize the access to
+         the buffer.
+
+        Buffer memory layout:
+                  data                                 metadata
+                    |                                      |
+                    | (current_idx)                        | (current_idx)
+                    v                                      v
+        +-------------------------------+----------------------------------------+
+        | chunk0 | chunk1 | ... | chunk | metadata0 | metadata1 | ... | metadata |
+        +-------------------------------+----------------------------------------+
+        | max_chunks x max_chunk_bytes  | max_chunks x (1 + n_reader) bytes      |
+
+        metadata memory layout: each byte is a flag, the first byte is the written
+        flag, and the rest are reader flags. The flags are set to 0 by default.
+        +--------------+--------------+--------------+-----+--------------+
+        | written_flag | reader0_flag | reader1_flag | ... | readerN_flag |
+        +--------------+--------------+--------------+-----+--------------+
+
+        The state of metadata is as follows:
+
+        (case 1) 0???...???: the block is not written yet, cannot read, can write
+        (case 2) 1000...000: the block is just written, can read, cannot write
+        (case 3) 1???...???: the block is written and read by some readers, can read if not read, cannot write
+        (case 4) 1111...111: the block is written and read by all readers, cannot read, can write
+
+        State transition for readers:
+
+        When a reader finds a block that it can read (case 2 or 3), it can yield the block for caller to read.
+        Only after the caller finishes reading the block, the reader can mark the block as read.
+        Readers only mark the block as read (from 0 to 1), the writer marks the block as ready to read (from 1 to 0).
+
+        State transition for writer:
+
+        When the writer writes to a block (case 1 or 4), it first resets the written flag to 0, converting either case
+        to case 1. Then it can yield the block for caller to write. After the caller finishes writing the block, the writer
+        can reset the reader flags to 0, and mark the block as written (from 0 to 1).
+        NOTE: the order is important here, first reset the reader flags (so that we are still in case 1), then mark the block as written. The state transition is atomic. If we do it in the reverse order, it will go through case 3 and then back to case 2, and readers might read the intermediate case 3, which is not correct.
+
+        During creation, `name` is None and the buffer is created. We can pass the
+        created object to other processes by pickling it. The other processes will
+        get the name of the shared memory and open it, so that they can access the
+        same shared memory buffer.
+        """  # noqa
+        self.n_reader = n_reader
+        self.metadata_size = 1 + n_reader
+        self.max_chunk_bytes = max_chunk_bytes
+        self.max_chunks = max_chunks
+        self.total_bytes_of_buffer = (
+            self.max_chunk_bytes + self.metadata_size
+        ) * self.max_chunks
+        self.data_offset = 0
+        self.metadata_offset = self.max_chunk_bytes * self.max_chunks
+
+        if name is None:
+            # we are creating a buffer
+            self.is_creator = True
+            self.shared_memory = shared_memory.SharedMemory(
+                create=True, size=self.total_bytes_of_buffer
+            )
+            assert self.shared_memory.buf is not None, "Buffer was not created"
+            # initialize the metadata section to 0
+            with self.shared_memory.buf[self.metadata_offset :] as metadata_buffer:
+                torch.frombuffer(metadata_buffer, dtype=torch.uint8).fill_(0)
+        else:
+            # we are opening an existing buffer
+            self.is_creator = False
+            # fix to https://stackoverflow.com/q/62748654/9191338
+            # Python incorrectly tracks shared memory even if it is not
+            # created by the process. The following patch is a workaround.
+            with patch(
+                "multiprocessing.resource_tracker.register",
+                lambda *args, **kwargs: None,
+            ):
+                try:
+                    self.shared_memory = shared_memory.SharedMemory(name=name)
+                    # See https://docs.python.org/3/library/multiprocessing.shared_memory.html # noqa
+                    # Some platforms allocate memory based on page size,
+                    # so the shared memory block size may be larger or equal
+                    # to the requested size. The size parameter is ignored
+                    # when attaching to an existing block.
+                    assert self.shared_memory.size >= self.total_bytes_of_buffer
+                except FileNotFoundError:
+                    # we might deserialize the object in a different node
+                    # in this case, this object is not used,
+                    # and we should suppress the error
+                    pass
+
+    def handle(self):
+        return (
+            self.n_reader,
+            self.max_chunk_bytes,
+            self.max_chunks,
+            self.shared_memory.name,
+        )
+
+    def __reduce__(self):
+        return (
+            self.__class__,
+            self.handle(),
+        )
+
+    def __del__(self):
+        if hasattr(self, "shared_memory"):
+            self.shared_memory.close()
+            if self.is_creator:
+                self.shared_memory.unlink()
+
+    @contextmanager
+    def get_data(self, current_idx: int):
+        start = self.data_offset + current_idx * self.max_chunk_bytes
+        end = start + self.max_chunk_bytes
+        assert self.shared_memory.buf is not None, "Buffer has been closed"
+        with self.shared_memory.buf[start:end] as buf:
+            yield buf
+
+    @contextmanager
+    def get_metadata(self, current_idx: int):
+        start = self.metadata_offset + current_idx * self.metadata_size
+        end = start + self.metadata_size
+        assert self.shared_memory.buf is not None, "Buffer has been closed"
+        with self.shared_memory.buf[start:end] as buf:
+            yield buf
+
+
+@dataclass
+class Handle:
+    local_reader_ranks: list[int] = field(default_factory=list)
+
+    buffer_handle: tuple[int, int, int, str] | None = None
+    local_subscribe_addr: str | None = None
+    local_notify_addr: str | None = None
+    remote_subscribe_addr: str | None = None
+    remote_addr_ipv6: bool = False
+
+
+class MessageQueue:
+    def __init__(
+        self,
+        n_reader,  # number of all readers
+        n_local_reader,  # number of local readers through shared memory
+        local_reader_ranks: list[int] | None = None,
+        # Default of 24MiB chosen to be large enough to accommodate grammar
+        # bitmask tensors for large batches (1024 requests).
+        max_chunk_bytes: int = 1024 * 1024 * 24,
+        max_chunks: int = 10,
+        connect_ip: str | None = None,
+    ):
+        if local_reader_ranks is None:
+            local_reader_ranks = list(range(n_local_reader))
+        else:
+            assert len(local_reader_ranks) == n_local_reader
+        self.n_local_reader = n_local_reader
+        n_remote_reader = n_reader - n_local_reader
+        self.n_remote_reader = n_remote_reader
+        self.shutting_down = False
+        context = Context()
+
+        if n_local_reader > 0:
+            # for local readers, we will:
+            # 1. create a shared memory ring buffer to communicate small data
+            # 2. create a publish-subscribe socket to communicate large data
+            self.buffer = ShmRingBuffer(n_local_reader, max_chunk_bytes, max_chunks)
+
+            # XPUB is very similar to PUB,
+            # except that it can receive subscription messages
+            # to confirm the number of subscribers
+            self.local_socket = context.socket(XPUB)
+            # set the verbose option so that we can receive every subscription
+            # message. otherwise, we will only receive the first subscription
+            # see http://api.zeromq.org/3-3:zmq-setsockopt for more details
+            self.local_socket.setsockopt(XPUB_VERBOSE, True)
+            local_subscribe_addr = get_open_zmq_ipc_path()
+            logger.debug("Binding to %s", local_subscribe_addr)
+            self.local_socket.bind(local_subscribe_addr)
+
+            self.current_idx = 0
+
+            # Create the notification side of the SpinCondition
+            local_notify_addr = get_open_zmq_ipc_path()
+            self._spin_condition = SpinCondition(
+                is_reader=False, context=context, notify_address=local_notify_addr
+            )
+        else:
+            self.buffer = None  # type: ignore
+            local_subscribe_addr = None
+            self.local_socket = None
+            self.current_idx = -1
+            local_notify_addr = None
+            self._spin_condition = None  # type: ignore
+
+        remote_addr_ipv6 = False
+        if n_remote_reader > 0:
+            # for remote readers, we will:
+            # create a publish-subscribe socket to communicate large data
+            if not connect_ip:
+                connect_ip = get_ip()
+            self.remote_socket = context.socket(XPUB)
+            self.remote_socket.setsockopt(XPUB_VERBOSE, True)
+            remote_subscribe_port = get_open_port()
+            if is_valid_ipv6_address(connect_ip):
+                self.remote_socket.setsockopt(IPV6, 1)
+                remote_addr_ipv6 = True
+                connect_ip = f"[{connect_ip}]"
+            socket_addr = f"tcp://{connect_ip}:{remote_subscribe_port}"
+            self.remote_socket.bind(socket_addr)
+            remote_subscribe_addr = f"tcp://{connect_ip}:{remote_subscribe_port}"
+        else:
+            remote_subscribe_addr = None
+            self.remote_socket = None
+
+        self._is_writer = True
+        self._is_local_reader = False
+        self.local_reader_rank = -1
+        # rank does not matter for remote readers
+        self._is_remote_reader = False
+
+        self.handle = Handle(
+            local_reader_ranks=local_reader_ranks,
+            buffer_handle=self.buffer.handle() if self.buffer is not None else None,
+            local_subscribe_addr=local_subscribe_addr,
+            local_notify_addr=local_notify_addr,
+            remote_subscribe_addr=remote_subscribe_addr,
+            remote_addr_ipv6=remote_addr_ipv6,
+        )
+
+        logger.debug("vLLM message queue communication handle: %s", self.handle)
+
+    def export_handle(self) -> Handle:
+        return self.handle
+
+    @staticmethod
+    def create_from_handle(handle: Handle, rank) -> "MessageQueue":
+        self = MessageQueue.__new__(MessageQueue)
+        self.handle = handle
+        self._is_writer = False
+
+        context = Context()
+
+        if rank in handle.local_reader_ranks:
+            assert handle.buffer_handle is not None
+            self.buffer = ShmRingBuffer(*handle.buffer_handle)
+            self.current_idx = 0
+            self.local_reader_rank = handle.local_reader_ranks.index(rank)
+            self._is_local_reader = True
+            self._is_remote_reader = False
+
+            self.local_socket = context.socket(SUB)
+            self.local_socket.setsockopt_string(SUBSCRIBE, "")
+            socket_addr = handle.local_subscribe_addr
+            logger.debug("Connecting to %s", socket_addr)
+            self.local_socket.connect(socket_addr)
+
+            self.remote_socket = None
+            assert isinstance(handle.local_notify_addr, str)
+            self._spin_condition = SpinCondition(
+                is_reader=True, context=context, notify_address=handle.local_notify_addr
+            )
+        else:
+            self.buffer = None  # type: ignore
+            self.current_idx = -1
+            self.local_reader_rank = -1
+            self._is_local_reader = False
+            self._is_remote_reader = True
+
+            self.local_socket = None
+
+            self.remote_socket = context.socket(SUB)
+            self.remote_socket.setsockopt_string(SUBSCRIBE, "")
+            if handle.remote_addr_ipv6:
+                self.remote_socket.setsockopt(IPV6, 1)
+            socket_addr = handle.remote_subscribe_addr
+            logger.debug("Connecting to %s", socket_addr)
+            self.remote_socket.connect(socket_addr)
+            self._spin_condition = None  # type: ignore
+
+        self.shutting_down = False
+        return self
+
+    def wait_until_ready(self):
+        """This is a collective operation. All processes (including the
+        readers and the writer) should call this function.
+        """
+        if self._is_writer:
+            # wait for all readers to connect
+
+            # local readers
+            for i in range(self.n_local_reader):
+                # wait for subscription messages from all local readers
+                self.local_socket.recv()
+            if self.n_local_reader > 0:
+                # send a message to all local readers
+                # to make sure the publish channel is working
+                self.local_socket.send(b"READY")
+
+            # remote readers
+            for i in range(self.n_remote_reader):
+                # wait for subscription messages from all remote readers
+                self.remote_socket.recv()
+            if self.n_remote_reader > 0:
+                # send a message to all remote readers
+                # to make sure the publish channel is working
+                self.remote_socket.send(b"READY")
+        elif self._is_local_reader:
+            # wait for the writer to send a message
+            recv = self.local_socket.recv()
+            assert recv == b"READY"
+        elif self._is_remote_reader:
+            # wait for the writer to send a message
+            recv = self.remote_socket.recv()
+            assert recv == b"READY"
+
+    def shutdown(self):
+        """If this is an idle reader, wakes it up so it can clean up and shut
+        down"""
+        self.shutting_down = True
+        if self._spin_condition is not None:
+            self._spin_condition.cancel()
+
+    @contextmanager
+    def acquire_write(self, timeout: float | None = None):
+        assert self._is_writer, "Only writers can acquire write"
+        start_time = time.monotonic()
+        n_warning = 1
+        while True:
+            with self.buffer.get_metadata(self.current_idx) as metadata_buffer:
+
+                def check():
+                    memory_fence()
+                    read_count = sum(metadata_buffer[1:])
+                    written_flag = metadata_buffer[0]
+                    return not (written_flag and read_count != self.buffer.n_reader)
+
+                if SPINLOOP_EXT_ENABLED and not check():
+                    spinloop(metadata_buffer, check, timeout=SPINLOOP_TIMEOUT_SECONDS)
+
+                if not check():
+                    # this block is written and not read by all readers
+                    # for writers, `self.current_idx` is the next block to write
+                    # if this block is not ready to write,
+                    # we need to wait until it is read by all readers
+
+                    # Release the processor to other threads
+                    sched_yield()
+
+                    # if we time out, raise an exception
+                    elapsed = time.monotonic() - start_time
+                    if timeout is not None and elapsed > timeout:
+                        raise TimeoutError
+
+                    # if we wait for a long time, log a message
+                    if elapsed > VLLM_RINGBUFFER_WARNING_INTERVAL * n_warning:
+                        logger.info(
+                            LONG_WAIT_TIME_LOG_MSG, VLLM_RINGBUFFER_WARNING_INTERVAL
+                        )
+                        n_warning += 1
+
+                    continue
+                # found a block that is either
+                # (1) not written
+                # (2) read by all readers
+
+                # mark the block as not written
+                metadata_buffer[0] = 0
+                # let caller write to the buffer
+                with self.buffer.get_data(self.current_idx) as buf:
+                    yield buf
+
+                # caller has written to the buffer
+                # NOTE: order is important here
+                # first set the read flags to 0
+                # then set the written flag to 1
+                # otherwise, the readers may think they already read the block
+                for i in range(1, self.buffer.n_reader + 1):
+                    # set read flag to 0, meaning it is not read yet
+                    metadata_buffer[i] = 0
+                # Memory fence here ensures the order of the buffer and flag
+                # writes. This guarantees that when `metadata_buffer[0] = 1` is
+                # visible to readers, `buf` can be completely ready. Without
+                # this, some CPU architectures with weak ordering may incur
+                # memory inconsistency.
+                memory_fence()
+                # mark the block as written
+                metadata_buffer[0] = 1
+                # Memory fence ensures the write is visible to readers on other cores
+                # before we proceed. Without this, readers may spin indefinitely
+                # waiting for a write that's stuck in our CPU's store buffer.
+                memory_fence()
+                self.current_idx = (self.current_idx + 1) % self.buffer.max_chunks
+                break
+
+    class ReadTimeoutWithWarnings:
+        def __init__(self, timeout: float | None, should_warn: bool) -> None:
+            self.started = time.monotonic()
+            self.deadline = sys.maxsize if timeout is None else self.started + timeout
+
+            # if should_warn, we need to wake up periodically to log
+            self.warning_wait_time_ms: int | None = (
+                VLLM_RINGBUFFER_WARNING_INTERVAL * 1000 if should_warn else None
+            )
+
+            self._should_warn = should_warn
+            self.n_warning = 1
+            self.timeout = timeout
+
+        def timeout_ms(self) -> int | None:
+            """Returns a timeout that is:
+            - min(time to deadline, time to next warning) if we're logging warnings
+            - time to deadline, if we're not logging warnings
+            - None if the timeout is None and we're not logging warnings
+            - raise TimeoutError if we are past the deadline
+            """
+            warning_wait_time = self.warning_wait_time_ms
+            if self.timeout is None:
+                return warning_wait_time
+
+            time_left_ms = int((self.deadline - time.monotonic()) * 1000)
+            if time_left_ms <= 0:
+                raise TimeoutError
+
+            if warning_wait_time and warning_wait_time < time_left_ms:
+                return warning_wait_time
+
+            return time_left_ms
+
+        def should_warn(self) -> bool:
+            """Returns true if it's time to log a warning for a timeout that is not
+            indefinite"""
+            if self._should_warn:
+                elapsed = time.monotonic() - self.started
+                if elapsed >= VLLM_RINGBUFFER_WARNING_INTERVAL * self.n_warning:
+                    self.n_warning += 1
+                    return True
+            return False
+
+    @contextmanager
+    def acquire_read(
+        self,
+        timeout: float | None = None,
+        indefinite: bool = False,
+    ):
+        assert self._is_local_reader, "Only readers can acquire read"
+        read_timeout = self.ReadTimeoutWithWarnings(
+            timeout=timeout, should_warn=not indefinite
+        )
+        with self.buffer.get_metadata(self.current_idx) as metadata_buffer:
+            while True:
+
+                def check():
+                    memory_fence()
+                    read_flag = metadata_buffer[self.local_reader_rank + 1]
+                    written_flag = metadata_buffer[0]
+                    return not (not written_flag or read_flag)
+
+                if SPINLOOP_EXT_ENABLED and not check():
+                    spinloop(
+                        metadata_buffer[0 : self.local_reader_rank + 1],
+                        check,
+                        timeout=SPINLOOP_TIMEOUT_SECONDS,
+                    )
+
+                if not check():
+                    # this block is either
+                    # (1) not written
+                    # (2) already read by this reader
+
+                    # for readers, `self.current_idx` is the next block to read
+                    # if this block is not ready,
+                    # we need to wait until it is written
+                    self._spin_condition.wait(timeout_ms=read_timeout.timeout_ms())
+
+                    if self.shutting_down:
+                        raise RuntimeError("cancelled")
+
+                    # if we wait for a long time, log a message
+                    if read_timeout.should_warn():
+                        logger.info(
+                            LONG_WAIT_TIME_LOG_MSG, VLLM_RINGBUFFER_WARNING_INTERVAL
+                        )
+
+                    continue
+                # found a block that is not read by this reader
+                # let caller read from the buffer
+                with self.buffer.get_data(self.current_idx) as buf:
+                    yield buf
+
+                # caller has read from the buffer
+                # set the read flag
+                metadata_buffer[self.local_reader_rank + 1] = 1
+                # Memory fence ensures the read flag is visible to the writer.
+                # Without this, writer may not see our read completion and
+                # could wait indefinitely for all readers to finish.
+                memory_fence()
+                self.current_idx = (self.current_idx + 1) % self.buffer.max_chunks
+
+                self._spin_condition.record_read()
+                break
+
+    def enqueue(self, obj, timeout: float | None = None):
+        """Write to message queue with optional timeout (in seconds)"""
+        assert self._is_writer, "Only writers can enqueue"
+        all_buffers: list[SizedBuffer] = [b""]
+        total_bytes = 6  # 2 bytes for oob buffer count, 4 for main buffer size
+
+        def oob_callback(buf: PickleBuffer) -> bool:
+            raw_buf = buf.raw()
+            if len(raw_buf) < 1024 * 1024:
+                # In-line buffers smaller than 1MiB.
+                return True
+            all_buffers.append(raw_buf)
+            nonlocal total_bytes
+            total_bytes += len(raw_buf) + 4
+            return False
+
+        all_buffers[0] = pickle.dumps(
+            obj, protocol=pickle.HIGHEST_PROTOCOL, buffer_callback=oob_callback
+        )
+        if self.n_local_reader > 0:
+            if total_bytes + len(all_buffers[0]) >= self.buffer.max_chunk_bytes:
+                with self.acquire_write(timeout) as buf:
+                    buf[0] = 1  # overflow
+                self.local_socket.send_multipart(all_buffers, copy=False)
+            else:
+                # Byte 0: 0
+                # Bytes 1-2: Count of buffers
+                # Then each buffer follows, preceded by 4 bytes containing its length:
+                # [4 byte int L][L bytes of buffer content] ...
+                with self.acquire_write(timeout) as buf:
+                    buf[0] = 0  # not overflow
+                    offset = 3
+                    buf[1:offset] = to_bytes_big(len(all_buffers), 2)  # oob buf count
+                    for buffer in all_buffers:
+                        buf_len = len(buffer)
+                        # prepend each buffer with 4 bytes containing its size.
+                        buf_offset = offset + 4
+                        buf[offset:buf_offset] = to_bytes_big(buf_len, 4)
+                        buf[buf_offset : (offset := buf_offset + buf_len)] = buffer
+
+            self._spin_condition.notify()
+
+        if self.n_remote_reader > 0:
+            self.remote_socket.send_multipart(all_buffers, copy=False)
+
+    def dequeue(
+        self,
+        timeout: float | None = None,
+        indefinite: bool = False,
+    ):
+        """Read from message queue with optional timeout (in seconds)"""
+        if self._is_local_reader:
+            with self.acquire_read(timeout, indefinite) as buf:
+                overflow = buf[0] == 1
+                if not overflow:
+                    offset = 3
+                    buf_count = from_bytes_big(buf[1:offset])
+                    all_buffers = []
+                    for i in range(buf_count):
+                        buf_offset = offset + 4
+                        buf_len = from_bytes_big(buf[offset:buf_offset])
+                        offset = buf_offset + buf_len
+                        all_buffers.append(buf[buf_offset:offset])
+                    obj = pickle.loads(all_buffers[0], buffers=all_buffers[1:])
+            if overflow:
+                obj = MessageQueue.recv(self.local_socket, timeout)
+        elif self._is_remote_reader:
+            obj = MessageQueue.recv(self.remote_socket, timeout)
+        else:
+            raise RuntimeError("Only readers can dequeue")
+        return obj
+
+    @staticmethod
+    def recv(socket: zmq.Socket, timeout: float | None) -> Any:
+        timeout_ms = None if timeout is None else int(timeout * 1000)
+        if not socket.poll(timeout=timeout_ms):
+            raise TimeoutError
+        recv, *recv_oob = socket.recv_multipart(copy=False)
+        return pickle.loads(recv, buffers=recv_oob)
+
+    def broadcast_object(self, obj=None):
+        if self._is_writer:
+            self.enqueue(obj)
+            return obj
+        return self.dequeue()
+
+    @staticmethod
+    def create_from_process_group_single_reader(
+        pg: ProcessGroup,
+        max_chunk_bytes,
+        max_chunks,
+        reader_rank: int = 0,
+        blocking: bool = False,
+    ) -> tuple["MessageQueue", list[Handle]]:
+        """
+        Creates a MessageQueue for a process group with a single reader.
+
+        This method is designed for scenarios where only one process (the reader)
+        will consume messages, and all other processes are writers. It sets up
+        the shared memory buffer and communication handles accordingly, and
+        gathers the handles from all processes to the reader.
+
+        Args:
+            pg (ProcessGroup): The torch distributed process group.
+            max_chunk_bytes (int): Maximum size in bytes for each chunk in the buffer.
+            max_chunks (int): Maximum number of chunks in the buffer.
+            reader_rank (int, optional): The global rank that will act as the reader.
+                Defaults to 0.
+            blocking (bool, optional): If True, blocks until all processes are ready.
+                Defaults to False.
+
+        Returns:
+            tuple[MessageQueue, list[Handle]]:
+            The MessageQueue instance for the calling process,
+            and a list of handles (only non-empty for the reader process).
+        """
+        from vllm.platforms.interface import get_assigned_physical_gpu_ids
+
+        assigned_physical_gpu_ids = get_assigned_physical_gpu_ids()
+        if assigned_physical_gpu_ids is not None:
+            local_size = len(assigned_physical_gpu_ids)
+        else:
+            local_size = current_platform.device_count()
+        rank = dist.get_rank()
+        same_node = rank // local_size == reader_rank // local_size
+        buffer_io = MessageQueue(
+            n_reader=1,
+            n_local_reader=1 if same_node else 0,
+            max_chunk_bytes=max_chunk_bytes,
+            max_chunks=max_chunks,
+        )
+        handle = buffer_io.export_handle()
+        handles = [None] * dist.get_world_size(pg) if rank == reader_rank else None
+        dist.gather_object(handle, handles, dst=reader_rank, group=pg)
+        if blocking:
+            buffer_io.wait_until_ready()
+        return buffer_io, cast(list[Handle], handles or [])
+
+    @staticmethod
+    def create_from_process_group(
+        pg: ProcessGroup | StatelessProcessGroup,
+        max_chunk_bytes,
+        max_chunks,
+        writer_rank: int = 0,
+        external_writer_handle=None,
+        blocking: bool = True,
+    ) -> "MessageQueue":
+        """
+        Creates a MessageQueue for a distributed process group with one writer and
+        multiple readers.
+
+        This method is designed for scenarios where one process (the writer) sends
+        messages, and all other processes (the readers) receive messages. It sets up
+        the shared memory buffer and socket communication handles accordingly, and
+        broadcasts the handle from the writer to all readers.
+
+        Args:
+            pg (ProcessGroup | StatelessProcessGroup): The torch distributed process
+                group.
+            max_chunk_bytes (int): Maximum size in bytes for each chunk in the buffer.
+            max_chunks (int): Maximum number of chunks in the buffer.
+            writer_rank (int, optional): The global rank that will act as the writer.
+                Defaults to 0.
+            external_writer_handle (Handle, optional): Used when there is a handle
+                from an external Message Queue. If provided, use this handle to init
+                PG writer message queue instead of creating a new one. Defaults to None.
+            blocking (bool, optional): If True, blocks until all processes are ready.
+                Defaults to True.
+
+        Returns:
+            MessageQueue: The MessageQueue instance for the calling process.
+
+        """
+        if isinstance(pg, ProcessGroup):
+            group_rank = dist.get_rank(pg)
+            group_world_size = dist.get_world_size(pg)
+            global_ranks = dist.get_process_group_ranks(pg)
+        else:
+            group_rank = pg.rank
+            group_world_size = pg.world_size
+            global_ranks = list(range(pg.world_size))
+        from vllm.distributed.parallel_state import in_the_same_node_as
+
+        status = in_the_same_node_as(pg, source_rank=writer_rank)
+        if group_rank == writer_rank:
+            if external_writer_handle is not None:
+                buffer_io = MessageQueue.create_from_handle(
+                    external_writer_handle, group_rank
+                )
+            else:
+                same_node_ranks = [i for i, s in enumerate(status) if s]
+                n_reader = group_world_size - 1
+                n_local_reader = len(same_node_ranks) - 1
+                local_reader_ranks = [i for i in same_node_ranks if i != writer_rank]
+                buffer_io = MessageQueue(
+                    n_reader=n_reader,
+                    n_local_reader=n_local_reader,
+                    local_reader_ranks=local_reader_ranks,
+                    max_chunk_bytes=max_chunk_bytes,
+                    max_chunks=max_chunks,
+                )
+            handle = buffer_io.export_handle()
+            if isinstance(pg, ProcessGroup):
+                dist.broadcast_object_list(
+                    [handle], src=global_ranks[writer_rank], group=pg
+                )
+            else:
+                pg.broadcast_obj(handle, writer_rank)
+        else:
+            if isinstance(pg, ProcessGroup):
+                recv = [None]
+                dist.broadcast_object_list(
+                    recv, src=global_ranks[writer_rank], group=pg
+                )
+                handle = recv[0]  # type: ignore
+            else:
+                handle = pg.broadcast_obj(None, writer_rank)
+            buffer_io = MessageQueue.create_from_handle(handle, group_rank)
+        if blocking:
+            buffer_io.wait_until_ready()
+        return buffer_io

vllm/distributed/device_communicators/shm_object_storage.py

@@ -0,0 +1,709 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pickle
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Iterable
+from contextlib import contextmanager, suppress
+from dataclasses import dataclass
+from itertools import chain
+from multiprocessing import shared_memory
+from multiprocessing.synchronize import Lock as LockType
+from typing import Any
+from unittest.mock import patch
+
+import torch
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class SingleWriterShmRingBuffer:
+    """
+    A single-writer, multiple-reader ring buffer implementation using shared
+    memory. This class provides a thread-safe ring buffer where one process
+    can write data while multiple processes/threads can read from it.
+
+    Architecture:
+    - Uses shared memory for cross-process communication
+    - Maintains metadata for each allocated buffer chunk in the writer process
+    - Supports custom "is_free_fn" functions to determine when buffers can be
+      reused
+    - Each buffer chunk contains: `[4-byte id][4-byte size][actual_data]`
+
+    Key Concepts:
+    - monotonic_id_start/end: Track the range of active buffer IDs
+    - data_buffer_start/end: Track the physical memory range in use
+    - Automatic wraparound when reaching buffer end
+    - Lazy garbage collection based on is_free_fn checks
+
+    Example Usage Scenarios:
+
+    Scenario 1: Simple Linear Allocation
+    ```
+    Buffer size: 100 bytes
+    Initial state: [................................................. ]
+                   ^start=end(0)
+
+    After allocating 20 bytes (id=0):
+    [id:0|size:20|data........][...................................]
+    ^start(0)                  ^end(28)
+
+    After allocating 30 bytes (id=1):
+    [id:0|size:20|data........][id:1|size:30|data..............][..]
+    ^start(0)                                                   ^end(66)
+    ```
+
+    Scenario 2: Memory Reclamation
+    ```
+    Before freeing (both buffers still in use):
+    [id:0|size:20|data........][id:1|size:30|data..............][..]
+    ^start(0)                                                   ^end(66)
+
+    After id:0 is marked free by readers:
+    [FREED.................... ][id:1|size:30|data..............][..]
+                                ^start(28)                       ^end(66)
+
+    After both are freed:
+    [FREED..............................................][..]
+                                                         ^start=end(66)
+    ```
+
+    Scenario 3: Wraparound Allocation (continuing from Scenario 2)
+    ```
+    Starting from after memory reclamation in Scenario 2:
+    [FREED..............................................][..]
+                                                         ^start=end(66)
+
+    Allocate 40 bytes (id=2) - only 34 bytes available at end, so wraparound:
+    [id:2|size:40|data........................][FREED.............][..]
+                                              ^end(148)            ^start(66)
+    ```
+
+    Scenario 4: Error Handling - Out of Space
+    ```
+    Starting from after wraparound allocation in Scenario 3:
+    [id:2|size:40|data........................][FREED.............][..]
+                                              ^end(148)            ^start(66)
+
+    Trying to allocate 20 more bytes:
+    occupied_size_new = end + size - start = 148 + 28 - 66 > buffer_size(100)
+    -> Raises MemoryError: "Not enough space in the data buffer"
+    ```
+
+    Thread Safety:
+    - Single writer: Only one process/thread should write (allocate_buf)
+    - Multiple readers: Multiple processes/threads can read (access_buf)
+    - Reader synchronization handled by is_free_fn callback
+    - Writer handles garbage collection (free_buf) based on reader feedback
+
+    Memory Layout per Buffer Chunk:
+    `[4-byte monotonic_id][4-byte chunk_size][actual_data...]`
+    ^metadata_start                         ^data_start
+
+    The monotonic_id ensures data integrity - readers can verify they're
+    accessing the correct data even after buffer wraparound or reuse.
+    """
+
+    def __init__(
+        self,
+        data_buffer_size: int,
+        name: str | None = None,
+        create: bool = False,
+    ):
+        self.data_buffer_size = data_buffer_size
+        self.is_writer = create
+
+        self.ID_NBYTES = 4
+        self.ID_MAX = 2**31  # exclusive, so 2**31 - 1 is the max value
+        self.SIZE_NBYTES = 4
+        # 4 bytes for id, 4 bytes for buffer size
+        self.MD_SIZE = self.ID_NBYTES + self.SIZE_NBYTES
+        self.monotonic_id_end = 0
+        self.monotonic_id_start = 0
+        self.data_buffer_start = 0
+        self.data_buffer_end = 0
+
+        if create:
+            logger.debug("Creating new shared memory buffer: %s", name)
+            # we are creating a buffer
+            self.metadata: dict[int, int] = {}  # monotonic_id -> start address
+            self.shared_memory = shared_memory.SharedMemory(
+                create=True, size=self.data_buffer_size, name=name
+            )
+        else:
+            # we are opening an existing buffer
+            # fix to https://stackoverflow.com/q/62748654/9191338
+            # Python incorrectly tracks shared memory even if it is not
+            # created by the process. The following patch is a workaround.
+            with patch(
+                "multiprocessing.resource_tracker.register",
+                lambda *args, **kwargs: None,
+            ):
+                self.shared_memory = shared_memory.SharedMemory(name=name)
+                # See https://docs.python.org/3/library/multiprocessing.shared_memory.html # noqa
+                # Some platforms allocate memory based on page size,
+                # so the shared memory block size may be larger or equal
+                # to the requested size. The size parameter is ignored
+                # when attaching to an existing block.
+                assert self.shared_memory.size >= self.data_buffer_size
+
+        logger.debug(
+            "Shared memory created/opened with name: %s, size: %d",
+            self.shared_memory.name,
+            self.data_buffer_size,
+        )
+
+    def handle(self):
+        return (
+            self.data_buffer_size,
+            self.shared_memory.name,
+        )
+
+    def clear(self) -> None:
+        """Clear the ring buffer."""
+        assert self.is_writer, "Only the writer can clear the buffer."
+        self.metadata.clear()
+        self.monotonic_id_end = 0
+        self.monotonic_id_start = 0
+        self.data_buffer_start = 0
+        self.data_buffer_end = 0
+
+    def close(self) -> None:
+        """Close the shared memory."""
+        if hasattr(self, "shared_memory"):
+            self.shared_memory.close()
+            if self.is_writer:
+                with suppress(FileNotFoundError):
+                    self.shared_memory.unlink()
+
+    def __del__(self):
+        self.close()
+
+    def int2byte(self, integer: int) -> bytes:
+        """Convert an integer to bytes."""
+        return integer.to_bytes(self.ID_NBYTES, "little", signed=True)
+
+    def byte2int(self, byte_data: bytes) -> int:
+        """Convert bytes back to an integer."""
+        return int.from_bytes(byte_data, "little", signed=True)
+
+    def allocate_buf(self, size: int) -> tuple[int, int]:
+        """
+        Allocate a buffer `MD_SIZE` + `size` bytes in the shared memory.
+        Memory layout:
+        `[4-byte monotonic_id][4-byte size][buffer data...]`
+        """
+        assert self.is_writer, "Only the writer can allocate buffers."
+        assert size > 0, "Size must be greater than 0"
+        assert self.shared_memory.buf is not None, "Buffer has been closed"
+        size += self.MD_SIZE  # add metadata size to the buffer size
+        # reset to beginning if the buffer does have enough contiguous space
+        buffer_end_reset = self.data_buffer_end % self.data_buffer_size
+        if buffer_end_reset + size > self.data_buffer_size:
+            buffer_end_reset = (
+                self.data_buffer_end // self.data_buffer_size + 1
+            ) * self.data_buffer_size
+        else:  # no reset needed
+            buffer_end_reset = self.data_buffer_end
+
+        # check if we have enough space in the data buffer
+        # i.e. if the new end (self.data_buffer_end + size)
+        # exceeds the start of the data buffer
+        occupied_size_new = buffer_end_reset + size - self.data_buffer_start
+        if occupied_size_new > self.data_buffer_size:
+            raise MemoryError(
+                "Not enough space in the data buffer, "
+                "try calling free_buf() to free up space"
+            )
+        self.data_buffer_end = buffer_end_reset
+
+        # first 4 bytes as the monotonic id
+        buf_idx = self.data_buffer_end % self.data_buffer_size
+        self.shared_memory.buf[buf_idx : buf_idx + self.ID_NBYTES] = self.int2byte(
+            self.monotonic_id_end
+        )
+        # next 4 bytes as the size of the data buffer
+        self.shared_memory.buf[buf_idx + self.ID_NBYTES : buf_idx + self.MD_SIZE] = (
+            self.int2byte(size)
+        )
+
+        # record metadata
+        self.metadata[self.monotonic_id_end % self.ID_MAX] = self.data_buffer_end
+        # update buffer and monotonic id indices
+        current_buffer_end = self.data_buffer_end
+        current_id_end = self.monotonic_id_end
+        self.data_buffer_end += size
+        self.monotonic_id_end = (self.monotonic_id_end + 1) % self.ID_MAX
+        return current_buffer_end, current_id_end
+
+    @contextmanager
+    def access_buf(self, address: int):
+        assert self.shared_memory.buf is not None, "Buffer has been closed"
+        buf_idx = address % self.data_buffer_size
+
+        # read metadata
+        metadata_buff = self.shared_memory.buf[buf_idx : buf_idx + self.MD_SIZE]
+        id = self.byte2int(metadata_buff[: self.ID_NBYTES])
+        size = self.byte2int(metadata_buff[self.ID_NBYTES : self.MD_SIZE])
+
+        # yield the data buffer and metadata
+        data_buff = self.shared_memory.buf[buf_idx + self.MD_SIZE : buf_idx + size]
+        with (
+            memoryview(data_buff) as data_view,
+        ):
+            yield data_view, (id, size)
+
+    def free_buf(
+        self,
+        is_free_fn: Callable[[int, memoryview], bool],
+        nbytes: int | None = None,
+    ) -> Iterable[int]:
+        """
+        Free a buffer of the given size. This is a no-op in shared memory,
+        but we need to keep track of the metadata.
+
+        If freed memory spreads across the end and start of the ring buffer,
+        the actual freed memory will be in two segments. In this case there
+        still might not be a contiguous space of `nbytes` available.
+
+        Args:
+            nbytes (int, optional): The size of the buffer to free. If None,
+                frees the maximum size of the ring buffer.
+        """
+
+        assert self.is_writer, "Only the writer can free buffers."
+        logger.debug(
+            "Freeing up space in the ring buffer, "
+            "monotonic_id_start: %d, monotonic_id_end: %d",
+            self.monotonic_id_start,
+            self.monotonic_id_end,
+        )
+        monotonic_id_before = self.monotonic_id_start
+        # if nbytes is None, free up the maximum size of the ring buffer
+        if nbytes is None:
+            nbytes = self.data_buffer_size
+        freed_bytes = 0
+        while self.monotonic_id_start in self.metadata and freed_bytes < nbytes:
+            address = self.metadata[self.monotonic_id_start]
+            with self.access_buf(address) as (data_buff, metadata):
+                if is_free_fn(self.monotonic_id_start, data_buff):
+                    # check passed, we can free the buffer
+                    del self.metadata[self.monotonic_id_start]
+                    self.monotonic_id_start = (
+                        self.monotonic_id_start + 1
+                    ) % self.ID_MAX
+                    if self.monotonic_id_start in self.metadata:
+                        # pointing to the start addr of next allocation
+                        self.data_buffer_start += (
+                            self.metadata[self.monotonic_id_start]
+                            - self.data_buffer_start
+                        ) % self.data_buffer_size
+                    else:
+                        # no remaining allocation, reset to zero
+                        self.data_buffer_start = self.data_buffer_end = 0
+                    freed_bytes += metadata[1]
+                else:
+                    # there are still readers, we cannot free the buffer
+                    break
+
+        logger.debug(
+            "Freed %d bytes from the ring buffer, "
+            "monotonic_id_start: %d, monotonic_id_end: %d",
+            freed_bytes,
+            self.monotonic_id_start,
+            self.monotonic_id_end,
+        )
+
+        # buffer wrap around
+        if self.data_buffer_start >= self.data_buffer_size:
+            self.data_buffer_start -= self.data_buffer_size
+            self.data_buffer_end -= self.data_buffer_size
+
+        monotonic_id_after = self.monotonic_id_start
+        # id wrap around
+        if monotonic_id_after >= monotonic_id_before:
+            return range(monotonic_id_before, monotonic_id_after)
+        else:
+            return chain(
+                range(monotonic_id_before, self.ID_MAX), range(0, monotonic_id_after)
+            )
+
+
+class ObjectSerde(ABC):
+    @abstractmethod
+    def serialize(self, value: Any) -> tuple[Any, int, bytes, int]:
+        """Serialize an object to bytes."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def deserialize(self, data: memoryview) -> Any:
+        """Deserialize bytes back to an object."""
+        raise NotImplementedError
+
+
+class MsgpackSerde(ObjectSerde):
+    def __init__(self):
+        # Delayed import to avoid circular dependency
+        from vllm.multimodal.inputs import MultiModalKwargsItem
+        from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
+
+        self.encoder = MsgpackEncoder()
+        self.tensor_decoder = MsgpackDecoder(torch.Tensor, share_mem=False)
+        self.mm_decoder = MsgpackDecoder(MultiModalKwargsItem, share_mem=False)
+        self._mm_kwargs_item_cls = MultiModalKwargsItem
+
+    def serialize(self, value: Any) -> tuple[bytes | list[bytes], int, bytes, int]:
+        len_arr = None
+        if isinstance(value, (torch.Tensor, self._mm_kwargs_item_cls)):
+            type_name = type(value).__name__
+            value = self.encoder.encode(value)
+            len_arr = [len(s) for s in value]
+            nbytes = sum(len_arr)
+        else:
+            value = pickle.dumps(value, protocol=pickle.HIGHEST_PROTOCOL)
+            type_name = type(value).__name__
+            nbytes = len(value)
+
+        object_metadata = (type_name, nbytes, len_arr)
+        serialized_metadata = pickle.dumps(
+            object_metadata, protocol=pickle.HIGHEST_PROTOCOL
+        )
+        return value, nbytes, serialized_metadata, len(serialized_metadata)
+
+    def deserialize(self, data_view: memoryview) -> Any:
+        # pickle.loads do not read past the end of a pickled object
+        # within a large buffer, so we can skip storing the metadata size
+        type_name, nbytes, len_arr = pickle.loads(data_view)
+        serialized_data = data_view[-nbytes:]
+
+        if type_name == torch.Tensor.__name__:
+            obj = []
+            start_idx = 0
+            for length in len_arr:
+                item_bytes = serialized_data[start_idx : start_idx + length]
+                obj.append(item_bytes)
+                start_idx += length
+            obj = self.tensor_decoder.decode(obj)
+        elif type_name == self._mm_kwargs_item_cls.__name__:
+            obj = []
+            start_idx = 0
+            for length in len_arr:
+                item_bytes = serialized_data[start_idx : start_idx + length]
+                obj.append(item_bytes)
+                start_idx += length
+            obj = self.mm_decoder.decode(obj)
+        elif type_name == bytes.__name__:
+            obj = pickle.loads(serialized_data)
+        else:
+            raise ValueError(f"Unsupported object type '{type_name}' in metadata")
+
+        return obj
+
+
+@dataclass
+class ShmObjectStorageHandle:
+    max_object_size: int
+    n_readers: int
+    ring_buffer_handle: tuple[int, str]
+    serde_class: type[ObjectSerde]
+    reader_lock: LockType | None
+
+
+class SingleWriterShmObjectStorage:
+    """
+    A single-writer, multiple-reader object storage system built on top of a
+    shared memory ring buffer. Provides key-value storage with automatic memory
+    management and cross-process serialization support.
+
+    This storage system follows a FIFO (First-In-First-Out) eviction policy
+    where the oldest objects are automatically freed when memory runs low.
+    Memory is reclaimed based on reader reference counting - objects are only
+    freed when all readers have finished accessing them.
+
+    Architecture:
+    - Single writer process can put(key, value) objects
+    - Multiple reader processes can get(address, monotonic_id) objects
+    - Built on SingleWriterShmRingBuffer for efficient shared memory management
+    - Thread-safe operations with reader synchronization via locks
+
+    Key Features:
+    - FIFO Eviction: Oldest objects are evicted first when memory is full
+    - Reference Counting: Objects are only freed when no readers are
+      accessing them
+    - Duplicate Key Handling: Existing keys are not overwritten, just
+      re-referenced
+    - Customized Serialization: By default uses Msgpack for efficient
+      serialization of Python objects, but can be extended for custom types
+    - Cross-Process Safety: Uses shared memory with proper synchronization
+    - Automatic Cleanup: Garbage collection happens transparently during
+      allocation
+
+    Memory Layout per Object:
+    `[4-byte reference_count][metadata_size][serialized_object_data]`
+
+    Thread Safety:
+    - Writer operations (put, clear) are single-threaded by design
+    - Reader operations (get) are thread-safe with lock-based reference
+      counting
+    - Memory reclamation is handled exclusively by the writer process
+    """
+
+    def __init__(
+        self,
+        max_object_size: int,
+        n_readers: int,
+        ring_buffer: SingleWriterShmRingBuffer,
+        serde_class: type[ObjectSerde] = MsgpackSerde,
+        reader_lock: LockType | None = None,
+    ):
+        """
+        Initialize the object storage.
+
+        Args:
+            max_object_size: Maximum size for a single object in bytes.
+            n_readers: Number of reader processes that can access the storage.
+            ring_buffer: The shared memory ring buffer for storing objects.
+            serde_class: Serializer/deserializer for objects.
+            reader_lock: Optional lock for synchronizing reader access.
+        Raises:
+            ValueError: If reader_lock is None for readers.
+        """
+
+        self.max_object_size = max_object_size
+        self.n_readers = n_readers
+        self.serde_class = serde_class
+        self.ser_de = serde_class()
+        self.ring_buffer = ring_buffer
+        self.is_writer = self.ring_buffer.is_writer
+
+        self.flag_bytes = 4  # for in-use flag
+
+        if self.is_writer:
+            # Key-value mapping: key -> (address, monotonic_id)
+            self.key_index: dict[str, tuple[int, int]] = {}
+            # Reverse mapping: monotonic_id -> key
+            self.id_index: dict[int, str] = {}
+            # Writer flag to track in-use status: monotonic_id -> count
+            self.writer_flag: dict[int, int] = {}
+        else:
+            if reader_lock is None:
+                raise ValueError("Lock must be provided for readers.")
+
+        self._reader_lock = reader_lock
+
+    def clear(self) -> None:
+        """Clear the object storage."""
+        if self.is_writer:
+            self.ring_buffer.clear()
+            self.key_index.clear()
+            self.id_index.clear()
+            self.writer_flag.clear()
+            logger.debug("Object storage cleared and reinitialized.")
+
+    def copy_to_buffer(
+        self,
+        data: bytes | list[bytes],
+        data_bytes: int,
+        metadata: bytes,
+        md_bytes: int,
+        data_view: memoryview,
+    ) -> None:
+        data_view[self.flag_bytes : self.flag_bytes + md_bytes] = metadata
+        if isinstance(data, bytes):
+            data_view[-data_bytes:] = data
+        elif isinstance(data, list):
+            start_idx = self.flag_bytes + md_bytes
+            for item_bytes in data:
+                item_size = len(item_bytes)
+                data_view[start_idx : start_idx + item_size] = item_bytes
+                start_idx += item_size
+        else:
+            raise ValueError(f"Unsupported data type for serialization: {type(data)}")
+
+    def increment_writer_flag(self, id: int) -> None:
+        """Set the in-use flag for the writer."""
+        self.writer_flag[id] = self.writer_flag.get(id, 0) + 1
+
+    def increment_reader_flag(self, data_view: memoryview) -> None:
+        """Set the in-use flag for the reader."""
+        # >0 for in-use flag
+        reader_count = self.ring_buffer.byte2int(data_view)
+        data_view[:] = self.ring_buffer.int2byte(reader_count + 1)
+
+    def free_unused(self) -> None:
+        """Free unused buffers in the ring buffer."""
+        # try to free up 2*max_object_size bytes of space in the ring buffer,
+        # since the buffer might be fragmented
+        freed_ids = self.ring_buffer.free_buf(
+            self.default_is_free_check, 2 * self.max_object_size
+        )
+        # update the metadata after freeing up space
+        for freed_id in freed_ids:
+            key_to_free = self.id_index[freed_id]
+            del self.key_index[key_to_free]
+            del self.id_index[freed_id]
+            del self.writer_flag[freed_id]
+
+    def is_cached(self, key: str) -> bool:
+        """
+        Check if the object with the given key is cached.
+        """
+        return key in self.key_index
+
+    def get_cached(self, key: str) -> tuple[int, int]:
+        """
+        Get the cached object by key if it exists.
+        """
+        address, monotonic_id = self.key_index[key]
+        self.increment_writer_flag(monotonic_id)
+        return address, monotonic_id
+
+    def put(self, key: str, value: Any) -> tuple[int, int]:
+        """
+        Store a key-value pair in the object storage.
+        Attempts to free max_object_size bytes using FIFO order
+        when the ring buffer runs out of space during a put() operation.
+
+        Args:
+            key: String key to identify the object
+            value: Any serializable Python object
+
+        Raises:
+            MemoryError: If there's not enough space in the buffer
+            ValueError: If the serialized object is too large
+            ValueError: If the key already exists in the storage
+        """
+        if key in self.key_index:
+            raise ValueError(f"Key '{key}' already exists in the storage.")
+
+        object_data, data_bytes, object_metadata, md_bytes = self.ser_de.serialize(
+            value
+        )
+        buffer_size = self.flag_bytes + data_bytes + md_bytes
+        # Sanity checks
+        if buffer_size > self.max_object_size:
+            raise ValueError(
+                f"Serialized object size ({buffer_size} bytes) exceeds "
+                f"max object size ({self.max_object_size} bytes)"
+            )
+
+        # Allocate new buffer
+        try:
+            address, monotonic_id = self.ring_buffer.allocate_buf(buffer_size)
+        except MemoryError:
+            self.free_unused()
+            # try again after freeing up space
+            address, monotonic_id = self.ring_buffer.allocate_buf(buffer_size)
+
+        # Write data to buffer
+        with self.ring_buffer.access_buf(address) as (data_view, metadata):
+            data_view[: self.flag_bytes] = self.ring_buffer.int2byte(0)
+            self.copy_to_buffer(
+                object_data, data_bytes, object_metadata, md_bytes, data_view
+            )
+        self.increment_writer_flag(monotonic_id)
+
+        # Update key index
+        self.key_index[key] = (address, monotonic_id)
+        self.id_index[monotonic_id] = key
+        return address, monotonic_id
+
+    def get(self, address: int, monotonic_id: int) -> Any:
+        # Read data from buffer
+        with self.ring_buffer.access_buf(address) as (data_view, buf_metadata):
+            # check id from metadata
+            if buf_metadata[0] != monotonic_id:
+                raise ValueError(
+                    f"Data for address:id '{address}:{monotonic_id}'"
+                    " has been modified or is invalid."
+                )
+
+            obj = self.ser_de.deserialize(data_view[self.flag_bytes :])
+
+            # decrease the in-use flag for reader reads
+            if self._reader_lock is not None:
+                with self._reader_lock:
+                    self.increment_reader_flag(data_view[: self.flag_bytes])
+            else:
+                # if self._reader_lock is None, it means we are the writer
+                # in this case, we do not need to decrease the reader count
+                assert self.is_writer
+
+        return obj
+
+    def touch(
+        self,
+        key: str,
+        address: int = 0,
+        monotonic_id: int = 0,
+    ) -> None:
+        """
+        Touch an existing cached item to update its eviction status.
+
+        For writers (ShmObjectStoreSenderCache): Increment writer_flag
+        For readers (ShmObjectStoreReceiverCache): Increment reader_count
+
+        Args:
+            key: String key of the object to touch
+            address: Address of the object (only for readers)
+            monotonic_id: Monotonic ID of the object (only for readers)
+
+        """
+        if self._reader_lock is None:
+            if key not in self.key_index:
+                return None
+            address, monotonic_id = self.key_index[key]
+            # Writer side: increment writer_flag to raise eviction threshold
+            self.increment_writer_flag(monotonic_id)
+        else:
+            with (
+                self._reader_lock,
+                self.ring_buffer.access_buf(address) as (data_view, _),
+            ):
+                reader_count = self.ring_buffer.byte2int(data_view[: self.flag_bytes])
+
+                # NOTE(Long):
+                # Avoid increasing flag on newly added item (sync with sender)
+                # Since when a new item is added
+                # pre-touch has no effect on writer side
+                if reader_count >= self.n_readers:
+                    self.increment_reader_flag(data_view[: self.flag_bytes])
+
+    def close(self) -> None:
+        """Close the shared memory."""
+        self.ring_buffer.close()
+
+    def handle(self):
+        """Get handle for sharing across processes."""
+        return ShmObjectStorageHandle(
+            max_object_size=self.max_object_size,
+            n_readers=self.n_readers,
+            ring_buffer_handle=self.ring_buffer.handle(),
+            serde_class=self.serde_class,
+            reader_lock=self._reader_lock,
+        )
+
+    @staticmethod
+    def create_from_handle(
+        handle: ShmObjectStorageHandle,
+    ) -> "SingleWriterShmObjectStorage":
+        logger.debug("Creating storage from handle: %s", handle)
+        ring_buffer = SingleWriterShmRingBuffer(*handle.ring_buffer_handle)
+        return SingleWriterShmObjectStorage(
+            max_object_size=handle.max_object_size,
+            n_readers=handle.n_readers,
+            ring_buffer=ring_buffer,
+            serde_class=handle.serde_class,
+            reader_lock=handle.reader_lock,
+        )
+
+    def default_is_free_check(self, id: int, buf: memoryview) -> bool:
+        """
+        Default is_free function that checks if the first 4 bytes are zero.
+        This indicates that the buffer is free.
+        """
+        reader_count = int.from_bytes(buf[0:4], "little", signed=True)
+        writer_count = self.writer_flag[id]
+        return reader_count >= writer_count * self.n_readers

vllm/distributed/device_communicators/symm_mem.py

@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+
+import vllm.envs as envs
+from vllm.distributed.device_communicators.all_reduce_utils import (
+    SYMM_MEM_ALL_REDUCE_MAX_SIZES,
+)
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+try:
+    import torch.distributed._symmetric_memory as torch_symm_mem
+
+    symm_mem_available = True
+except ImportError:
+    symm_mem_available = False
+
+logger = init_logger(__name__)
+
+
+class SymmMemCommunicator:
+    _WORLD_SIZES_MULTIMEM = {
+        "9.0": [4, 6, 8],
+        "10.0": [6, 8],
+        "10.3": [6, 8],
+    }
+
+    def __init__(
+        self,
+        group: ProcessGroup,
+        device: int | str | torch.device,
+        # add options for testing
+        force_multimem: bool | None = None,
+        max_size_override: int | None = None,
+    ):
+        self.disabled = True
+
+        if not symm_mem_available:
+            return
+
+        if not current_platform.is_cuda():
+            logger.warning("SymmMemCommunicator: symmetric memory is not available.")
+            return
+        if isinstance(device, int):
+            device = torch.device(f"cuda:{device}")
+        elif isinstance(device, str):
+            device = torch.device(device)
+        torch.accelerator.set_device_index(device)
+        self.dtype = torch.bfloat16
+        self.device = device
+        self.group = group
+        self.world_size = dist.get_world_size(self.group)
+        capability = current_platform.get_device_capability()
+        if capability is None:
+            logger.warning(
+                "SymmMemCommunicator: device capability is unknown, "
+                "communicator is not available."
+            )
+            return
+        self.device_capability = capability.as_version_str()
+        if self.device_capability not in SYMM_MEM_ALL_REDUCE_MAX_SIZES:
+            logger.warning(
+                "SymmMemCommunicator: Device capability %s not supported, "
+                "communicator is not available.",
+                self.device_capability,
+            )
+            return
+        if self.world_size not in SYMM_MEM_ALL_REDUCE_MAX_SIZES[self.device_capability]:
+            logger.warning(
+                "SymmMemCommunicator: World size %d not supported, "
+                "communicator is not available.",
+                self.world_size,
+            )
+            return
+        # Use override max_size if provided, otherwise use default
+        if max_size_override is not None:
+            self.max_size = max_size_override
+            logger.info(
+                "SymmMemCommunicator: Using override max_size: %s bytes",
+                self.max_size,
+            )
+        else:
+            self.max_size = SYMM_MEM_ALL_REDUCE_MAX_SIZES[self.device_capability][
+                self.world_size
+            ]
+        try:
+            self.buffer = torch_symm_mem.empty(
+                self.max_size // self.dtype.itemsize,
+                device=self.device,
+                dtype=self.dtype,
+            )
+            handle = torch_symm_mem.rendezvous(self.buffer, self.group.group_name)
+        except RuntimeError as e:
+            logger.warning_once(
+                "SymmMemCommunicator: symmetric memory initialization failed: %s "
+                "Communicator is not available. To suppress this warning set "
+                "VLLM_ALLREDUCE_USE_SYMM_MEM=0",
+                str(e),
+            )
+            return
+        if handle.multicast_ptr == 0:
+            logger.warning(
+                "SymmMemCommunicator: symmetric memory "
+                "multicast operations are not supported."
+            )
+            return
+        self.force_multimem = force_multimem
+        self.disabled = False
+        if envs.VLLM_BATCH_INVARIANT:
+            self.disabled = True
+
+    def should_use_symm_mem(self, inp: torch.Tensor):
+        if self.disabled:
+            return False
+        if inp.dtype != self.dtype:
+            return False
+        inp_size = inp.numel() * inp.element_size()
+        if inp_size % 4 != 0:
+            return False
+        return inp_size <= self.max_size
+
+    def all_reduce(
+        self, inp: torch.Tensor, *, out: torch.Tensor | None = None
+    ) -> torch.Tensor | None:
+        if not self.should_use_symm_mem(inp):
+            return None
+        if out is None:
+            out = torch.empty_like(inp)
+        self.buffer[: inp.numel()].copy_(inp.view(-1))
+
+        # Determine which algorithm to use
+        use_multimem = False
+        if self.force_multimem is not None:
+            # Test override: use forced setting
+            use_multimem = self.force_multimem
+        else:
+            # Normal logic: use multimem for supported world sizes
+            use_multimem = (
+                self.world_size in self._WORLD_SIZES_MULTIMEM[self.device_capability]
+            )
+
+        if use_multimem:
+            torch.ops.symm_mem.multimem_all_reduce_(
+                self.buffer[: inp.numel()], "sum", self.group.group_name
+            )
+        else:
+            torch.ops.symm_mem.two_shot_all_reduce_(
+                self.buffer[: inp.numel()], "sum", self.group.group_name
+            )
+        out.copy_(self.buffer[: inp.numel()].view(out.shape))
+        return out

vllm/distributed/device_communicators/xpu_communicator.py

@@ -0,0 +1,253 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+
+from vllm.logger import init_logger
+
+from .base_device_communicator import DeviceCommunicatorBase
+
+logger = init_logger(__name__)
+
+
+class XpuCommunicator(DeviceCommunicatorBase):
+    def __init__(
+        self,
+        cpu_group: ProcessGroup,
+        device: torch.device | None = None,
+        device_group: ProcessGroup | None = None,
+        unique_name: str = "",
+    ):
+        super().__init__(cpu_group, device, device_group, unique_name)
+        self.ca_comm: None = None
+        if self.use_all2all:
+            if self.all2all_backend in ("naive", "allgather_reducescatter"):
+                from .all2all import AgRsAll2AllManager
+
+                self.all2all_manager = AgRsAll2AllManager(self.cpu_group)
+                logger.info("Using AgRs manager on XPU device.")
+
+            else:  # type: ignore[has-type]
+                logger.warning(
+                    "`%s` all2all manager is not supported on XPU. "
+                    "Falling back to AgRs manager for XPU, "
+                    "which is the Default backend",
+                    self.all2all_backend,  # type: ignore[has-type]
+                )
+                from .all2all import AgRsAll2AllManager
+
+                self.all2all_manager = AgRsAll2AllManager(self.cpu_group)
+                logger.info("Using AgRs manager on XPU device.")
+
+    def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
+        output = input_.clone()
+        dist.all_reduce(output, group=self.device_group)
+        return output
+
+    def reduce_scatter(self, input_: torch.Tensor, dim: int = -1):
+        world_size = self.world_size
+
+        if dim < 0:
+            # Convert negative dim to positive.
+            dim += input_.dim()
+
+        # Note: This will produce an incorrect answer if we don't make
+        # the input_tensor contiguous. Possible bug in reduce_scatter_tensor?
+        input_tensor = input_.movedim(0, dim).contiguous()
+
+        assert input_tensor.shape[0] % world_size == 0
+        chunk_size = input_tensor.shape[0] // world_size
+        output_shape = (chunk_size,) + input_tensor.shape[1:]
+
+        output = torch.empty(
+            output_shape, dtype=input_tensor.dtype, device=input_tensor.device
+        )
+
+        dist.reduce_scatter_tensor(output, input_tensor, group=self.device_group)
+
+        # Reshape before returning
+        return output.movedim(0, dim).contiguous()
+
+    def reduce_scatterv(
+        self, input_: torch.Tensor, dim: int = -1, sizes: list[int] | None = None
+    ):
+        world_size = self.world_size
+
+        if dim < 0:
+            # Convert negative dim to positive.
+            dim += input_.dim()
+
+        # Note: This will produce an incorrect answer if we don't make
+        # the input_tensor contiguous. Possible bug in reduce_scatter_tensor?
+        input_tensor = input_.movedim(0, dim).contiguous()
+
+        if sizes is not None:
+            assert len(sizes) == world_size
+            assert input_tensor.shape[0] == sum(sizes)
+            chunk_size = sizes[self.rank_in_group]
+        else:
+            assert input_tensor.shape[0] % world_size == 0
+            chunk_size = input_tensor.shape[0] // world_size
+        output_shape = (chunk_size,) + input_tensor.shape[1:]
+
+        output = torch.empty(
+            output_shape, dtype=input_tensor.dtype, device=input_tensor.device
+        )
+        if sizes is not None and sizes.count(sizes[0]) != len(sizes):
+            # if inputs shape in different ranks is not the same using reduce_scatter
+            input_splits = list(input_tensor.split(sizes, dim=0))
+            dist.reduce_scatter(output, input_splits, group=self.device_group)
+        else:
+            dist.reduce_scatter_tensor(output, input_tensor, group=self.device_group)
+        # Reshape before returning
+        return output.movedim(0, dim).contiguous()
+
+    def all_gatherv(
+        self,
+        input_: torch.Tensor | list[torch.Tensor],
+        dim: int = 0,
+        sizes: list[int] | None = None,
+    ):
+        if dim != 0:
+            raise NotImplementedError("only dim 0 all-gatherv is supported")
+        world_size = self.world_size
+
+        # 'sizes' is not needed if all inputs in the same group have the same
+        # shape
+        if sizes is not None and all(s == sizes[0] for s in sizes):
+            sizes = None
+
+        def _all_gather_single(input_: torch.Tensor, sizes: list[int] | None = None):
+            input_size = input_.size()
+            if sizes is not None:
+                assert len(sizes) == world_size
+                assert input_.shape[dim] == sizes[self.rank_in_group], (
+                    f"{input_.shape[dim]} != {sizes[self.rank_in_group]}"
+                )
+                output_size = (sum(sizes),) + input_size[1:]
+            else:
+                output_size = (input_size[0] * world_size,) + input_size[1:]
+            # Allocate output tensor.
+            output_tensor = torch.empty(
+                output_size, dtype=input_.dtype, device=input_.device
+            )
+
+            if sizes is not None:
+                all_gather_list = []
+                for size in sizes:
+                    all_gather_list.append(
+                        torch.empty(
+                            (size,) + input_.shape[1:],
+                            dtype=input_.dtype,
+                            device=input_.device,
+                        )
+                    )
+                dist.all_gather(all_gather_list, input_, group=self.device_group)
+                output_tensor = torch.cat(all_gather_list, dim=0)
+            else:
+                dist.all_gather([output_tensor], input_, group=self.device_group)
+            return output_tensor
+
+        if isinstance(input_, torch.Tensor):
+            return _all_gather_single(input_, sizes)
+
+        output_list = []
+        for inp in input_:
+            output_list.append(_all_gather_single(inp, sizes=sizes))
+        return output_list
+
+    def gather(
+        self, input_: torch.Tensor, dst: int = 0, dim: int = -1
+    ) -> torch.Tensor | None:
+        assert -input_.dim() <= dim < input_.dim(), (
+            f"Invalid dim ({dim}) for input tensor with shape {input_.size()}"
+        )
+        if dim < 0:
+            # Convert negative dim to positive.
+            dim += input_.dim()
+        # For xpu path, gather doesn't work properly together with ray
+        # cluster so we use all_gather instead for now.
+        input_size = input_.size()
+        # Allocate output tensor.
+        output_tensor = torch.empty(
+            (self.world_size,) + input_size, dtype=input_.dtype, device=input_.device
+        )
+        # All-gather.
+        dist.all_gather_into_tensor(output_tensor, input_, group=self.device_group)
+        if self.rank_in_group == dst:
+            # Reshape
+            output_tensor = output_tensor.movedim(0, dim)
+            output_tensor = output_tensor.reshape(
+                input_size[:dim]
+                + (self.world_size * input_size[dim],)
+                + input_size[dim + 1 :]
+            )
+        else:
+            output_tensor = None
+        return output_tensor
+
+    def broadcast(self, input_: torch.Tensor, src: int = 0) -> None:
+        dist.broadcast(input_, src=src, group=self.device_group)
+
+    def dispatch_router_logits(
+        self,
+        hidden_states: torch.Tensor,
+        router_logits: torch.Tensor,
+        is_sequence_parallel: bool = False,
+        extra_tensors: list[torch.Tensor] | None = None,
+    ) -> (
+        tuple[torch.Tensor, torch.Tensor]
+        | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
+    ):
+        """
+        Dispatch the hidden states and router logits to the appropriate device.
+        This is a no-op in the base class.
+        """
+
+        assert self.all2all_manager is not None
+        return self.all2all_manager.dispatch_router_logits(
+            hidden_states,
+            router_logits,
+            is_sequence_parallel,
+            extra_tensors,
+        )
+
+    def dispatch(
+        self,
+        hidden_states: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        is_sequence_parallel: bool = False,
+        extra_tensors: list[torch.Tensor] | None = None,
+    ) -> (
+        tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+        | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
+    ):
+        """
+        Dispatch the hidden states and topk weights/ids to the appropriate device.
+        This is a no-op in the base class.
+        """
+        assert self.all2all_manager is not None
+        return self.all2all_manager.dispatch(
+            hidden_states,
+            topk_weights,
+            topk_ids,
+            is_sequence_parallel,
+            extra_tensors=extra_tensors,
+        )
+
+    def combine(
+        self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
+    ) -> torch.Tensor:
+        """
+        Combine the hidden states and router logits from the appropriate device.
+        This is a no-op in the base class.
+        """
+        assert self.all2all_manager is not None
+        return self.all2all_manager.combine(
+            hidden_states,
+            is_sequence_parallel,
+        )

vllm/distributed/ec_transfer/__init__.py

@@ -0,0 +1,16 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.distributed.ec_transfer.ec_transfer_state import (
+    ensure_ec_transfer_initialized,
+    ensure_ec_transfer_shutdown,
+    get_ec_transfer,
+    has_ec_transfer,
+)
+
+__all__ = [
+    "get_ec_transfer",
+    "ensure_ec_transfer_initialized",
+    "ensure_ec_transfer_shutdown",
+    "has_ec_transfer",
+]

vllm/distributed/ec_transfer/ec_connector/base.py

@@ -0,0 +1,277 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+ECConnectorBase Class for Distributed Encoder Cache &
+P2P Encoder cache communication in V1
+
+The class provides the following primitives:
+    Scheduler-side: runs in the scheduler, binds metadata, which
+    is used by the worker-side to load/save Encoder cache.
+        check_caches_exist() - Check whether Encoder cache of requests exist
+        update_state_after_alloc() - update ECConnector state after
+        allocate. This will decide to load the cache or not
+        request_finished() - called when a request is finished,
+        free the cache with the requests
+
+    Worker-side: runs in each worker, loads/saves Encoder Cache to/from
+    the Connector based on the metadata.
+        start_load_ec() - starts loading all ECs (maybe async)
+        wait_for_save() - blocks until all saves are done
+
+        get_finished() - called with ids of finished requests, returns
+            ids of requests that have completed async sending/recving.
+"""
+
+import enum
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from vllm.logger import init_logger
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.outputs import ECConnectorOutput
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.v1.request import Request
+
+logger = init_logger(__name__)
+
+
+class ECConnectorRole(enum.Enum):
+    # Connector running in the scheduler process
+    SCHEDULER = 0
+
+    # Connector running in the worker process
+    WORKER = 1
+
+
+class ECConnectorMetadata(ABC):  # noqa: B024
+    """
+    Abstract Metadata used to communicate between the
+    Scheduler ECConnector and Worker ECConnector.
+    """
+
+    pass
+
+
+class ECConnectorBase(ABC):
+    def __init__(self, vllm_config: "VllmConfig", role: ECConnectorRole):
+        self._connector_metadata: ECConnectorMetadata | None = None
+        self._vllm_config = vllm_config
+        self._role = role
+        if vllm_config.ec_transfer_config is not None:
+            self._is_producer = vllm_config.ec_transfer_config.is_ec_producer
+            self._is_consumer = vllm_config.ec_transfer_config.is_ec_consumer
+        else:
+            raise ValueError("ec_transfer_config must be set for ECConnectorBase")
+
+    @property
+    def role(self) -> ECConnectorRole:
+        return self._role
+
+    @property
+    def is_producer(self) -> bool:
+        return self._is_producer
+
+    @property
+    def is_consumer(self) -> bool:
+        return self._is_consumer
+
+    def shutdown(self) -> None:
+        """
+        Shutdown the connector. This is called when the process
+        is shutting down to ensure that all the async operations are
+        completed and the connector is cleaned up properly.
+        """
+        return None
+
+    # ==============================
+    # Worker-side methods
+    # ==============================
+
+    def bind_connector_metadata(self, connector_metadata: ECConnectorMetadata) -> None:
+        """Set the connector metadata from the scheduler.
+
+        This function should be called by the model runner every time
+        before the model execution. The metadata will be used for runtime
+        EC cache loading.
+
+        Args:
+            connector_metadata (dict): the connector metadata.
+        """
+        self._connector_metadata = connector_metadata
+
+    def clear_connector_metadata(self) -> None:
+        """Clear the connector metadata.
+
+        This function should be called by the model runner every time
+        after the model execution.
+        """
+        self._connector_metadata = None
+
+    def _get_connector_metadata(self) -> ECConnectorMetadata:
+        """Get the connector metadata.
+
+        This function should only be called inside the connector.
+
+        Returns:
+            ConnectorMetadata: the connector metadata.
+        """
+
+        # Should only be called while set to valid metadata.
+        assert self._connector_metadata is not None
+        return self._connector_metadata
+
+    def register_caches(
+        self,
+        ec_caches: dict[str, torch.Tensor],
+    ):
+        """
+        Initialize with the EC caches.
+        Args:
+            ec_caches: dictionary of encoder cache
+        """
+        # TODO: Implement this later for P2P feature
+        return
+
+    @abstractmethod
+    def start_load_caches(
+        self, encoder_cache: dict[str, torch.Tensor], **kwargs
+    ) -> None:
+        """
+        Start loading the cache from the connector into vLLM's encoder cache.
+
+        This method loads the encoder cache based on metadata provided by the scheduler.
+        It is called before `_gather_mm_embeddings` for the EC Connector. For EC,
+        the `encoder_cache` and `mm_hash` are stored in `kwargs`.
+
+        Args:
+            encoder_cache (dict[str, torch.Tensor]): A dictionary mapping multimodal
+                data hashes (`mm_hash`) to encoder cache tensors.
+            kwargs (dict): Additional keyword arguments for the connector.
+        """
+        pass
+
+    @abstractmethod
+    def save_caches(
+        self, encoder_cache: dict[str, torch.Tensor], mm_hash: str, **kwargs
+    ) -> None:
+        """
+        Save the encoder cache to the connector.
+
+        This method saves the encoder cache from the worker's local storage
+        to shared storage or another external connector.
+
+        Args:
+            encoder_cache (dict[str, torch.Tensor]): A dictionary mapping multimodal
+                data hashes (`mm_hash`) to encoder cache tensors.
+            mm_hash (str): The hash of the multimodal data whose cache is being saved.
+            kwargs (dict): Additional keyword arguments for the connector.
+        """
+        pass
+
+    def get_finished(
+        self, finished_req_ids: set[str]
+    ) -> tuple[set[str] | None, set[str] | None]:
+        """
+        Notifies worker-side connector ids of requests that have
+        finished generating tokens on the worker.
+        The scheduler process (via the Executors) will use this output
+        to track which workers are done.
+
+        Returns:
+            ids of requests that have finished asynchronous transfer
+            (requests that previously returned True from request_finished()),
+            tuple of (sending/saving ids, recving/loading ids).
+            The finished saves/sends req ids must belong to a set provided in a
+            call to this method (this call or a prior one).
+        """
+        return None, None
+
+    # ==============================
+    # Scheduler-side methods
+    # ==============================
+
+    @abstractmethod
+    def has_cache_item(
+        self,
+        identifier: str,
+    ) -> bool:
+        """
+        Check if a single encoder cache exists
+
+        Args:
+            identifier (str): the identifier of the media.
+
+        Returns:
+            A bool where value is True if cache exist for
+            the media
+        """
+        pass
+
+    def ensure_cache_available(
+        self, request: "Request", num_computed_tokens: int
+    ) -> bool:
+        """
+        Ensure encoder cache items are available for the given request.
+        May initiate asynchronous transfers for items not yet local.
+
+        Args:
+            request: the request whose multimodal features to check.
+            num_computed_tokens: tokens already covered by cached KV blocks.
+
+        Returns:
+            True if all items are ready or no transfer is needed.
+            False if any items are still in transit (request should be deferred).
+        """
+        return True
+
+    @abstractmethod
+    def update_state_after_alloc(self, request: "Request", index: int):
+        """
+        Update ECConnector state to decide allocate cache for requests
+
+        Args:
+            request (Request): the request object.
+        """
+        pass
+
+    @abstractmethod
+    def build_connector_meta(
+        self, scheduler_output: SchedulerOutput
+    ) -> ECConnectorMetadata:
+        """
+        Build the connector metadata for this step.
+
+        This function should NOT modify fields in the scheduler_output.
+        Also, calling this function will reset the state of the connector.
+
+        Args:
+            scheduler_output (SchedulerOutput): the scheduler output object.
+        """
+        pass
+
+    def update_connector_output(self, connector_output: ECConnectorOutput):
+        """
+        Update ECConnector state from worker-side connectors output.
+
+        Args:
+            connector_output (ECConnectorOutput): the worker-side
+                connectors output.
+        """
+        return
+
+    def request_finished(
+        self, request: "Request"
+    ) -> tuple[bool, dict[str, Any] | None]:
+        """
+        Called when a request has finished, before its encoder cache is freed.
+
+        Returns:
+            True if the request is being saved/sent asynchronously and cached
+            should not be freed until the request_id is returned from
+            get_finished().
+        """
+        return False, None

vllm/distributed/ec_transfer/ec_connector/example_connector.py

@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import os
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import safetensors
+
+from vllm.config import VllmConfig
+from vllm.distributed.ec_transfer.ec_connector.base import (
+    ECConnectorBase,
+    ECConnectorMetadata,
+    ECConnectorRole,
+)
+from vllm.logger import init_logger
+from vllm.v1.core.sched.output import SchedulerOutput
+
+if TYPE_CHECKING:
+    from vllm.v1.request import Request
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class MMMeta:
+    mm_hash: str
+    num_token: int
+
+    @staticmethod
+    def make_meta(mm_hash, num_token) -> "MMMeta":
+        return MMMeta(mm_hash=mm_hash, num_token=num_token)
+
+
+@dataclass
+class ECExampleConnectorMetadata(ECConnectorMetadata):
+    mm_datas: list[MMMeta]
+
+    def __init__(self):
+        self.mm_datas = []
+
+    def add_mm_data(self, mm_data: MMMeta):
+        self.mm_datas.append(mm_data)
+
+
+class ECExampleConnector(ECConnectorBase):
+    # NOTE: This is Simple debug implementation of the EC connector.
+    # It save / load the EC cache to / from the disk.
+
+    def __init__(self, vllm_config: "VllmConfig", role: ECConnectorRole):
+        super().__init__(vllm_config=vllm_config, role=role)
+        # req_id -> index
+        self._mm_datas_need_loads: dict[str, int] = {}
+        transfer_config = vllm_config.ec_transfer_config
+        if transfer_config is not None:
+            self._storage_path = transfer_config.get_from_extra_config(
+                "shared_storage_path", "/tmp"
+            )
+            logger.debug(transfer_config)
+            logger.debug("Shared storage path is %s", self._storage_path)
+        else:
+            raise ValueError("ec_transfer_config must be set for ECConnectorBase")
+
+    def start_load_caches(self, encoder_cache, **kwargs) -> None:
+        """
+        Start loading the cache from the connector into vLLM's encoder cache.
+
+        This method loads the encoder cache based on metadata provided by the scheduler.
+        It is called before `_gather_mm_embeddings` for the EC Connector. For EC,
+        the `encoder_cache` and `mm_hash` are stored in `kwargs`.
+
+        Args:
+            encoder_cache (dict[str, torch.Tensor]): A dictionary mapping multimodal
+                data hashes (`mm_hash`) to encoder cache tensors.
+            kwargs (dict): Additional keyword arguments for the connector.
+        """
+        from vllm.platforms import current_platform
+
+        # Get the metadata
+        metadata: ECConnectorMetadata = self._get_connector_metadata()
+        assert isinstance(metadata, ECExampleConnectorMetadata)
+        assert encoder_cache is not None
+        if metadata is None:
+            logger.warning(
+                "In connector.start_load_caches, but the connector metadata is None"
+            )
+            return
+        # Load the EC for each mm data
+        for mm_data in metadata.mm_datas:
+            if mm_data.mm_hash in encoder_cache:
+                continue
+            filename = self._generate_filename_debug(mm_data.mm_hash)
+            ec_cache = safetensors.torch.load_file(
+                filename, device=current_platform.device_type
+            )["ec_cache"]
+            encoder_cache[mm_data.mm_hash] = ec_cache
+            logger.debug("Success load encoder cache for hash %s", mm_data.mm_hash)
+
+    def save_caches(self, encoder_cache, mm_hash, **kwargs) -> None:
+        """
+        Save the encoder cache to the connector.
+
+        This method saves the encoder cache from the worker's local storage
+        to shared storage or another external connector.
+
+        Args:
+            encoder_cache (dict[str, torch.Tensor]): A dictionary mapping multimodal
+                data hashes (`mm_hash`) to encoder cache tensors.
+            mm_hash (str): The hash of the multimodal data whose cache is being saved.
+            kwargs (dict): Additional keyword arguments for the connector.
+        """
+        # Return if it is PD Instance
+        if not self.is_producer:
+            return
+        filename = self._generate_filename_debug(mm_hash)
+        ec_cache = encoder_cache[mm_hash]
+        tensors = {"ec_cache": ec_cache.detach().cpu()}
+        safetensors.torch.save_file(tensors, filename)
+        logger.debug("Save cache successful for mm_hash %s", mm_hash)
+
+    def has_cache_item(
+        self,
+        identifier: str,
+    ) -> bool:
+        """
+        Check if cache exist externally for the media
+
+        Args:
+            identifier (str): the identifier of the media.
+
+        Returns:
+            Bool indicate that media exists in cache or not
+        """
+        return self._found_match_for_mm_data(identifier)
+
+    def update_state_after_alloc(
+        self,
+        request: "Request",
+        index: int,
+    ) -> None:
+        """
+        Update ECConnector state after encoder cache allocation.
+        """
+        mm_hash = request.mm_features[index].identifier
+        # Only load cache if it is consumer and cache exists
+        if not self.is_consumer or not self.has_cache_item(mm_hash):
+            return
+        num_encoder_token = request.get_num_encoder_embeds(index)
+        self._mm_datas_need_loads[mm_hash] = num_encoder_token
+
+    def build_connector_meta(
+        self,
+        scheduler_output: SchedulerOutput,
+    ) -> ECConnectorMetadata:
+        """Build the connector metadata for this step.
+
+        This function should NOT modify any fields in the scheduler_output.
+        Also, calling this function will reset the state of the connector.
+        This only build for load mm_data only
+        Args:
+            scheduler_output (SchedulerOutput): the scheduler output object.
+        """
+        meta = ECExampleConnectorMetadata()
+        for mm_hash, num_encoder_token in self._mm_datas_need_loads.items():
+            meta.add_mm_data(MMMeta.make_meta(mm_hash, num_encoder_token))
+        self._mm_datas_need_loads.clear()
+        return meta
+
+    # ==============================
+    # Helper functions
+    # ==============================
+
+    def _found_match_for_mm_data(self, mm_hash) -> bool:
+        """Check if the cache is hit for the request."""
+        filename = self._generate_filename_debug(mm_hash)
+        return os.path.exists(filename)
+
+    def _generate_foldername_debug(
+        self,
+        mm_hash: str,
+        create_folder: bool = True,  # <- now defaults to True
+    ) -> str:
+        """
+        Return the folder in which the cache for this mm_hash lives.
+        If `create_folder` is True (default) the directory is created
+        recursively the first time it is needed.
+        """
+        foldername = os.path.join(self._storage_path, mm_hash)
+        if create_folder:
+            os.makedirs(foldername, exist_ok=True)
+        return foldername
+
+    def _generate_filename_debug(self, mm_hash: str) -> str:
+        """
+        Return the full path of the safetensors file for this mm_hash.
+        Ensures the parent directory exists because
+        `_generate_foldername_debug` is called with its default
+        (`create_folder=True`).
+        """
+        foldername = self._generate_foldername_debug(mm_hash)  # <- folder auto-created
+        return os.path.join(foldername, "encoder_cache.safetensors")

vllm/distributed/ec_transfer/ec_connector/factory.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import importlib
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+from vllm.distributed.ec_transfer.ec_connector.base import (
+    ECConnectorBase,
+    ECConnectorRole,
+)
+from vllm.logger import init_logger
+
+if TYPE_CHECKING:
+    from vllm.config import ECTransferConfig, VllmConfig
+
+logger = init_logger(__name__)
+
+
+class ECConnectorFactory:
+    _registry: dict[str, Callable[[], type[ECConnectorBase]]] = {}
+
+    @classmethod
+    def register_connector(cls, name: str, module_path: str, class_name: str) -> None:
+        """Register a connector with a lazy-loading module and class name."""
+        if name in cls._registry:
+            raise ValueError(f"Connector '{name}' is already registered.")
+
+        def loader() -> type[ECConnectorBase]:
+            module = importlib.import_module(module_path)
+            return getattr(module, class_name)
+
+        cls._registry[name] = loader
+
+    @classmethod
+    def create_connector(
+        cls,
+        config: "VllmConfig",
+        role: ECConnectorRole,
+    ) -> ECConnectorBase:
+        ec_transfer_config = config.ec_transfer_config
+        if ec_transfer_config is None:
+            raise ValueError("ec_transfer_config must be set to create a connector")
+        connector_cls = cls.get_connector_class(ec_transfer_config)
+        logger.info(
+            "Creating connector with name: %s and engine_id: %s",
+            connector_cls.__name__,
+            ec_transfer_config.engine_id,
+        )
+        # Connector is explicitly separated into two roles.
+        # Scheduler connector:
+        # - Co-locate with scheduler process
+        # - Should only be used inside the Scheduler class
+        # Worker connector:
+        # - Co-locate with worker process
+        return connector_cls(config, role)
+
+    @classmethod
+    def get_connector_class(
+        cls, ec_transfer_config: "ECTransferConfig"
+    ) -> type[ECConnectorBase]:
+        """Get the connector class by name."""
+        connector_name = ec_transfer_config.ec_connector
+        if connector_name is None:
+            raise ValueError("EC connect must not be None")
+        elif connector_name in cls._registry:
+            connector_cls = cls._registry[connector_name]()
+        else:
+            connector_module_path = ec_transfer_config.ec_connector_module_path
+            if connector_module_path is None:
+                raise ValueError(f"Unsupported connector type: {connector_name}")
+            connector_module = importlib.import_module(connector_module_path)
+            connector_cls = getattr(connector_module, connector_name)
+        return connector_cls
+
+
+# Register various connectors here.
+# The registration should not be done in each individual file, as we want to
+# only load the files corresponding to the current connector.
+
+ECConnectorFactory.register_connector(
+    "ECExampleConnector",
+    "vllm.distributed.ec_transfer.ec_connector.example_connector",
+    "ECExampleConnector",
+)

vllm/distributed/ec_transfer/ec_transfer_state.py

@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import TYPE_CHECKING
+
+from vllm.distributed.ec_transfer.ec_connector.base import (
+    ECConnectorBase,
+    ECConnectorRole,
+)
+from vllm.distributed.ec_transfer.ec_connector.factory import ECConnectorFactory
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+
+_EC_CONNECTOR_AGENT: ECConnectorBase | None = None
+
+
+def get_ec_transfer() -> ECConnectorBase:
+    assert _EC_CONNECTOR_AGENT is not None, "disaggregated EC cache is not initialized"
+    return _EC_CONNECTOR_AGENT
+
+
+def has_ec_transfer() -> bool:
+    return _EC_CONNECTOR_AGENT is not None
+
+
+def ensure_ec_transfer_initialized(vllm_config: "VllmConfig") -> None:
+    """
+    Initialize EC cache connector.
+    """
+
+    global _EC_CONNECTOR_AGENT
+
+    if vllm_config.ec_transfer_config is None:
+        return
+
+    if (
+        vllm_config.ec_transfer_config.is_ec_transfer_instance
+        and _EC_CONNECTOR_AGENT is None
+    ):
+        _EC_CONNECTOR_AGENT = ECConnectorFactory.create_connector(
+            config=vllm_config, role=ECConnectorRole.WORKER
+        )
+
+
+def ensure_ec_transfer_shutdown() -> None:
+    global _EC_CONNECTOR_AGENT
+    if _EC_CONNECTOR_AGENT is not None:
+        _EC_CONNECTOR_AGENT.shutdown()
+        _EC_CONNECTOR_AGENT = None

vllm/distributed/elastic_ep/elastic_execute.py

@@ -0,0 +1,686 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import copy
+import gc
+import weakref
+from collections.abc import Iterable, Sequence
+from dataclasses import replace
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.distributed import P2POp
+
+from vllm.compilation.counter import compilation_counter
+from vllm.compilation.cuda_graph import CUDAGraphWrapper
+from vllm.compilation.wrapper import reset_compile_wrapper
+from vllm.config import (
+    CompilationMode,
+    set_current_vllm_config,
+)
+from vllm.distributed import (
+    get_dp_group,
+    get_ep_group,
+    get_pcp_group,
+    get_tp_group,
+)
+from vllm.distributed.elastic_ep.standby_state import (
+    create_standby_groups,
+    get_standby_dp_group,
+    get_standby_ep_group,
+    pop_standby_groups,
+)
+from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator
+from vllm.distributed.parallel_state import (
+    _replace_active_groups,
+    get_eplb_group,
+    prepare_communication_buffer_for_model,
+)
+from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
+from vllm.model_executor.layers.fused_moe.eep_reconfigure import (
+    make_eep_staged_quant_method,
+)
+from vllm.utils import is_moe_layer
+from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
+from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper
+from vllm.v1.worker.workspace import lock_workspace, unlock_workspace
+
+logger = init_logger(__name__)
+
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
+        FusedMoEMethodBase,
+    )
+
+
+def batch_transfer_weights(
+    model: nn.Module,
+    is_sender: bool,
+    peer_rank: int,
+    dp_group: StatelessGroupCoordinator,
+    expert_weights: Sequence[Iterable[torch.Tensor]],
+) -> None:
+    device_comm = dp_group.device_communicator
+    if device_comm is None:
+        raise ValueError("No device communicator found")
+
+    expert_weights_set = set()
+    for weight_group in expert_weights:
+        for weight in weight_group:
+            expert_weights_set.add(weight.data_ptr())
+
+    state_dict = model.state_dict()
+    all_params = []
+
+    for name, param in state_dict.items():
+        if name.endswith("expert_map") or name.find("._shared_experts") != -1:
+            continue
+        if param.data_ptr() not in expert_weights_set:
+            all_params.append(param.data)
+
+    assert len(all_params) > 0
+    p2p_ops = []
+    for param in all_params:
+        op = object.__new__(P2POp)
+        if is_sender:
+            op.op = torch.distributed.isend
+            op.tensor = param
+        else:
+            op.op = torch.distributed.irecv
+            op.tensor = param
+        op.group_peer = peer_rank
+        p2p_ops.append(op)
+    device_comm.batch_isend_irecv(p2p_ops)
+
+
+def broadcast_expert_mapping(
+    physical_to_logical: torch.Tensor | None,
+    num_local_physical_experts: int | None,
+    num_logical_experts: int | None,
+    dp_group: StatelessGroupCoordinator,
+    device: torch.device,
+    src_rank: int = 0,
+) -> tuple[torch.Tensor, int, int]:
+    if dp_group.rank_in_group == src_rank:
+        assert physical_to_logical is not None
+        assert num_local_physical_experts is not None
+        assert num_logical_experts is not None
+        assert physical_to_logical.dtype == torch.int64
+        shape_tensor = torch.tensor(
+            list(physical_to_logical.shape), dtype=torch.int64, device="cpu"
+        )
+        metadata_tensor = torch.tensor(
+            [num_local_physical_experts, num_logical_experts],
+            dtype=torch.int64,
+            device="cpu",
+        )
+    else:
+        shape_tensor = torch.empty(2, dtype=torch.int64, device="cpu")
+        metadata_tensor = torch.empty(2, dtype=torch.int64, device="cpu")
+
+    shape_tensor = dp_group.tcp_store_group.broadcast(shape_tensor, src_rank)
+    metadata_tensor = dp_group.tcp_store_group.broadcast(metadata_tensor, src_rank)
+
+    if dp_group.rank_in_group != src_rank:
+        assert device is not None
+        physical_to_logical = torch.empty(
+            tuple(shape_tensor.tolist()),
+            dtype=torch.int64,
+            device=device,
+        )
+
+    assert physical_to_logical is not None
+    physical_to_logical = dp_group.broadcast(physical_to_logical, src_rank)
+    num_local_physical_experts = int(metadata_tensor[0].item())
+    num_logical_experts = int(metadata_tensor[1].item())
+
+    return physical_to_logical, num_local_physical_experts, num_logical_experts
+
+
+class ElasticEPScalingExecutor:
+    def __init__(self, worker):
+        self.worker_ref = weakref.ref(worker)
+        self.reconfig_request = None
+        self._staged_moe_quant_methods: dict[nn.Module, FusedMoEMethodBase] = {}
+
+    @property
+    def worker(self):
+        worker = self.worker_ref()
+        if worker is None:
+            raise RuntimeError("Worker has been garbage collected")
+        return worker
+
+    def execute(self, execute_method: str, *args, **kwargs):
+        method = getattr(self, execute_method, None)
+        if method is None:
+            raise ValueError(f"Unknown execute method: {execute_method}")
+        return method(*args, **kwargs)
+
+    def _set_eplb_suppressed(self, suppressed: bool) -> None:
+        self.worker.model_runner.eep_eplb_suppressed = suppressed
+        ep_group = get_standby_ep_group() or get_ep_group()
+        if ep_group.rank == 0:
+            logger.info(
+                "[Elastic EP] EPLB %s elastic scaling transition",
+                "disabled during" if suppressed else "re-enabled after",
+            )
+
+    def load_model(self) -> None:
+        (
+            expanded_physical_to_logical,
+            num_logical_experts,
+            old_num_physical_experts,
+        ) = self.receive_expert_mapping()
+        num_physical_experts = expanded_physical_to_logical.shape[1]
+        self.worker.parallel_config.eplb_config.num_redundant_experts = (
+            num_physical_experts - num_logical_experts
+        )
+        self.worker.load_model(load_dummy_weights=True)
+        self.worker.model_runner.setup_eplb_from_mapping(
+            expanded_physical_to_logical, old_num_physical_experts
+        )
+        self._set_eplb_suppressed(True)
+
+    def create_standby_groups(
+        self, reconfig_request: ReconfigureDistributedRequest
+    ) -> None:
+        self.reconfig_request = reconfig_request
+        new_dp_size = reconfig_request.new_data_parallel_size
+        old_dp_size = get_dp_group().world_size
+        world_size = self.worker.vllm_config.parallel_config.world_size
+        new_world_size_across_dp = world_size * new_dp_size
+        updated_config = copy.copy(self.worker.vllm_config)
+        updated_config.parallel_config = copy.deepcopy(
+            self.worker.vllm_config.parallel_config
+        )
+        updated_config.parallel_config.data_parallel_size = new_dp_size
+        with set_current_vllm_config(updated_config):
+            create_standby_groups(
+                new_dp_size=new_dp_size,
+                new_world_size_across_dp=new_world_size_across_dp,
+                master_ip=reconfig_request.new_data_parallel_master_ip,
+                coord_store_port=reconfig_request.coord_store_port,
+                enable_eplb=updated_config.parallel_config.enable_eplb,
+            )
+        if new_dp_size > old_dp_size:
+            self._set_eplb_suppressed(True)
+            eplb_state = self.worker.model_runner.eplb_state
+            if eplb_state is not None:
+                eplb_state.drain_async()
+        elif new_dp_size < old_dp_size:
+            self._stage_standby_moe_quant_methods()
+
+    def transfer_weights(self, old_dp_size: int, new_dp_size: int) -> None:
+        standby_dp_group = get_standby_dp_group()
+        assert standby_dp_group is not None
+        # Broadcast old_dp_size to all workers in standby group
+        if standby_dp_group.rank_in_group < old_dp_size:
+            old_dp_size_tensor = torch.tensor(
+                [old_dp_size], dtype=torch.int64, device="cpu"
+            )
+        else:
+            old_dp_size_tensor = torch.empty(1, dtype=torch.int64, device="cpu")
+        old_dp_size_tensor = standby_dp_group.tcp_store_group.broadcast(
+            old_dp_size_tensor, 0
+        )
+
+        num_new_workers = new_dp_size - old_dp_size
+        dp_rank = self.worker.vllm_config.parallel_config.data_parallel_rank
+
+        # Sender-receiver pairing: the first new_workers % old_dp_size
+        # senders get (k+1) contiguous receivers, the rest get k
+        # receivers.
+        num_dst_per_sender = num_new_workers // old_dp_size
+        remainder = num_new_workers % old_dp_size
+
+        if dp_rank < remainder:
+            recv_begin = dp_rank * (num_dst_per_sender + 1)
+            recv_end = recv_begin + num_dst_per_sender + 1
+        else:
+            recv_begin = (
+                remainder * (num_dst_per_sender + 1)
+                + (dp_rank - remainder) * num_dst_per_sender
+            )
+            recv_end = recv_begin + num_dst_per_sender
+
+        ranks_to_send = list(range(old_dp_size + recv_begin, old_dp_size + recv_end))
+
+        model = self.worker.model_runner.get_model()
+        for new_worker_rank in sorted(ranks_to_send):
+            batch_transfer_weights(
+                model=model,
+                is_sender=True,
+                peer_rank=new_worker_rank,
+                dp_group=standby_dp_group,
+                expert_weights=model.expert_weights,
+            )
+        torch.accelerator.synchronize()
+
+    def broadcast_expert_mapping(self) -> None:
+        standby_dp_group = get_standby_dp_group()
+        assert standby_dp_group is not None
+        model_config = self.worker.model_runner.model_config
+        eplb_state = self.worker.model_runner.eplb_state
+        assert eplb_state is not None
+        eplb_model_state = eplb_state.model_states[model_config.compute_hash()]
+        physical_to_logical = eplb_model_state.physical_to_logical_map
+        num_physical_experts = physical_to_logical.shape[1]
+        num_local_physical_experts = num_physical_experts // get_ep_group().world_size
+        num_logical_experts = eplb_model_state.logical_replica_count.shape[1]
+        broadcast_expert_mapping(
+            physical_to_logical=physical_to_logical,
+            num_local_physical_experts=num_local_physical_experts,
+            num_logical_experts=num_logical_experts,
+            dp_group=standby_dp_group,
+            src_rank=0,
+            device=self.worker.device,
+        )
+        # New workers enter load_model after receiving the expert mapping.
+        # Stage replacement MoE kernels before returning to the state machine
+        # so existing ranks can participate in collective EP comm creation.
+        self._stage_standby_moe_quant_methods()
+
+    def _make_eep_moe_config(self, module, dp_group, ep_group):
+        parallel_config = self.worker.vllm_config.parallel_config
+        tp_size = get_tp_group().world_size
+        sp_size = tp_size if parallel_config.use_sequence_parallel_moe else 1
+        moe_parallel_config = FusedMoEParallelConfig.make(
+            tp_size_=tp_size,
+            pcp_size_=get_pcp_group().world_size,
+            dp_size_=dp_group.world_size,
+            sp_size_=sp_size,
+            vllm_parallel_config=parallel_config,
+        )
+        return replace(
+            module.moe_config,
+            num_experts=module.moe_config.num_local_experts * ep_group.world_size,
+            moe_parallel_config=moe_parallel_config,
+        )
+
+    def _stage_standby_moe_quant_methods(self) -> None:
+        standby_dp_group = get_standby_dp_group()
+        standby_ep_group = get_standby_ep_group()
+        model = self.worker.model_runner.get_model()
+        moe_modules = [module for module in model.modules() if is_moe_layer(module)]
+        self._staged_moe_quant_methods.clear()
+        with set_current_vllm_config(self.worker.vllm_config):
+            for module in moe_modules:
+                staged_quant_method = make_eep_staged_quant_method(
+                    module,
+                    self._make_eep_moe_config(
+                        module,
+                        standby_dp_group,
+                        standby_ep_group,
+                    ),
+                )
+                if staged_quant_method is not None:
+                    self._staged_moe_quant_methods[module] = staged_quant_method
+
+    def _commit_staged_moe_quant_methods(self) -> None:
+        model = self.worker.model_runner.get_model()
+        moe_modules = [module for module in model.modules() if is_moe_layer(module)]
+        for module in moe_modules:
+            staged_quant_method = self._staged_moe_quant_methods.pop(module, None)
+            if staged_quant_method is None:
+                continue
+            assert staged_quant_method.moe_kernel is not None
+            module._replace_quant_method(staged_quant_method)
+            staged_quant_method.moe_kernel.prepare_finalize.on_commit()
+        self._staged_moe_quant_methods.clear()
+
+    def _release_cuda_graphs(self) -> None:
+        if isinstance(self.worker.model_runner.model, CUDAGraphWrapper):
+            wrapper = self.worker.model_runner.model
+            wrapper.concrete_cudagraph_entries = {}
+
+        elif isinstance(self.worker.model_runner.model, UBatchWrapper):
+            raise RuntimeError("DBO is not yet supported in elastic EP")
+
+        torch.compiler.reset()
+        with set_current_vllm_config(self.worker.vllm_config):
+            reset_compile_wrapper(self.worker.model_runner.get_model())
+
+        gc.collect()
+        torch.accelerator.synchronize()
+        torch.accelerator.empty_cache()
+
+    def switch_and_remove(self) -> None:
+        self._release_cuda_graphs()
+        _replace_active_groups(world=None, dp=None, ep=None, eplb=None, node_count=None)
+
+    def switch_and_prepare(self) -> None:
+        old_dp_size = get_dp_group().world_size
+        old_ep_size = get_ep_group().world_size
+
+        self._release_cuda_graphs()
+        _replace_active_groups(**pop_standby_groups())
+
+        parallel_config = self.worker.vllm_config.parallel_config
+        reconfig_request = self.reconfig_request
+        assert reconfig_request is not None
+        new_dp_size = reconfig_request.new_data_parallel_size
+        new_ep_size = get_ep_group().world_size
+
+        parallel_config.data_parallel_size = new_dp_size
+        if (
+            reconfig_request.new_data_parallel_rank
+            != ReconfigureRankType.KEEP_CURRENT_RANK
+        ):
+            parallel_config.data_parallel_rank = reconfig_request.new_data_parallel_rank
+        if (
+            reconfig_request.new_data_parallel_rank_local
+            != ReconfigureRankType.KEEP_CURRENT_RANK
+        ):
+            parallel_config.data_parallel_rank_local = (
+                reconfig_request.new_data_parallel_rank_local
+            )
+        parallel_config.data_parallel_master_ip = (
+            reconfig_request.new_data_parallel_master_ip
+        )
+        parallel_config.data_parallel_master_port = (
+            reconfig_request.new_data_parallel_master_port
+        )
+
+        # Reconfigure MoE modules with new EP size
+        moe_modules = [
+            module
+            for module in self.worker.model_runner.model.modules()
+            if is_moe_layer(module)
+        ]
+        num_local_experts = moe_modules[0].moe_config.num_local_experts
+        assert all(
+            module.moe_config.num_local_experts == num_local_experts
+            for module in moe_modules
+        ), "All MoE modules must have the same number of experts"
+        dp_group = get_dp_group()
+        ep_group = get_ep_group()
+        for module in moe_modules:
+            new_moe_config = self._make_eep_moe_config(module, dp_group, ep_group)
+            module._set_moe_config(new_moe_config)
+
+        # Update EPLB state
+        eplb_state = self.worker.model_runner.eplb_state
+        assert eplb_state is not None
+        model_config = self.worker.model_runner.model_config
+        eplb_model_state = eplb_state.model_states[model_config.compute_hash()]
+
+        num_physical_experts = num_local_experts * new_ep_size
+        num_logical_experts = eplb_model_state.logical_replica_count.shape[1]
+        parallel_config.eplb_config.num_redundant_experts = (
+            num_physical_experts - num_logical_experts
+        )
+        old_physical_to_logical = eplb_model_state.physical_to_logical_map
+        num_moe_layers = old_physical_to_logical.shape[0]
+        num_local_experts = eplb_model_state.expert_load_pass.shape[1] // old_ep_size
+        if new_dp_size > old_dp_size:
+            expanded_physical_to_logical = torch.full(
+                (num_moe_layers, num_local_experts * new_ep_size),
+                -1,
+                dtype=old_physical_to_logical.dtype,
+                device=old_physical_to_logical.device,
+            )
+            expanded_physical_to_logical[:, : num_local_experts * old_ep_size] = (
+                old_physical_to_logical
+            )
+            eplb_model_state.physical_to_logical_map = expanded_physical_to_logical
+
+        old_num_physical_experts = eplb_model_state.expert_load_pass.shape[1]
+        pad_size = num_physical_experts - old_num_physical_experts
+        if new_dp_size > old_dp_size:
+            assert pad_size > 0
+            expanded_expert_load_pass = F.pad(
+                eplb_model_state.expert_load_pass, (0, pad_size), value=0
+            )
+            expanded_expert_load_window = F.pad(
+                eplb_model_state.expert_load_window, (0, pad_size), value=0
+            )
+            eplb_model_state.expert_load_pass = expanded_expert_load_pass
+            eplb_model_state.expert_load_window = expanded_expert_load_window
+            eplb_state.num_valid_physical_experts = old_num_physical_experts
+        else:
+            assert pad_size < 0
+            eplb_model_state.expert_load_pass = eplb_model_state.expert_load_pass[
+                :, :num_physical_experts
+            ]
+            eplb_model_state.expert_load_window = eplb_model_state.expert_load_window[
+                :, :, :num_physical_experts
+            ]
+            eplb_state.num_valid_physical_experts = num_physical_experts
+
+        model = self.worker.model_runner.get_model()
+        model.expert_weights = []
+        with set_current_vllm_config(self.worker.vllm_config):
+            model.set_eplb_state(
+                eplb_model_state.expert_load_pass,
+                eplb_model_state.logical_to_physical_map,
+                eplb_model_state.logical_replica_count,
+            )
+            eplb_state._propagate_shared_tensors(
+                model, eplb_model_state.num_unpadded_tokens_tensors
+            )
+            model.update_physical_experts_metadata(
+                num_physical_experts=num_physical_experts,
+                num_local_physical_experts=num_local_experts,
+            )
+            self._commit_staged_moe_quant_methods()
+            # Legacy modular methods need to be recreated for the new EP size.
+            for module in moe_modules:
+                if getattr(module._quant_method, "wraps_legacy_quant_method", False):
+                    module._replace_quant_method(module._quant_method.old_quant_method)
+            prepare_communication_buffer_for_model(self.worker.model_runner.model)
+
+        eplb_model_state.expert_buffer = [
+            torch.empty_like(w) for w in model.expert_weights[0]
+        ]
+        assert parallel_config.eplb_config.communicator is not None, (
+            "EPLB communicator backend must be set by ParallelConfig"
+        )
+        eplb_model_state.communicator = create_eplb_communicator(
+            group_coordinator=get_eplb_group(),
+            backend=parallel_config.eplb_config.communicator,
+            expert_weights=model.expert_weights,
+            expert_buffer=eplb_model_state.expert_buffer,
+        )
+
+        if (
+            self.worker.vllm_config.compilation_config.mode
+            == CompilationMode.STOCK_TORCH_COMPILE
+        ):
+            # NOTE(yongji): when using stock torch.compile,
+            # torch.compile is triggered during GPUModelRunner's load_model()
+            # TODO(yongji):check do we need to re-trigger torch.compile here?
+            # any changes to the tensor shapes in execution should already
+            # be handled internally by torch.compile.
+            backend = self.worker.vllm_config.compilation_config.init_backend(
+                self.worker.vllm_config
+            )
+            compilation_counter.stock_torch_compile_count += 1
+            self.worker.model_runner.model.compile(fullgraph=True, backend=backend)
+
+        multi_block_table = self.worker.model_runner.input_batch.block_table
+        saved_block_tables: list[tuple[torch.Tensor, torch.Tensor]] = []
+        for bt in multi_block_table.block_tables:
+            saved_block_tables.append(
+                (bt.block_table.gpu.clone(), bt.block_table.cpu.clone())
+            )
+        multi_block_table.clear()
+
+        unlock_workspace()
+        self.worker.compile_or_warm_up_model()
+        lock_workspace()
+
+        for bt, (saved_gpu, saved_cpu) in zip(
+            multi_block_table.block_tables, saved_block_tables
+        ):
+            bt.block_table.gpu.copy_(saved_gpu)
+            bt.block_table.cpu.copy_(saved_cpu)
+        if new_dp_size < old_dp_size:
+            self._set_eplb_suppressed(False)
+
+    def _perform_eplb_reshuffle(
+        self, rank_mapping: dict[int, int] | None = None
+    ) -> None:
+        if get_ep_group().rank == 0:
+            logger.info("[Elastic EP] Starting expert resharding...")
+
+        eplb_state = self.worker.model_runner.eplb_state
+        assert eplb_state is not None
+
+        model_config = self.worker.model_runner.model_config
+        eplb_model_state = eplb_state.model_states[model_config.compute_hash()]
+        is_async_enabled = eplb_state.is_async
+        eplb_state.is_async = False
+        if rank_mapping is None:
+            eplb_state.rearrange()
+        else:
+            eplb_state.rearrange(rank_mapping=rank_mapping)
+        # NOTE(yongji): check whether we need to synchronize here
+        torch.accelerator.synchronize()
+        # reset expert_rearrangement_step to ensure all ranks are synchronized
+        eplb_state.expert_rearrangement_step = 0
+        eplb_state.num_valid_physical_experts = (
+            eplb_model_state.physical_to_logical_map.shape[1]
+        )
+        eplb_state.is_async = is_async_enabled
+        # Start the async worker thread if it doesn't exist yet (idempotent).
+        # This is needed for new workers after scale-up: they create EplbState
+        # in setup_eplb_from_mapping() but don't start the thread there because
+        # groups aren't ready yet.
+        eplb_state.start_async_loop()
+        if get_ep_group().rank == 0:
+            logger.info("[Elastic EP] Expert resharding completed")
+
+    def perform_eplb_reshuffle(self) -> None:
+        self._perform_eplb_reshuffle()
+        self._set_eplb_suppressed(False)
+
+    def perform_scale_down_eplb_reshuffle(self, new_dp_size: int) -> None:
+        self._set_eplb_suppressed(True)
+        eplb_state = self.worker.model_runner.eplb_state
+        if eplb_state is not None:
+            eplb_state.drain_async()
+        parallel_config = self.worker.vllm_config.parallel_config
+        tp_size = parallel_config.tensor_parallel_size
+        old_ep_size = parallel_config.data_parallel_size * tp_size
+        new_ep_size = new_dp_size * tp_size
+        rank_mapping = {
+            old_ep_rank: old_ep_rank if old_ep_rank < new_ep_size else -1
+            for old_ep_rank in range(old_ep_size)
+        }
+        self._perform_eplb_reshuffle(rank_mapping=rank_mapping)
+
+    def receive_weights(self) -> None:
+        dp_group = get_dp_group()
+        assert isinstance(dp_group, StatelessGroupCoordinator)
+        new_dp_size = dp_group.world_size
+        dp_rank = self.worker.vllm_config.parallel_config.data_parallel_rank
+
+        # Receive old_dp_size broadcasted during transfer_weights
+        old_dp_size_tensor = torch.empty(1, dtype=torch.int64, device="cpu")
+        old_dp_size_tensor = dp_group.tcp_store_group.broadcast(old_dp_size_tensor, 0)
+        old_dp_size = int(old_dp_size_tensor[0].item())
+
+        # Calculate which existing worker will send to this new worker
+        num_new_workers = new_dp_size - old_dp_size
+        new_worker_idx = dp_rank - old_dp_size
+        num_dst_per_sender = num_new_workers // old_dp_size
+        remainder = num_new_workers % old_dp_size
+
+        if new_worker_idx < remainder * (num_dst_per_sender + 1):
+            sender_rank = new_worker_idx // (num_dst_per_sender + 1)
+        else:
+            sender_rank = (
+                remainder
+                + (new_worker_idx - remainder * (num_dst_per_sender + 1))
+                // num_dst_per_sender
+            )
+
+        model = self.worker.model_runner.get_model()
+        batch_transfer_weights(
+            model=model,
+            is_sender=False,
+            peer_rank=sender_rank,
+            dp_group=dp_group,
+            expert_weights=model.expert_weights,
+        )
+        torch.accelerator.synchronize()
+
+    def receive_expert_mapping(self) -> tuple[torch.Tensor, int, int]:
+        dp_group = get_dp_group()
+        assert isinstance(dp_group, StatelessGroupCoordinator)
+        physical_to_logical, num_local_physical_experts, num_logical_experts = (
+            broadcast_expert_mapping(
+                physical_to_logical=None,
+                num_local_physical_experts=None,
+                num_logical_experts=None,
+                dp_group=dp_group,
+                src_rank=0,
+                device=self.worker.device,
+            )
+        )
+        num_moe_layers = physical_to_logical.shape[0]
+        new_dp_size = get_dp_group().world_size
+        tp_size = self.worker.vllm_config.parallel_config.tensor_parallel_size
+        new_ep_size = new_dp_size * tp_size
+        expanded_physical_to_logical = torch.full(
+            (num_moe_layers, num_local_physical_experts * new_ep_size),
+            -1,
+            dtype=physical_to_logical.dtype,
+            device=physical_to_logical.device,
+        )
+        old_num_physical_experts = physical_to_logical.shape[1]
+        expanded_physical_to_logical[:, :old_num_physical_experts] = physical_to_logical
+        return (
+            expanded_physical_to_logical,
+            num_logical_experts,
+            old_num_physical_experts,
+        )
+
+    def prepare_new_worker(self) -> None:
+        with set_current_vllm_config(self.worker.vllm_config):
+            prepare_communication_buffer_for_model(self.worker.model_runner.get_model())
+
+    def rewarm_workspace(self) -> None:
+        # Must run on every DP sibling in lockstep: _dummy_run calls
+        # coordinate_batch_across_dp whenever data_parallel_size > 1
+        # (gpu_model_runner.py:3663), which deadlocks if any rank skips it.
+
+        # Save and clear block tables so profile_run/compile_or_warm_up_model
+        # don't write dummy slot mappings into real KV-cache blocks (mirrors
+        # switch_and_prepare's pattern).
+        multi_block_table = self.worker.model_runner.input_batch.block_table
+        saved_block_tables: list[tuple[torch.Tensor, torch.Tensor]] = []
+        for bt in multi_block_table.block_tables:
+            saved_block_tables.append(
+                (bt.block_table.gpu.clone(), bt.block_table.cpu.clone())
+            )
+        multi_block_table.clear()
+
+        # _ensure_workspace_size allocates a fresh tensor on grow, leaving
+        # captured CUDA graphs with stale data pointers; drop graphs before
+        # re-warm so captures realign with the resized buffer.
+        self._release_cuda_graphs()
+        unlock_workspace()
+
+        # Grow the MoE workspace at max_num_tokens.
+        # compile_or_warm_up_model alone only exercises cudagraph-capture
+        # sizes (≤64 tokens for this test) and leaves the workspace at
+        # ~10-14 MB; the post-all-to-all per-rank token count under real
+        # post-reshuffle routing needs hundreds of MB. Use _dummy_run
+        # directly (rather than profile_run) with skip_eplb=True so dummy
+        # routing doesn't pollute the just-rebalanced EPLB stats — same
+        # convention compile_or_warm_up_model itself uses.
+        runner = self.worker.model_runner
+        runner._dummy_run(runner.max_num_tokens, is_profile=True, skip_eplb=True)
+        self.worker.compile_or_warm_up_model()
+
+        lock_workspace()
+
+        for bt, (saved_gpu, saved_cpu) in zip(
+            multi_block_table.block_tables, saved_block_tables
+        ):
+            bt.block_table.gpu.copy_(saved_gpu)
+            bt.block_table.cpu.copy_(saved_cpu)

vllm/distributed/elastic_ep/elastic_state.py

@@ -0,0 +1,593 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import enum
+import time
+import weakref
+from datetime import timedelta
+from typing import TYPE_CHECKING, Literal, TypeAlias
+
+import torch.distributed
+
+from vllm.config import ParallelConfig
+from vllm.distributed import (
+    sched_yield,
+    stateless_destroy_torch_distributed_process_group,
+)
+from vllm.logger import init_logger
+from vllm.v1.engine import (
+    EEPNotificationType,
+    ReconfigureDistributedRequest,
+    ReconfigureRankType,
+)
+from vllm.v1.engine.core import DPEngineCoreProc
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.v1.executor.abstract import Executor
+
+logger = init_logger(__name__)
+
+WorkerType = Literal["existing", "new", "removing"]
+
+
+class ScaleUpExistingEngineState(enum.IntEnum):
+    WAIT_NEW_CORE_ENGINES_INIT = 0
+    CREATE_STANDBY_GROUPS = 1
+    TRANSFER_EXPERT_MAPPING = 2
+    WAIT_NEW_CORE_ENGINES_WEIGHTS_INIT = 3
+    TRANSFER_WEIGHTS = 4
+    SYNC_KV_CACHE_MEMORY_SIZE = 5
+    SWITCH_AND_PREPARE = 6
+    EPLB_RESHUFFLE = 7
+    COMPLETE = 8
+
+
+class ScaleUpNewEngineState(enum.IntEnum):
+    PRE_KV_INIT = 0
+    PREPARE = 1
+    EPLB_RESHUFFLE = 2
+    COMPLETE = 3
+
+
+class ScaleDownRemainingEngineState(enum.IntEnum):
+    PREPARE = 0
+    EPLB_RESHUFFLE = 1
+    SWITCH_AND_PREPARE = 2
+    COMPLETE = 3
+
+
+class ScaleDownRemovingEngineState(enum.IntEnum):
+    PREPARE = 0
+    EPLB_RESHUFFLE = 1
+    COMPLETE = 2
+
+
+EngineState: TypeAlias = (
+    ScaleUpExistingEngineState
+    | ScaleUpNewEngineState
+    | ScaleDownRemainingEngineState
+    | ScaleDownRemovingEngineState
+)
+
+
+class _BarrierTimeoutError(RuntimeError):
+    """
+    Exception raised for timeout
+    in the first stage of our two-staged
+    TCPStore based barrier to synchronize the
+    execution of all engines in the DP group.
+    """
+
+
+class ElasticEPScalingState:
+    def __init__(
+        self,
+        model_executor: "Executor",
+        engine_core: "DPEngineCoreProc",
+        vllm_config: "VllmConfig",
+        new_parallel_config: ParallelConfig,
+        worker_type: WorkerType,
+        scale_type: Literal["scale_up", "scale_down"],
+        reconfig_request: ReconfigureDistributedRequest | None = None,
+    ):
+        self.model_executor_ref = weakref.ref(model_executor)
+        self.engine_core_ref = weakref.ref(engine_core)
+        self.vllm_config = vllm_config
+        self.old_dp_group = self.engine_core.dp_group if worker_type != "new" else None
+        self.old_dp_store = self.engine_core.dp_store if worker_type != "new" else None
+        self.new_parallel_config: ParallelConfig = new_parallel_config
+        self.new_dp_group = self.engine_core.dp_group if worker_type == "new" else None
+        self.new_dp_store = self.engine_core.dp_store if worker_type == "new" else None
+        self.worker_type = worker_type
+        self.scale_type = scale_type
+        self.reconfig_request = reconfig_request
+
+        self.state: EngineState
+        if scale_type == "scale_up":
+            self.state = (
+                ScaleUpNewEngineState.PRE_KV_INIT
+                if worker_type == "new"
+                else ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_INIT
+            )
+        else:
+            self.state = (
+                ScaleDownRemovingEngineState.PREPARE
+                if worker_type == "removing"
+                else ScaleDownRemainingEngineState.PREPARE
+            )
+
+    @property
+    def model_executor(self) -> "Executor":
+        model_executor = self.model_executor_ref()
+        if model_executor is None:
+            raise RuntimeError("Model executor has been garbage collected")
+        return model_executor
+
+    @property
+    def engine_core(self) -> "DPEngineCoreProc":
+        engine_core = self.engine_core_ref()
+        if engine_core is None:
+            raise RuntimeError("Engine core has been garbage collected")
+        return engine_core
+
+    def progress(self) -> bool:
+        if self.scale_type == "scale_up":
+            return (
+                self._progress_new_engine()
+                if self.worker_type == "new"
+                else self._progress_existing_engine()
+            )
+        return (
+            self._progress_removing_engine()
+            if self.worker_type == "removing"
+            else self._progress_remaining_engine()
+        )
+
+    def run_pre_kv_init_states(self) -> None:
+        assert self.scale_type == "scale_up" and self.worker_type == "new"
+        assert self.state == ScaleUpNewEngineState.PRE_KV_INIT
+        assert self.progress()
+        assert self.state == ScaleUpNewEngineState.PREPARE
+
+    def _execute_tcp_store_barrier(
+        self, dp_store, group_rank, group_size, barrier_id, timeout=None
+    ):
+        arrival_key = f"arrival_{barrier_id}_{group_rank}"
+        dp_store.set(arrival_key, b"1")
+
+        start_time = time.time()
+        processes_arrived: set[int] = set()
+
+        while len(processes_arrived) < group_size:
+            if (
+                timeout is not None
+                and time.time() - start_time > timeout.total_seconds()
+            ):
+                raise _BarrierTimeoutError(
+                    f"Barrier timed out after {timeout.total_seconds()} seconds"
+                )
+
+            for i in range(group_size):
+                if i in processes_arrived:
+                    continue
+
+                key = f"arrival_{barrier_id}_{i}"
+                present = dp_store.check([key])
+                if present:
+                    processes_arrived.add(i)
+
+            if len(processes_arrived) < group_size:
+                sched_yield()
+
+    def _staged_barrier(self, use_new_group: bool, barrier_name: str) -> bool:
+        """
+        Execute a two-staged barrier to synchronize all engines in the DP group.
+
+        Some DP EngineCores may receive the reconfiguration notifications
+        later than others, and already proceed to engine step (model forward)
+        in the busy loop.
+        In this case, EngineCores that already proceed to reconfiguration
+        should skip reconfiguration and execute model forward for one more
+        step, so in the next step, all EngineCores will be synchronized.
+        We use a two-staged barrier to achieve this. The first time each
+        EngineCore executes the barrier, if a timeout is reached before the
+        barrier completes, that means some EngineCores have already entered
+        engine step. The EngineCores that timed out will then proceed to
+        engine step, and will synchronize with the other EngineCores in the
+        next step with a barrier without timeout.
+        """
+        dp_group = self.new_dp_group if use_new_group else self.old_dp_group
+        dp_store = self.new_dp_store if use_new_group else self.old_dp_store
+        assert dp_group is not None and dp_store is not None
+
+        group_rank = dp_group.rank()
+        group_size = dp_group.size()
+        barrier_id = f"eep_barrier_{barrier_name}"
+        sync_key = f"{barrier_id}_sync"
+
+        # TODO(yongji): figure out appropriate timeout for the barrier
+        timeout = None if dp_store.check([sync_key]) else timedelta(seconds=5)
+
+        try:
+            self._execute_tcp_store_barrier(
+                dp_store, group_rank, group_size, barrier_id, timeout=timeout
+            )
+            torch.distributed.barrier(dp_group)
+            if group_rank == 0:
+                dp_store.delete_key(sync_key)
+                for i in range(group_size):
+                    dp_store.delete_key(f"arrival_{barrier_id}_{i}")
+            return True
+        except _BarrierTimeoutError as e:
+            if timeout is None:
+                raise RuntimeError("Unexpected timeout encountered") from e
+            dp_store.compare_set(sync_key, "", b"1")
+            return False
+
+    def _progress_existing_engine(self) -> bool:
+        state = self.state
+        assert self.old_dp_group is not None and self.old_dp_store is not None
+
+        if state == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_INIT:
+            return False
+
+        elif state == ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS:
+            # NOTE(yongji): wait for all existing workers to receive the request
+            if (
+                int(self.old_dp_store.get("eep_barrier_engine_count"))
+                < self.old_dp_group.size()
+            ):
+                return False
+            if not self._staged_barrier(
+                use_new_group=False, barrier_name="create_standby_groups"
+            ):
+                return False
+            if self.old_dp_group.rank() == 0:
+                self.old_dp_store.delete_key("eep_barrier_engine_count")
+            self._create_standby_groups()
+            self.state = ScaleUpExistingEngineState.TRANSFER_EXPERT_MAPPING
+            return True
+
+        elif state == ScaleUpExistingEngineState.TRANSFER_EXPERT_MAPPING:
+            self._transfer_expert_mapping()
+            self.state = ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_WEIGHTS_INIT
+            return True
+
+        elif state == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_WEIGHTS_INIT:
+            return False
+
+        elif state == ScaleUpExistingEngineState.TRANSFER_WEIGHTS:
+            if (
+                int(self.old_dp_store.get("eep_barrier_engine_count"))
+                < self.old_dp_group.size()
+            ):
+                return False
+            if not self._staged_barrier(
+                use_new_group=False, barrier_name="transfer_weights"
+            ):
+                return False
+            if self.old_dp_group.rank() == 0:
+                self.old_dp_store.delete_key("eep_barrier_engine_count")
+            self._transfer_weights()
+            self.state = ScaleUpExistingEngineState.SYNC_KV_CACHE_MEMORY_SIZE
+            return True
+
+        elif state == ScaleUpExistingEngineState.SYNC_KV_CACHE_MEMORY_SIZE:
+            self._sync_kv_cache_memory_size()
+            self.state = ScaleUpExistingEngineState.SWITCH_AND_PREPARE
+            return True
+
+        elif state == ScaleUpExistingEngineState.SWITCH_AND_PREPARE:
+            self._switch_and_prepare()
+            self.state = ScaleUpExistingEngineState.EPLB_RESHUFFLE
+            assert self.new_dp_store is not None
+            self.new_dp_store.add("eep_barrier_engine_count", 1)
+            return True
+
+        elif state == ScaleUpExistingEngineState.EPLB_RESHUFFLE:
+            assert self.new_dp_group is not None and self.new_dp_store is not None
+            if (
+                int(self.new_dp_store.get("eep_barrier_engine_count"))
+                < self.new_dp_group.size()
+            ):
+                return False
+            if not self._staged_barrier(
+                use_new_group=True, barrier_name="eplb_reshuffle"
+            ):
+                return False
+            if self.new_dp_group.rank() == 0:
+                self.new_dp_store.delete_key("eep_barrier_engine_count")
+            self._eplb_reshuffle()
+            self.state = ScaleUpExistingEngineState.COMPLETE
+            self._update_parallel_config()
+            return True
+
+        else:
+            assert self.state == ScaleUpExistingEngineState.COMPLETE
+            return True
+
+    def _progress_new_engine(self) -> bool:
+        state = self.state
+        assert self.new_dp_group is not None and self.new_dp_store is not None
+
+        if state == ScaleUpNewEngineState.PRE_KV_INIT:
+            self.engine_core._eep_send_engine_core_notification(
+                EEPNotificationType.NEW_CORE_ENGINES_WEIGHTS_INIT_READY
+            )
+            self.model_executor.collective_rpc(
+                "elastic_ep_execute", args=("receive_weights",)
+            )
+            self.engine_core.available_gpu_memory_for_kv_cache = (
+                ParallelConfig.sync_kv_cache_memory_size(self.new_dp_group, -1)
+            )
+            self.model_executor.collective_rpc(
+                "elastic_ep_execute", args=("prepare_new_worker",)
+            )
+            self.state = ScaleUpNewEngineState.PREPARE
+            return True
+
+        elif state == ScaleUpNewEngineState.PREPARE:
+            tensor = torch.tensor([0, 0, 0], dtype=torch.int32, device="cpu")
+            torch.distributed.all_reduce(
+                tensor,
+                op=torch.distributed.ReduceOp.MAX,
+                group=self.new_dp_group,
+            )
+            data = tensor.tolist()
+            self.engine_core.engines_running = bool(data[0])
+            self.engine_core.current_wave = int(data[1])
+            self.engine_core.step_counter = int(data[2])
+            self.state = ScaleUpNewEngineState.EPLB_RESHUFFLE
+            self.new_dp_store.add("eep_barrier_engine_count", 1)
+            return True
+
+        elif state == ScaleUpNewEngineState.EPLB_RESHUFFLE:
+            if (
+                int(self.new_dp_store.get("eep_barrier_engine_count"))
+                < self.new_dp_group.size()
+            ):
+                return False
+            if not self._staged_barrier(
+                use_new_group=True, barrier_name="eplb_reshuffle"
+            ):
+                return False
+            assert self.new_dp_group.rank() > 0
+            self._eplb_reshuffle()
+            self.state = ScaleUpNewEngineState.COMPLETE
+            return True
+
+        else:
+            assert self.state == ScaleUpNewEngineState.COMPLETE
+            return True
+
+    def _progress_remaining_engine(self) -> bool:
+        state = self.state
+        assert self.old_dp_group is not None and self.old_dp_store is not None
+
+        if state == ScaleDownRemainingEngineState.PREPARE:
+            self.state = ScaleDownRemainingEngineState.EPLB_RESHUFFLE
+            self.old_dp_store.add("eep_barrier_engine_count", 1)
+            return True
+
+        elif state == ScaleDownRemainingEngineState.EPLB_RESHUFFLE:
+            if (
+                int(self.old_dp_store.get("eep_barrier_engine_count"))
+                < self.old_dp_group.size()
+            ):
+                return False
+            if not self._staged_barrier(
+                use_new_group=False, barrier_name="eplb_reshuffle"
+            ):
+                return False
+            if self.old_dp_group.rank() == 0:
+                self.old_dp_store.delete_key("eep_barrier_engine_count")
+            self._eplb_reshuffle_before_scale_down()
+            self.state = ScaleDownRemainingEngineState.SWITCH_AND_PREPARE
+            # NOTE(yongji): currently, after EPLB reshuffle
+            # that redistributes experts to remaining workers, workers
+            # to be removed will immediately initiate shutdown;
+            # existing workers can no longer execute forward steps using
+            # the old setup. In the future, we may keep
+            # the removing workers alive a bit longer,
+            # e.g., to drain in-batch requests.
+            self._create_standby_groups()
+            self._switch_and_prepare()
+            self._update_parallel_config()
+            self.state = ScaleDownRemainingEngineState.COMPLETE
+            return True
+
+        else:
+            assert self.state == ScaleDownRemainingEngineState.COMPLETE
+            return True
+
+    def _progress_removing_engine(self) -> bool:
+        state = self.state
+        assert self.old_dp_group is not None and self.old_dp_store is not None
+
+        if state == ScaleDownRemovingEngineState.PREPARE:
+            self.state = ScaleDownRemovingEngineState.EPLB_RESHUFFLE
+            self.old_dp_store.add("eep_barrier_engine_count", 1)
+            return True
+
+        if state == ScaleDownRemovingEngineState.EPLB_RESHUFFLE:
+            if (
+                int(self.old_dp_store.get("eep_barrier_engine_count"))
+                < self.old_dp_group.size()
+            ):
+                return False
+            if not self._staged_barrier(
+                use_new_group=False, barrier_name="eplb_reshuffle"
+            ):
+                return False
+            assert self.old_dp_group.rank() > 0
+            self._eplb_reshuffle_before_scale_down()
+            self._switch_and_remove()
+            self.state = ScaleDownRemovingEngineState.COMPLETE
+            self.engine_core._eep_send_engine_core_notification(
+                EEPNotificationType.SHUTDOWN_COMPLETE
+            )
+            return True
+
+        else:
+            assert self.state == ScaleDownRemovingEngineState.COMPLETE
+            return True
+
+    def handle_notification(self, notification_type: EEPNotificationType):
+        assert self.worker_type != "new"
+        assert self.old_dp_store is not None
+        if (
+            notification_type == EEPNotificationType.NEW_CORE_ENGINES_INIT_READY
+            and self.state == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_INIT
+        ):
+            self.old_dp_store.add("eep_barrier_engine_count", 1)
+            self.state = ScaleUpExistingEngineState.CREATE_STANDBY_GROUPS
+        elif (
+            notification_type == EEPNotificationType.NEW_CORE_ENGINES_WEIGHTS_INIT_READY
+            and self.state
+            == ScaleUpExistingEngineState.WAIT_NEW_CORE_ENGINES_WEIGHTS_INIT
+        ):
+            self.old_dp_store.add("eep_barrier_engine_count", 1)
+            self.state = ScaleUpExistingEngineState.TRANSFER_WEIGHTS
+
+    def is_complete(self) -> bool:
+        if self.scale_type == "scale_up":
+            return (
+                self.state == ScaleUpNewEngineState.COMPLETE
+                if self.worker_type == "new"
+                else self.state == ScaleUpExistingEngineState.COMPLETE
+            )
+        return (
+            self.state == ScaleDownRemovingEngineState.COMPLETE
+            if self.worker_type == "removing"
+            else self.state == ScaleDownRemainingEngineState.COMPLETE
+        )
+
+    def _create_standby_groups(self):
+        assert self.old_dp_group is not None
+        self.new_dp_group, self.new_dp_store = (
+            self.new_parallel_config.stateless_init_dp_group(return_store=True)
+        )
+        self.model_executor.collective_rpc(
+            "elastic_ep_execute", args=("create_standby_groups", self.reconfig_request)
+        )
+        if self.old_dp_group.rank() == 0:
+            logger.info("[Elastic EP] Created standby communication groups")
+
+    def _transfer_weights(self):
+        assert self.reconfig_request is not None and self.old_dp_group is not None
+        old_dp_size = self.old_dp_group.size()
+        new_dp_size = self.reconfig_request.new_data_parallel_size
+
+        self.model_executor.collective_rpc(
+            "elastic_ep_execute", args=("transfer_weights", old_dp_size, new_dp_size)
+        )
+        if self.old_dp_group.rank() == 0:
+            logger.info("[Elastic EP] Transferred weights to new workers")
+
+    def _transfer_expert_mapping(self):
+        assert self.old_dp_group is not None
+        self.model_executor.collective_rpc(
+            "elastic_ep_execute", args=("broadcast_expert_mapping",)
+        )
+        if self.old_dp_group.rank() == 0:
+            logger.info("[Elastic EP] Broadcasted expert mapping to new workers")
+
+    def _sync_kv_cache_memory_size(self):
+        assert self.engine_core.available_gpu_memory_for_kv_cache > 0
+        assert self.new_dp_group is not None and self.old_dp_group is not None
+        ParallelConfig.sync_kv_cache_memory_size(
+            self.new_dp_group,
+            self.engine_core.available_gpu_memory_for_kv_cache,
+        )
+        if self.old_dp_group.rank() == 0:
+            logger.info("[Elastic EP] Synced KV cache memory size to new workers")
+
+    def _switch_and_prepare(self):
+        self.model_executor.collective_rpc(
+            "elastic_ep_execute", args=("switch_and_prepare",)
+        )
+        old_dp_group = self.old_dp_group
+        stateless_destroy_torch_distributed_process_group(old_dp_group)
+        assert self.new_dp_group is not None
+        new_dp_group = self.new_dp_group
+        self.engine_core.dp_group = new_dp_group
+        self.engine_core.dp_rank = new_dp_group.rank()
+        self.engine_core.dp_store = self.new_dp_store
+        engines_running = int(self.engine_core.engines_running)
+        current_wave = self.engine_core.current_wave
+        step_counter = self.engine_core.step_counter
+        tensor = torch.tensor(
+            [engines_running, current_wave, step_counter],
+            dtype=torch.int32,
+            device="cpu",
+        )
+        torch.distributed.all_reduce(
+            tensor, op=torch.distributed.ReduceOp.MAX, group=new_dp_group
+        )
+        data = tensor.tolist()
+        self.engine_core.engines_running = bool(data[0])
+        self.engine_core.current_wave = int(data[1])
+        self.engine_core.step_counter = int(data[2])
+        if new_dp_group.rank() == 0:
+            self.engine_core._eep_send_engine_core_notification(
+                EEPNotificationType.RECONFIGURE_FINISHED
+            )
+            logger.info("[Elastic EP] Switched to new setup")
+
+    def _eplb_reshuffle(self):
+        self.model_executor.collective_rpc(
+            "elastic_ep_execute", args=("perform_eplb_reshuffle",)
+        )
+        # Reshuffle changes per-rank token routing; the locked MoE workspace
+        # may now be too small. Rewarm covers both new and existing engines.
+        self.model_executor.collective_rpc(
+            "elastic_ep_execute", args=("rewarm_workspace",)
+        )
+        assert self.new_dp_group is not None
+        if self.new_dp_group.rank() == 0:
+            logger.info("[Elastic EP] EPLB reshuffle completed")
+
+    def _eplb_reshuffle_before_scale_down(self):
+        assert self.reconfig_request is not None and self.old_dp_group is not None
+        self.model_executor.collective_rpc(
+            "elastic_ep_execute",
+            args=(
+                "perform_scale_down_eplb_reshuffle",
+                self.reconfig_request.new_data_parallel_size,
+            ),
+        )
+        if self.old_dp_group.rank() == 0:
+            logger.info("[Elastic EP] EPLB reshuffle completed")
+
+    def _switch_and_remove(self):
+        self.model_executor.collective_rpc(
+            "elastic_ep_execute", args=("switch_and_remove",)
+        )
+
+    def _update_parallel_config(self):
+        assert self.reconfig_request is not None
+        reconfig_request = self.reconfig_request
+        parallel_config = self.vllm_config.parallel_config
+        parallel_config.data_parallel_size = reconfig_request.new_data_parallel_size
+        if (
+            reconfig_request.new_data_parallel_rank
+            != ReconfigureRankType.KEEP_CURRENT_RANK
+        ):
+            parallel_config.data_parallel_rank = reconfig_request.new_data_parallel_rank
+        if (
+            reconfig_request.new_data_parallel_rank_local
+            != ReconfigureRankType.KEEP_CURRENT_RANK
+        ):
+            parallel_config.data_parallel_rank_local = (
+                reconfig_request.new_data_parallel_rank_local
+            )
+        parallel_config.data_parallel_master_ip = (
+            reconfig_request.new_data_parallel_master_ip
+        )
+        parallel_config.data_parallel_master_port = (
+            reconfig_request.new_data_parallel_master_port
+        )
+        parallel_config._data_parallel_master_port_list = (
+            reconfig_request.new_data_parallel_master_port_list
+        )
+        parallel_config._coord_store_port = reconfig_request.coord_store_port

vllm/distributed/elastic_ep/standby_state.py

@@ -0,0 +1,123 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+
+from vllm.distributed.parallel_state import (
+    _init_stateless_group,
+    _node_count,
+    get_pp_group,
+    get_tp_group,
+    get_world_group,
+)
+from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator
+
+_STANDBY_WORLD: StatelessGroupCoordinator | None = None
+_STANDBY_WORLD_NODE_COUNT: int | None = None
+_STANDBY_DP: StatelessGroupCoordinator | None = None
+_STANDBY_EP: StatelessGroupCoordinator | None = None
+_STANDBY_EPLB: StatelessGroupCoordinator | None = None
+
+
+def get_standby_dp_group() -> StatelessGroupCoordinator | None:
+    return _STANDBY_DP
+
+
+def get_standby_ep_group() -> StatelessGroupCoordinator | None:
+    return _STANDBY_EP
+
+
+def get_standby_eplb_group() -> StatelessGroupCoordinator | None:
+    return _STANDBY_EPLB
+
+
+def get_standby_world_group() -> StatelessGroupCoordinator | None:
+    return _STANDBY_WORLD
+
+
+def create_standby_groups(
+    new_dp_size: int,
+    new_world_size_across_dp: int,
+    master_ip: str,
+    coord_store_port: int,
+    enable_eplb: bool = True,
+    backend: str | None = None,
+) -> None:
+    global \
+        _STANDBY_WORLD, \
+        _STANDBY_WORLD_NODE_COUNT, \
+        _STANDBY_DP, \
+        _STANDBY_EP, \
+        _STANDBY_EPLB
+
+    from vllm.distributed.utils import get_cached_tcp_store_client
+
+    assert new_world_size_across_dp == torch.distributed.get_world_size() * new_dp_size
+    world_group = get_world_group()
+    assert isinstance(world_group, StatelessGroupCoordinator)
+    backend = backend or world_group.backend
+
+    coord_store = get_cached_tcp_store_client(master_ip, coord_store_port)
+
+    standby_world_ranks = [list(range(new_world_size_across_dp))]
+    _STANDBY_WORLD = _init_stateless_group(
+        standby_world_ranks,
+        "world",
+        master_ip,
+        backend,
+        use_device_communicator=False,
+        coord_store=coord_store,
+    )
+    _STANDBY_WORLD_NODE_COUNT = _node_count(_STANDBY_WORLD.tcp_store_group)
+
+    tp_size = get_tp_group().world_size
+    pp_size = get_pp_group().world_size
+
+    all_ranks = torch.arange(new_world_size_across_dp).reshape(
+        -1, new_dp_size, pp_size, tp_size
+    )
+    standby_dp_ranks = all_ranks.transpose(1, 3).reshape(-1, new_dp_size).unbind(0)
+    standby_dp_ranks = [x.tolist() for x in standby_dp_ranks]
+    _STANDBY_DP = _init_stateless_group(
+        standby_dp_ranks, "dp", master_ip, backend, coord_store=coord_store
+    )
+
+    standby_ep_ranks = (
+        all_ranks.transpose(1, 2).reshape(-1, new_dp_size * tp_size).unbind(0)
+    )
+    standby_ep_ranks = [x.tolist() for x in standby_ep_ranks]
+    _STANDBY_EP = _init_stateless_group(
+        standby_ep_ranks, "ep", master_ip, backend, coord_store=coord_store
+    )
+
+    if enable_eplb:
+        _STANDBY_EPLB = _init_stateless_group(
+            standby_ep_ranks,
+            "eplb",
+            master_ip,
+            backend,
+            coord_store=coord_store,
+        )
+
+
+def pop_standby_groups() -> dict:
+    """Return all standby groups and clear the standby state."""
+    global \
+        _STANDBY_WORLD, \
+        _STANDBY_WORLD_NODE_COUNT, \
+        _STANDBY_DP, \
+        _STANDBY_EP, \
+        _STANDBY_EPLB
+
+    result = dict(
+        world=_STANDBY_WORLD,
+        dp=_STANDBY_DP,
+        ep=_STANDBY_EP,
+        eplb=_STANDBY_EPLB,
+        node_count=_STANDBY_WORLD_NODE_COUNT,
+    )
+    _STANDBY_WORLD = None
+    _STANDBY_WORLD_NODE_COUNT = None
+    _STANDBY_DP = None
+    _STANDBY_EP = None
+    _STANDBY_EPLB = None
+    return result

vllm/distributed/eplb/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Expert parallelism load balancer (EPLB)."""

vllm/distributed/eplb/async_worker.py

@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+The async worker that transfers experts in the background.
+"""
+
+import threading
+from typing import TYPE_CHECKING
+
+import torch
+
+from vllm.distributed.parallel_state import get_eplb_group
+from vllm.logger import init_logger
+
+from .eplb_utils import CpuGpuEvent
+from .rebalance_execute import AsyncEplbLayerResult, transfer_layer
+
+if TYPE_CHECKING:
+    from .eplb_state import EplbModelState, EplbState
+
+logger = init_logger(__name__)
+
+
+def start_async_worker(
+    state: "EplbState",
+    is_profile: bool = False,
+) -> threading.Thread:
+    rank = get_eplb_group().device_group.rank()
+    device_index = state.cuda_device_index
+    assert state.is_async
+
+    def thread_target() -> None:
+        assert device_index is not None
+        torch.accelerator.set_device_index(device_index)
+        cuda_stream = torch.cuda.Stream(device=device_index)
+        try:
+            transfer_run_periodically(
+                state=state,
+                cuda_stream=cuda_stream,
+                is_profile=is_profile,
+            )
+        except Exception as exc:  # pragma: no cover - diagnostic path
+            logger.exception("async loop error (Rank %d): %s", rank, str(exc))
+
+    thread = threading.Thread(target=thread_target, daemon=True)
+    thread.start()
+    return thread
+
+
+def run_rebalance_experts(
+    model_state: "EplbModelState",
+    eplb_state: "EplbState",
+    physical_to_logical_map_cpu: torch.Tensor,
+    cuda_stream: torch.cuda.Stream,
+) -> torch.Tensor:
+    assert model_state.eplb_stats is not None
+    eplb_stats = model_state.eplb_stats
+
+    # Move the global expert load window to CPU for computation.
+    with torch.cuda.stream(cuda_stream):
+        global_expert_load_window = eplb_stats.global_expert_load_window.cpu()
+    # Compute new expert mappings for the model
+    new_physical_to_logical_map = eplb_state.policy.rebalance_experts(
+        global_expert_load_window,
+        eplb_stats.num_replicas,
+        eplb_stats.num_groups,
+        eplb_stats.num_nodes,
+        eplb_stats.num_gpus,
+        physical_to_logical_map_cpu,
+    )
+    assert new_physical_to_logical_map.device == torch.device("cpu")
+
+    return new_physical_to_logical_map
+
+
+def transfer_run_periodically(
+    state: "EplbState",
+    cuda_stream: torch.cuda.Stream,
+    is_profile: bool = False,
+) -> None:
+    while True:
+        state.rearrange_event.wait(stream=cuda_stream)
+
+        eplb_group = get_eplb_group().device_group
+        eplb_cpu_group = get_eplb_group().cpu_group
+        ep_rank = eplb_group.rank()
+
+        assert state.is_async
+        for model_state in state.model_states.values():
+            layer_idx = 0
+            # Set the async worker's CUDA stream on the communicator
+            model_state.communicator.set_stream(cuda_stream)
+            num_layers = model_state.model.num_moe_layers
+
+            # Snapshot the physical_to_logical_map (synchronized with
+            # rearrange_event) and copy it to CPU
+            with torch.cuda.stream(cuda_stream):
+                physical_to_logical_map_cpu = model_state.physical_to_logical_map.cpu()
+
+            new_physical_to_logical_map = run_rebalance_experts(
+                model_state, state, physical_to_logical_map_cpu, cuda_stream
+            )
+
+            # Execute one EPLB layer transfer per model forward pass. Each iteration
+            # of this loop will copy the new set of expert weights into
+            # model_state.expert_buffer, which will be consumed by the main thread in
+            # move_to_workspace.
+            # We sync the rebalanced flag across ranks before each iteration so
+            # all ranks make a coordinated decision to continue or stop.
+            while layer_idx < num_layers:
+                flag = torch.tensor(
+                    [int(model_state.rebalanced)],
+                    dtype=torch.int32,
+                    device="cpu",
+                )
+                torch.distributed.all_reduce(flag, group=eplb_cpu_group)
+                if int(flag.item()) != eplb_cpu_group.size():
+                    logger.warning(
+                        "async worker (rank=%d): layer %d coordinated stop "
+                        "(flag_sum=%d, group_size=%d)",
+                        ep_rank,
+                        layer_idx,
+                        int(flag.item()),
+                        eplb_cpu_group.size(),
+                    )
+                    model_state.rebalanced = False
+                    break
+
+                transfer_metadata = transfer_layer(
+                    old_layer_indices=physical_to_logical_map_cpu[layer_idx],
+                    new_layer_indices=new_physical_to_logical_map[layer_idx],
+                    expert_weights=model_state.model.expert_weights[layer_idx],
+                    expert_weights_buffer=model_state.expert_buffer,
+                    communicator=model_state.communicator,
+                    ep_group=eplb_group,
+                    is_profile=is_profile,
+                    cuda_stream=cuda_stream,
+                    layer_idx=layer_idx,
+                )
+
+                # Wait until all writes to expert_buffer have finished before making the
+                # AsyncEplbLayerResult visible to the main thread.
+                cuda_stream.synchronize()
+
+                # This event guarantees that expert_buffer will not be overwritten by
+                # subsequent iterations of this loop until the main thread has consumed
+                # it. Record is called by the main thread after move_from_buffer().
+                consumed_event = CpuGpuEvent()
+
+                model_state.pending_result = AsyncEplbLayerResult(
+                    layer_idx=layer_idx,
+                    new_physical_to_logical_map=new_physical_to_logical_map[layer_idx],
+                    transfer_metadata=transfer_metadata,
+                    consumed_event=consumed_event,
+                )
+
+                # Block this thread until the main thread and main stream
+                # finish copying model_state.expert_buffer into
+                # model_state.model.expert_weights[layer_idx]
+                consumed_event.wait(stream=cuda_stream)
+                assert model_state.pending_result is None
+                layer_idx += 1

vllm/distributed/eplb/eplb_communicator.py

@@ -0,0 +1,776 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+EPLB communicator implementations and factory.
+"""
+
+import contextlib
+import time
+import uuid
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from datetime import timedelta
+
+import numpy as np
+import torch
+from torch.distributed import (
+    P2POp,
+    ProcessGroup,
+    batch_isend_irecv,
+)
+
+import vllm.distributed.nixl_utils as nixl_utils
+from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
+from vllm.distributed.device_communicators.pynccl_wrapper import (
+    ncclDataTypeEnum,
+)
+from vllm.distributed.parallel_state import (
+    GroupCoordinator,
+    get_pp_group,
+    is_local_first_rank,
+)
+from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator
+from vllm.distributed.utils import is_weak_contiguous
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+logger = init_logger(__name__)
+
+
+def has_nixl() -> bool:
+    """Whether the optional NIXL / RIXL package is available."""
+    return nixl_utils.NixlWrapper is not None
+
+
+class EplbCommunicator(ABC):
+    """Abstract EPLB communicator for expert weight transfers."""
+
+    @abstractmethod
+    def add_send(
+        self,
+        tensors: list[torch.Tensor],
+        dst_rank: int,
+        expert_id: int,
+    ) -> None:
+        pass
+
+    @abstractmethod
+    def add_recv(
+        self,
+        tensors: list[torch.Tensor],
+        src_rank: int,
+        expert_id: int,
+    ) -> None:
+        pass
+
+    @abstractmethod
+    def execute(self) -> None:
+        """Complete all enqueued transfers.
+
+        Some backends perform communication here; others (e.g. NIXL)
+        issue transfers eagerly in add_recv and only wait here.
+        On return, all data is available in the destination buffers.
+        """
+
+    def set_transfer_context(  # noqa: B027
+        self, old_indices: np.ndarray, layer_idx: int
+    ) -> None:
+        """Pre-set layer context before add_recv calls.
+
+        Default is a no-op; overridden by backends (e.g. NIXL) that need
+        layer-level context to issue transfers inside add_recv.
+        """
+
+    @property
+    def needs_profile_buffer_reservation(self) -> bool:
+        """Whether the profile path must run a dummy collective operation to reserve
+        communication buffers."""
+        return True
+
+    def set_stream(self, cuda_stream: torch.cuda.Stream | None) -> None:
+        self._cuda_stream = cuda_stream
+
+    def _log_initialized(self) -> None:
+        if is_local_first_rank():
+            logger.info("Initialized EPLB communicator: %s.", self.__class__.__name__)
+
+
+class TorchDistNcclEplbCommunicator(EplbCommunicator):
+    """EPLB communicator backed by torch.distributed isend/irecv."""
+
+    def __init__(
+        self,
+        ep_group: ProcessGroup,
+        cuda_stream: torch.cuda.Stream | None = None,
+    ) -> None:
+        self._ep_group = ep_group
+        self._cuda_stream = cuda_stream
+        self._p2p_ops: list[P2POp] = []
+        self._log_initialized()
+
+    def add_send(
+        self,
+        tensors: list[torch.Tensor],
+        dst_rank: int,
+        expert_id: int,  # unused by this backend
+    ) -> None:
+        for tensor in tensors:
+            self._p2p_ops.append(
+                P2POp(
+                    torch.distributed.isend,
+                    tensor,
+                    dst_rank,
+                    self._ep_group,
+                )
+            )
+
+    def add_recv(
+        self,
+        tensors: list[torch.Tensor],
+        src_rank: int,
+        expert_id: int,  # unused by this backend
+    ) -> None:
+        for tensor in tensors:
+            self._p2p_ops.append(
+                P2POp(
+                    torch.distributed.irecv,
+                    tensor,
+                    src_rank,
+                    self._ep_group,
+                )
+            )
+
+    def execute(self) -> None:
+        if not self._p2p_ops:
+            return
+        try:
+            with torch.cuda.stream(self._cuda_stream):
+                reqs = batch_isend_irecv(self._p2p_ops)
+                for req in reqs:
+                    req.wait()
+        finally:
+            self._p2p_ops.clear()
+
+
+class TorchDistGlooStagedEplbCommunicator(EplbCommunicator):
+    """EPLB communicator using gloo P2P with CPU staging."""
+
+    def __init__(
+        self,
+        cpu_group: ProcessGroup,
+        cuda_stream: torch.cuda.Stream | None = None,
+    ) -> None:
+        self._cpu_group = cpu_group
+        self._cuda_stream = cuda_stream
+        self._ops: list[tuple[str, torch.Tensor, int]] = []
+        self._log_initialized()
+
+    def add_send(
+        self,
+        tensors: list[torch.Tensor],
+        dst_rank: int,
+        expert_id: int,  # unused by this backend
+    ) -> None:
+        for tensor in tensors:
+            self._ops.append(("send", tensor, dst_rank))
+
+    def add_recv(
+        self,
+        tensors: list[torch.Tensor],
+        src_rank: int,
+        expert_id: int,  # unused by this backend
+    ) -> None:
+        for tensor in tensors:
+            self._ops.append(("recv", tensor, src_rank))
+
+    def execute(self) -> None:
+        if not self._ops:
+            return
+
+        p2p_ops: list[P2POp] = []
+        recv_staging: list[tuple[torch.Tensor, torch.Tensor]] = []
+
+        def build_ops() -> None:
+            for op, tensor, peer_rank in self._ops:
+                if op == "send":
+                    cpu_tensor = tensor.to(device="cpu", non_blocking=True)
+                    p2p_ops.append(
+                        P2POp(
+                            torch.distributed.isend,
+                            cpu_tensor,
+                            peer_rank,
+                            self._cpu_group,
+                        )
+                    )
+                    continue
+                cpu_tensor = torch.empty_like(tensor, device="cpu")
+                p2p_ops.append(
+                    P2POp(
+                        torch.distributed.irecv,
+                        cpu_tensor,
+                        peer_rank,
+                        self._cpu_group,
+                    )
+                )
+                recv_staging.append((tensor, cpu_tensor))
+
+        try:
+            with torch.cuda.stream(self._cuda_stream):
+                build_ops()
+        finally:
+            self._ops.clear()
+
+        # Wait for all D2H copies to finish
+        # before issuing gloo batch_isend_irecv operations.
+        if self._cuda_stream is not None:
+            self._cuda_stream.synchronize()
+        else:
+            torch.cuda.current_stream().synchronize()
+
+        reqs = batch_isend_irecv(p2p_ops)
+        for req in reqs:
+            req.wait()
+
+        if not recv_staging:
+            return
+        with torch.cuda.stream(self._cuda_stream):
+            for dst_tensor, cpu_tensor in recv_staging:
+                dst_tensor.copy_(cpu_tensor, non_blocking=True)
+
+
+class NixlEplbCommunicator(EplbCommunicator):
+    """EPLB communicator backed by NIXL READ transfers."""
+
+    def __init__(
+        self,
+        cpu_group: ProcessGroup,
+        all_expert_weights: Sequence[Sequence[torch.Tensor]],
+        expert_buffer: Sequence[torch.Tensor],
+        defer_remote_setup: bool = False,
+    ) -> None:
+        """Create a NIXL-backed EPLB communicator.
+
+        Args:
+            cpu_group: CPU process group for metadata exchange.
+            all_expert_weights: Expert weight tensors for all MoE layers.
+            expert_buffer: Pre-allocated receive buffer tensors.
+            defer_remote_setup: If True, postpone the collective
+                all-gather of NIXL agent metadata until the first
+                ``set_transfer_context`` call.  Required for elastic EP
+                where ranks join asynchronously and cannot participate
+                in collectives at construction time.
+        """
+        assert all_expert_weights, (
+            "NixlEplbCommunicator requires non-empty all_expert_weights."
+        )
+        assert expert_buffer, "NixlEplbCommunicator requires non-empty expert_buffer."
+        nixl_wrapper_cls = nixl_utils.NixlWrapper
+        if nixl_wrapper_cls is None:
+            raise RuntimeError("NIXL/ RIXL is unavailable.")
+
+        self._cpu_group = cpu_group
+        self._world_size = cpu_group.size()
+        self._rank = cpu_group.rank()
+
+        self._all_expert_weights = all_expert_weights
+        self._expert_buffer = expert_buffer
+        self._num_local_experts: int = all_expert_weights[0][0].shape[0]
+        self._device = all_expert_weights[0][0].device
+
+        for layer_tensors in all_expert_weights:
+            for tensor in layer_tensors:
+                assert is_weak_contiguous(tensor), (
+                    "Expert weight tensors must be contiguous in memory"
+                )
+                assert tensor.device == self._device, (
+                    "All local EPLB tensors are expected to be on the same "
+                    f"device: expected={self._device}, got={tensor.device}"
+                )
+        for tensor in expert_buffer:
+            assert is_weak_contiguous(tensor), (
+                "expert_buffer tensors must be contiguous in memory"
+            )
+
+        # (local_dlist, remote_dlist, xfer_handle) for in-flight READs;
+        # accumulated by add_recv, drained by execute.
+        self._xfer_entries: list[tuple[int, int, int]] = []
+        # Per-rank expert_id -> physical row; set by set_transfer_context.
+        self._expert_to_src_row: list[dict[int, int]] | None = None
+        self._layer_idx: int | None = None
+
+        nixl_agent_config = nixl_utils.nixl_agent_config
+        config = (
+            nixl_agent_config(capture_telemetry=False)
+            if nixl_agent_config is not None
+            else None
+        )
+        self._nixl_wrapper = nixl_wrapper_cls(self._make_agent_name(), config)
+        self._nixl_memory_type = "VRAM"
+        # NIXL registration handles; deregistered in __del__.
+        self._registered_descs: list[object] = []
+        self._remote_agents: dict[int, str] = {}
+        # peer -> (layer, tensor) -> (base_ptr, bytes_per_expert, dev_id).
+        self._remote_send_meta: dict[
+            int, dict[tuple[int, int], tuple[int, int, int]]
+        ] = {}
+
+        self._cuda_device_id = int(self._device.index or 0)
+        self._remote_state_initialized = False
+        self._init_step("buffers", self._init_registered_buffers)
+        if defer_remote_setup:
+            logger.info_once("NIXL EPLB: deferring remote agent setup (elastic EP).")
+        else:
+            self._init_remote_state()
+        self._log_initialized()
+
+    def _init_remote_state(self) -> None:
+        """Exchange NIXL agent metadata and RDMA pointer info with all peers.
+
+        This is a collective operation (uses ``all_gather_object`` twice).
+        Under elastic EP the call is deferred to the first
+        ``set_transfer_context`` invocation, where all ranks are
+        guaranteed to be synchronized.
+        """
+        self._init_step("agents", self._init_remote_agents)
+        self._init_step("send meta", self._exchange_remote_send_meta)
+        self._remote_state_initialized = True
+
+    def _ensure_remote_state(self) -> None:
+        if not self._remote_state_initialized:
+            self._init_remote_state()
+
+    @property
+    def needs_profile_buffer_reservation(self) -> bool:
+        return False
+
+    @staticmethod
+    def _init_step(name: str, fn: object, *args: object, **kwargs: object) -> None:
+        try:
+            fn(*args, **kwargs)  # type: ignore[operator]
+        except Exception as exc:
+            raise RuntimeError(f"NIXL EPLB init failed: {name}") from exc
+
+    def _make_agent_name(self) -> str:
+        """Build a deployment-unique nixl agent name."""
+        pp_size = get_pp_group().world_size
+        pp_suffix = f"-pp{get_pp_group().rank_in_group}" if pp_size > 1 else ""
+        uid = uuid.uuid4().hex[:8]
+        return f"eplb-{self._rank}{pp_suffix}-{uid}"
+
+    def set_stream(self, cuda_stream: torch.cuda.Stream | None) -> None:
+        pass
+
+    def add_send(
+        self,
+        tensors: list[torch.Tensor],
+        dst_rank: int,
+        expert_id: int,
+    ) -> None:
+        # No-op: NIXL READ is receiver-initiated. The sender's expert
+        # weights are pre-registered and always readable in-place.
+        pass
+
+    def set_transfer_context(self, old_indices: np.ndarray, layer_idx: int) -> None:
+        self._ensure_remote_state()
+        assert not self._xfer_entries, (
+            f"set_transfer_context() called with {len(self._xfer_entries)} "
+            f"pending transfers from layer {self._layer_idx}; "
+            f"execute() was not called after previous add_recv() calls"
+        )
+        self._layer_idx = layer_idx
+        n = self._num_local_experts
+        rank_experts = old_indices[: self._world_size * n].reshape(self._world_size, n)
+        self._expert_to_src_row = [
+            {int(eid): i for i, eid in enumerate(row) if eid != -1}
+            for row in rank_experts
+        ]
+
+    def add_recv(
+        self,
+        tensors: list[torch.Tensor],
+        src_rank: int,
+        expert_id: int,
+    ) -> None:
+        # Build NIXL descriptors and issue the RDMA READ immediately,
+        # overlapping the transfer with the remaining Python loop in
+        # move_to_buffer.
+        assert self._expert_to_src_row is not None and self._layer_idx is not None, (
+            "set_transfer_context() must be called before add_recv()"
+        )
+        src_row = self._expert_to_src_row[src_rank][expert_id]
+        layer_idx = self._layer_idx
+
+        local_descs: list[tuple[int, int, int]] = []
+        remote_descs: list[tuple[int, int, int]] = []
+        for t_idx, t in enumerate(tensors):
+            send_base, send_stride, remote_dev = self._remote_send_meta[src_rank][
+                (layer_idx, t_idx)
+            ]
+            assert t.nbytes == send_stride, (
+                f"tensor {t_idx} size {t.nbytes} != remote stride {send_stride}"
+            )
+            local_descs.append(
+                (
+                    t.data_ptr(),
+                    t.nbytes,
+                    self._cuda_device_id,
+                )
+            )
+            remote_descs.append(
+                (
+                    send_base + src_row * send_stride,
+                    send_stride,
+                    remote_dev,
+                )
+            )
+
+        local_h, remote_h, xfer_h = self._create_peer_xfer(
+            src_rank, local_descs, remote_descs
+        )
+        self._nixl_wrapper.transfer(xfer_h)
+        self._xfer_entries.append((local_h, remote_h, xfer_h))
+
+    def _init_remote_agents(self) -> None:
+        local_metadata = self._nixl_wrapper.get_agent_metadata()
+        gathered_metadata: list[bytes | None] = [None] * self._world_size
+        torch.distributed.all_gather_object(
+            gathered_metadata, local_metadata, group=self._cpu_group
+        )
+        for peer in range(self._world_size):
+            if peer == self._rank:
+                continue
+            peer_metadata = gathered_metadata[peer]
+            assert peer_metadata is not None
+            self._remote_agents[peer] = self._nixl_wrapper.add_remote_agent(
+                peer_metadata
+            )
+
+    def _init_registered_buffers(self) -> None:
+        all_tensors: list[torch.Tensor] = []
+        for layer_tensors in self._all_expert_weights:
+            all_tensors.extend(layer_tensors)
+        all_tensors.extend(self._expert_buffer)
+
+        descs = self._nixl_wrapper.get_reg_descs(all_tensors)
+        self._nixl_wrapper.register_memory(descs)
+        self._registered_descs.append(descs)
+
+    def _exchange_remote_send_meta(self) -> None:
+        """Exchange per-layer per-tensor metadata so receivers can compute
+        remote RDMA addresses at transfer time."""
+        local_meta: dict[tuple[int, int], tuple[int, int, int]] = {}
+        for layer_idx, layer_tensors in enumerate(self._all_expert_weights):
+            for t_idx, t in enumerate(layer_tensors):
+                nbytes_per_expert = t.nbytes // self._num_local_experts
+                local_meta[(layer_idx, t_idx)] = (
+                    t.data_ptr(),
+                    nbytes_per_expert,
+                    self._cuda_device_id,
+                )
+
+        # Per-rank map: (layer_idx, tensor_idx) -> (base_ptr, bytes_per_expert, dev_id).
+        # add_recv uses base_ptr + src_row * bytes_per_expert to compute
+        # the remote RDMA address for each expert.
+        gathered_meta: list[dict[tuple[int, int], tuple[int, int, int]] | None] = [
+            None
+        ] * self._world_size
+        torch.distributed.all_gather_object(
+            gathered_meta, local_meta, group=self._cpu_group
+        )
+
+        local_keys = set(local_meta.keys())
+        for peer in self._remote_agents:
+            peer_meta = gathered_meta[peer]
+            assert peer_meta is not None
+            peer_keys = set(peer_meta.keys())
+            if peer_keys != local_keys:
+                raise RuntimeError(
+                    f"NIXL EPLB metadata key mismatch with rank {peer}: "
+                    f"local={sorted(local_keys)}, peer={sorted(peer_keys)}"
+                )
+            for key in local_keys:
+                _, local_stride, _ = local_meta[key]
+                _, peer_stride, _ = peer_meta[key]
+                if local_stride != peer_stride:
+                    raise RuntimeError(
+                        f"NIXL EPLB nbytes_per_expert mismatch for {key} "
+                        f"with rank {peer}: "
+                        f"local={local_stride}, peer={peer_stride}"
+                    )
+            self._remote_send_meta[peer] = peer_meta
+
+    def _wait_for_all_transfers(self, handles: list[int]) -> None:
+        pending = set(handles)
+        while pending:
+            completed: list[int] = []
+            for handle in pending:
+                state = self._nixl_wrapper.check_xfer_state(handle)
+                if state == "DONE":
+                    completed.append(handle)
+                    continue
+                if state != "PROC":
+                    raise RuntimeError(f"NIXL transfer failed with state={state}")
+            for handle in completed:
+                pending.remove(handle)
+            if pending:
+                time.sleep(0.0005)
+
+    def _create_peer_xfer(
+        self,
+        src: int,
+        local_descs: list[tuple[int, int, int]],
+        remote_descs: list[tuple[int, int, int]],
+    ) -> tuple[int, int, int]:
+        """Create a batched xfer for multiple descriptors from one peer.
+
+        Each element in *local_descs* / *remote_descs* is an
+        ``(address, size, device_id)`` tuple.
+
+        Returns ``(local_dlist, remote_dlist, xfer_handle)``.
+        """
+        local_desc = self._nixl_wrapper.get_xfer_descs(
+            local_descs, self._nixl_memory_type
+        )
+        local_handle = self._nixl_wrapper.prep_xfer_dlist(
+            "NIXL_INIT_AGENT",
+            local_desc,
+        )
+
+        remote_desc = self._nixl_wrapper.get_xfer_descs(
+            remote_descs, self._nixl_memory_type
+        )
+        remote_handle = self._nixl_wrapper.prep_xfer_dlist(
+            self._remote_agents[src],
+            remote_desc,
+        )
+
+        indices = list(range(len(local_descs)))
+        xfer_handle = self._nixl_wrapper.make_prepped_xfer(
+            "READ",
+            local_handle,
+            indices,
+            remote_handle,
+            indices,
+        )
+        return (local_handle, remote_handle, xfer_handle)
+
+    def _post_read_barrier(self) -> None:
+        """Correctness fence: prevents overwrite-while-remote-read race.
+
+        We avoid ``torch.distributed.monitored_barrier`` because it
+        calls ``get_backend(group)`` which fails for stateless groups
+        (elastic EP).  An async ``all_reduce`` + ``wait(timeout)``
+        works with both regular and stateless groups and provides
+        equivalent timeout detection.
+        """
+        _dummy = torch.zeros(1, dtype=torch.int32)
+        work = torch.distributed.all_reduce(
+            _dummy, group=self._cpu_group, async_op=True
+        )
+        work.wait(timeout=timedelta(minutes=5))
+
+    def execute(self) -> None:
+        assert self._layer_idx is not None or not self._xfer_entries, (
+            "set_transfer_context() must be called before execute() "
+            "if any add_recv() calls were made"
+        )
+        try:
+            self._wait_for_all_transfers([x[2] for x in self._xfer_entries])
+
+            self._post_read_barrier()
+        finally:
+            for local_h, remote_h, xfer_h in self._xfer_entries:
+                with contextlib.suppress(Exception):
+                    self._nixl_wrapper.release_xfer_handle(xfer_h)
+                with contextlib.suppress(Exception):
+                    self._nixl_wrapper.release_dlist_handle(local_h)
+                with contextlib.suppress(Exception):
+                    self._nixl_wrapper.release_dlist_handle(remote_h)
+            self._xfer_entries.clear()
+            self._expert_to_src_row = None
+            self._layer_idx = None
+
+    def __del__(self) -> None:
+        with contextlib.suppress(Exception):
+            for local_h, remote_h, xfer_h in self._xfer_entries:
+                with contextlib.suppress(Exception):
+                    self._nixl_wrapper.release_xfer_handle(xfer_h)
+                with contextlib.suppress(Exception):
+                    self._nixl_wrapper.release_dlist_handle(local_h)
+                with contextlib.suppress(Exception):
+                    self._nixl_wrapper.release_dlist_handle(remote_h)
+        with contextlib.suppress(Exception):
+            for descs in self._registered_descs:
+                with contextlib.suppress(Exception):
+                    self._nixl_wrapper.deregister_memory(descs)
+            self._registered_descs.clear()
+        with contextlib.suppress(Exception):
+            for agent_name in self._remote_agents.values():
+                with contextlib.suppress(Exception):
+                    self._nixl_wrapper.remove_remote_agent(agent_name)
+            self._remote_agents.clear()
+
+
+class PyNcclEplbCommunicator(EplbCommunicator):
+    """EPLB communicator backed by PyNcclCommunicator using ncclSend/ncclRecv."""
+
+    def __init__(
+        self,
+        pynccl_comm: PyNcclCommunicator,
+        cuda_stream: torch.cuda.Stream | None = None,
+    ) -> None:
+        self._pynccl_comm = pynccl_comm
+        self._cuda_stream = cuda_stream
+        self._group_started = False
+        self._log_initialized()
+
+    def _ensure_group_started(self) -> None:
+        if not self._group_started:
+            self._pynccl_comm.group_start()
+            self._group_started = True
+
+    def add_send(
+        self,
+        tensors: list[torch.Tensor],
+        dst_rank: int,
+        expert_id: int,  # unused by this backend
+    ) -> None:
+        self._ensure_group_started()
+        for tensor in tensors:
+            self._pynccl_comm.send(tensor, dst_rank, stream=self._cuda_stream)
+
+    def add_recv(
+        self,
+        tensors: list[torch.Tensor],
+        src_rank: int,
+        expert_id: int,  # unused by this backend
+    ) -> None:
+        self._ensure_group_started()
+        for tensor in tensors:
+            self._pynccl_comm.recv(tensor, src_rank, stream=self._cuda_stream)
+
+    def execute(self) -> None:
+        if self._group_started:
+            self._pynccl_comm.group_end()
+            self._group_started = False
+
+
+def create_eplb_communicator(
+    group_coordinator: GroupCoordinator,
+    backend: str,
+    expert_weights: Sequence[Sequence[torch.Tensor]],
+    expert_buffer: Sequence[torch.Tensor],
+) -> EplbCommunicator:
+    """Create an EPLB communicator for the given backend.
+
+    Args:
+        group_coordinator: Process-group coordinator that provides the
+            device and CPU communication groups.
+        backend: Communicator backend name (``"torch_nccl"``,
+            ``"torch_gloo"``, ``"pynccl"``, or ``"nixl"``).
+            Falls back to ``"torch_nccl"`` when *None*.
+            Stateless (elastic EP) groups support ``"torch_nccl"``,
+            ``"pynccl"``, and ``"nixl"``; ``"torch_nccl"`` is silently
+            promoted to ``"pynccl"``.  ``"nixl"`` uses deferred remote
+            agent setup to avoid collective deadlocks during elastic
+            scaling.  When tensors reside on CPU, ``"torch_gloo"`` or
+            ``"torch_nccl"`` are used via the CPU process group.
+        expert_weights: Expert weight tensors for *all* MoE layers.
+            Shape ``(num_layers)(num_tensors_per_layer)``.
+            NixlEplbCommunicator registers all layers with NIXL for
+            zero-copy RDMA reads.
+        expert_buffer: Pre-allocated receive buffer tensors (one per
+            weight tensor in a single layer).
+    """
+    first_layer = expert_weights[0] if expert_weights else []
+    tensor_device_type = first_layer[0].device.type if first_layer else "cpu"
+    torch_group = (
+        group_coordinator.cpu_group
+        if tensor_device_type == "cpu"
+        else group_coordinator.device_group
+    )
+
+    def _create_pynccl() -> EplbCommunicator:
+        if tensor_device_type == "cpu":
+            raise RuntimeError(
+                "EPLB communicator 'pynccl' supports only cuda-like devices "
+                f"(got {tensor_device_type})."
+            )
+        unsupported_dtypes = sorted(
+            {
+                tensor.dtype
+                for tensor in first_layer
+                if not ncclDataTypeEnum.supports_torch_dtype(tensor.dtype)
+            },
+            key=str,
+        )
+        if unsupported_dtypes:
+            raise RuntimeError(
+                "EPLB communicator 'pynccl' requested but expert weights contain "
+                "unsupported dtypes: "
+                f"({', '.join(str(dtype) for dtype in unsupported_dtypes)})."
+            )
+
+        device_comm = group_coordinator.device_communicator
+        pynccl_comm = (
+            getattr(device_comm, "pynccl_comm", None)
+            if device_comm is not None
+            else None
+        )
+        if pynccl_comm is None or pynccl_comm.disabled or not pynccl_comm.available:
+            raise RuntimeError("EPLB communicator 'pynccl' requested but unavailable.")
+        try:
+            return PyNcclEplbCommunicator(pynccl_comm=pynccl_comm)
+        except Exception as exc:
+            raise RuntimeError(
+                f"Failed to initialize PyNcclEplbCommunicator ({exc})."
+            ) from exc
+
+    is_stateless = isinstance(group_coordinator, StatelessGroupCoordinator)
+    if is_stateless:
+        if backend == "nixl":
+            pass  # handled below with defer_remote_setup=True
+        elif backend not in ("torch_nccl", "pynccl"):
+            raise ValueError(
+                f"Elastic EP requires 'torch_nccl', 'pynccl', or 'nixl' "
+                f"EPLB communicator (got '{backend}')."
+            )
+        else:
+            if backend == "torch_nccl":
+                logger.warning(
+                    "Stateless elastic EP requires PyNCCL backend. "
+                    "Forcing EPLB communicator to 'pynccl'."
+                )
+                backend = "pynccl"
+            return _create_pynccl()
+
+    if backend == "nixl":
+        if not has_nixl():
+            raise RuntimeError(
+                "EPLB communicator 'nixl' requested but NIXL is unavailable."
+            )
+        if not (current_platform.is_cuda_alike() and tensor_device_type != "cpu"):
+            raise RuntimeError(
+                "EPLB communicator 'nixl' supports only cuda-like devices "
+                f"(got {tensor_device_type})."
+            )
+        try:
+            return NixlEplbCommunicator(
+                cpu_group=group_coordinator.cpu_group,
+                all_expert_weights=expert_weights,
+                expert_buffer=expert_buffer,
+                defer_remote_setup=is_stateless,
+            )
+        except Exception as exc:
+            raise RuntimeError(
+                f"Failed to initialize NixlEplbCommunicator ({exc})."
+            ) from exc
+    elif backend == "torch_gloo":
+        return TorchDistGlooStagedEplbCommunicator(
+            cpu_group=group_coordinator.cpu_group,
+        )
+    elif backend == "torch_nccl":
+        return TorchDistNcclEplbCommunicator(ep_group=torch_group)
+    elif backend == "pynccl":
+        return _create_pynccl()
+    raise ValueError(f"Unknown EPLB communicator backend: {backend}")

vllm/distributed/eplb/eplb_utils.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Utility functions for EPLB (Expert Parallel Load Balancing)."""
+
+import os
+import threading
+
+import torch
+
+from vllm.config import ParallelConfig
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class CpuGpuEvent:
+    """
+    Combines a CUDA event with a CPU threading event to enforce record->wait
+    ordering across two threads.
+
+    This class is designed for exactly two threads: one producer that calls
+    record() and one consumer that calls wait(). Using it with more than two
+    threads is not supported and will produce undefined behavior.
+
+    CUDA events alone are insufficient for cross-thread synchronization because
+    waiting on an unrecorded CUDA event is a no-op. The wait will return
+    immediately instead of blocking. This class adds a threading.Event so
+    that the waiting thread blocks on the CPU side until record() is called, at
+    which point the CUDA event is guaranteed to be in-flight and event.wait() will
+    correctly synchronize the GPU stream.
+    """
+
+    def __init__(self):
+        self._event = torch.cuda.Event()
+        self._recorded = threading.Event()
+
+    def wait(self, stream: torch.cuda.Stream | None = None):
+        """
+        Blocks the calling thread until record finishes. Used to guarantee that the
+        record kernel is called before wait.
+
+        Should only be called by the Async Eplb thread.
+        """
+        self._recorded.wait()
+        self._event.wait(stream)
+        self._recorded.clear()
+
+    def record(self, stream: torch.cuda.Stream | None = None):
+        """
+        Unblocks the waiting thread after calling event.record().
+
+        Should only be called by the main thread.
+        """
+        if self._recorded.is_set():
+            raise RuntimeError(
+                "CpuGpuEvent.record() called before the previous event was "
+                "consumed by wait()"
+            )
+        self._event = torch.cuda.Event()
+        self._event.record(stream)
+        self._recorded.set()
+
+
+def override_envs_for_eplb(
+    parallel_config: ParallelConfig,
+    moe_backend: str | None = None,
+) -> None:
+    """
+    Override environment variables for EPLB when specific conditions are met.
+
+    Args:
+        parallel_config: The parallel configuration object.
+        moe_backend: The configured MoE backend (e.g. ``deep_gemm_mega_moe``).
+    """
+    is_data_parallel = parallel_config.data_parallel_size > 1
+    is_eplb_enabled = parallel_config.enable_eplb
+    is_mega_moe = moe_backend == "deep_gemm_mega_moe"
+    is_nccl_based_eplb_communicator = parallel_config.eplb_config.communicator in (
+        "torch_nccl",
+        "pynccl",
+    )
+
+    # Override NCCL_MAX_CTAS to avoid hangs when EPLB's NCCL weight exchange
+    # contends with MoE backend's cooperative-launch on GPU SMs.
+    #
+    # DeepGEMM Mega MoE uses cooperative launch, which tries to reserve a
+    # large fraction of the GPU's SMs. If those SMs are occupied by NCCL,
+    # the cooperative launch blocks until enough SMs are freed, causing a
+    # deadlock. Limiting NCCL occupancy via NCCL_MAX_CTAS leaves space for
+    # the cooperative kernel to launch and complete.
+    if (
+        is_data_parallel
+        and is_eplb_enabled
+        and is_nccl_based_eplb_communicator
+        and is_mega_moe
+    ):
+        current_value_str = os.getenv("NCCL_MAX_CTAS")
+
+        if current_value_str and current_value_str.isdigit():
+            return
+
+        override_value = 8
+        os.environ["NCCL_MAX_CTAS"] = str(override_value)
+        logger.info_once(
+            f"EPLB: Setting NCCL_MAX_CTAS={override_value} "
+            f"for expert parallel with NCCL-based EPLB communicator and "
+            f"cooperative MoE backend (deep_gemm_mega_moe)",
+            scope="global",
+        )

vllm/distributed/eplb/policy/__init__.py

@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import get_args
+
+from vllm.config.parallel import EPLBPolicyOption
+
+from .abstract import AbstractEplbPolicy
+from .default import DefaultEplbPolicy
+
+EPLB_POLICIES = {"default": DefaultEplbPolicy}
+
+# Ensure that the EPLB_POLICIES keys match the EPLBPolicyOption values
+assert set(EPLB_POLICIES.keys()) == set(get_args(EPLBPolicyOption))
+
+__all__ = [
+    "AbstractEplbPolicy",
+    "DefaultEplbPolicy",
+    "EPLB_POLICIES",
+]

vllm/distributed/eplb/policy/abstract.py

@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from abc import ABC, abstractmethod
+
+import torch
+
+
+class AbstractEplbPolicy(ABC):
+    @classmethod
+    @abstractmethod
+    def rebalance_experts(
+        cls,
+        weight: torch.Tensor,
+        num_replicas: int,
+        num_groups: int,
+        num_nodes: int,
+        num_ranks: int,
+        old_global_expert_indices: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """
+        Entry point for expert-parallelism load balancer.
+
+        Parameters:
+            weight: [layers, num_logical_experts], the load statistics
+                for all logical experts
+            num_replicas: number of physical experts, must be a multiple of
+                `num_ranks`
+            num_groups: number of expert groups
+            num_nodes: number of server nodes
+            num_ranks: number of ranks, must be a multiple of `num_nodes`
+            old_global_expert_indices: [layers, num_logical_experts], the old global
+                expert indices. Used to avoid unnecessary weight copying
+                for experts moving within one rank.
+        Returns:
+            physical_to_logical_map: [layers, num_replicas], the expert
+                index of each replica
+        """
+        raise NotImplementedError

vllm/distributed/eplb/policy/default.py

@@ -0,0 +1,332 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Expert parallelism load balancer (EPLB) for vLLM.
+
+This module implements the core rearrangement algorithm.
+
+The rearrangement algorithm is adapted from
+[DeepSeek EPLB](https://github.com/deepseek-ai/eplb).
+
+Please find at [#12](https://github.com/deepseek-ai/EPLB/issues/12) an example
+on how the EPLB algorithm works.
+"""
+
+import numpy as np
+import torch
+
+from .abstract import AbstractEplbPolicy
+
+
+class DefaultEplbPolicy(AbstractEplbPolicy):
+    @classmethod
+    def balanced_packing(
+        cls, weight: np.ndarray, num_packs: int
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """
+        Pack n weighted objects to m packs, such that each bin contains exactly
+        n/m objects and the weights of all packs are as balanced as possible.
+
+        Parameters:
+            weight: [X, n], the weight of each item
+            num_packs: number of packs
+
+        Returns:
+            pack_index: [X, n], the pack index of each item
+            rank_in_pack: [X, n], the rank of the item in the pack
+        """
+        num_layers, num_groups = weight.shape
+        assert num_groups % num_packs == 0
+        groups_per_pack = num_groups // num_packs
+
+        if groups_per_pack == 1:
+            pack_index = np.tile(np.arange(num_groups, dtype=np.int64), (num_layers, 1))
+            rank_in_pack = np.zeros_like(pack_index, dtype=np.int64)
+            return pack_index, rank_in_pack
+
+        # Sort and get indices in descending order
+        indices = np.argsort(-weight, axis=-1)
+
+        pack_index = np.full((num_layers, num_groups), -1, dtype=np.int64)
+        rank_in_pack = np.full((num_layers, num_groups), -1, dtype=np.int64)
+
+        pack_weights = np.zeros((num_layers, num_packs), dtype=np.float64)
+        pack_items = np.zeros((num_layers, num_packs), dtype=np.int64)
+
+        # Run the packing algorithm
+        for layer_idx in range(num_layers):
+            weights_row = pack_weights[layer_idx]
+            items_row = pack_items[layer_idx]
+
+            for group in indices[layer_idx]:
+                # Pick the lightest pack; full packs are masked out by inf.
+                pack = int(np.argmin(weights_row))
+
+                pack_index[layer_idx, group] = pack
+                rank_in_pack[layer_idx, group] = items_row[pack]
+                weights_row[pack] += weight[layer_idx, group]
+                items_row[pack] += 1
+                if items_row[pack] == groups_per_pack:
+                    # Mark as unavailable for future selections.
+                    weights_row[pack] = np.inf
+
+        return pack_index, rank_in_pack
+
+    @classmethod
+    def replicate_experts(
+        cls, weight: np.ndarray, num_phy: int
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """
+        Replicate `num_log` experts to `num_phy` replicas, such that the maximum
+        load of all replicas is minimized.
+
+        Parameters:
+            weight: [X, num_log]
+            num_phy: total number of experts after replication
+
+        Returns:
+            phy2log: [X, num_phy], logical expert id of each physical expert
+            logcnt: [X, num_log], number of replicas for each logical expert
+        """
+        n, num_log = weight.shape
+        num_redundant = num_phy - num_log
+        assert num_redundant >= 0
+        phy2log = np.tile(np.arange(num_phy, dtype=np.int64), (n, 1))
+        logcnt = np.ones((n, num_log), dtype=np.int64)
+        arangen = np.arange(n, dtype=np.int64)
+        for i in range(num_log, num_phy):
+            redundant_indices = np.argmax(weight / logcnt, axis=-1)
+            phy2log[:, i] = redundant_indices
+            logcnt[arangen, redundant_indices] += 1
+        return phy2log, logcnt
+
+    @classmethod
+    def rebalance_experts_hierarchical(
+        cls,
+        weight: np.ndarray,
+        num_physical_experts: int,
+        num_groups: int,
+        num_nodes: int,
+        num_gpus: int,
+    ) -> np.ndarray:
+        """
+        Parameters:
+            weight: [num_moe_layers, num_logical_experts]
+            num_physical_experts: number of physical experts after replication
+            num_groups: number of expert groups
+            num_nodes: number of server nodes, where the intra-node network
+                (e.g, NVLink) is faster
+            num_gpus: number of GPUs, must be a multiple of `num_nodes`
+
+        Returns:
+            phy2log: [layers, num_replicas], the expert
+                index of each replica
+        """
+        num_layers, num_logical_experts = weight.shape
+        assert num_logical_experts % num_groups == 0
+        group_size = num_logical_experts // num_groups
+        assert num_groups % num_nodes == 0
+        groups_per_node = num_groups // num_nodes
+        assert num_gpus % num_nodes == 0
+        assert num_physical_experts % num_gpus == 0
+        phy_experts_per_gpu = num_physical_experts // num_gpus
+
+        def inverse(perm: np.ndarray) -> np.ndarray:
+            inv = np.empty_like(perm)
+            row_idx = np.arange(perm.shape[0])[:, None]
+            col_idx = np.arange(perm.shape[1], dtype=np.int64)
+            inv[row_idx, perm] = col_idx
+            return inv
+
+        # Step 1: pack groups to nodes
+        tokens_per_group = weight.reshape(num_layers, num_groups, group_size).sum(
+            axis=-1
+        )
+        group_pack_index, group_rank_in_pack = cls.balanced_packing(
+            tokens_per_group, num_nodes
+        )
+        # Map each logical expert into a node-local ordering based on packed groups.
+        log2mlog = (
+            (
+                (group_pack_index * groups_per_node + group_rank_in_pack)[..., None]
+                * group_size
+            )
+            + np.arange(group_size, dtype=np.int64)
+        ).reshape(num_layers, num_logical_experts)
+        mlog2log = inverse(log2mlog)
+
+        # Step 2: construct redundant experts within nodes
+        # Reorder weights into the node-local layout so replication is done per node.
+        tokens_per_mlog = np.take_along_axis(weight, mlog2log, axis=1).reshape(
+            -1, num_logical_experts // num_nodes
+        )
+        phy2mlog, mlogcnt = cls.replicate_experts(
+            tokens_per_mlog, num_physical_experts // num_nodes
+        )
+
+        # Step 3: pack physical_experts to GPUs
+        # Effective per-physical load = logical load divided by replica count.
+        tokens_per_phy = np.take_along_axis(tokens_per_mlog / mlogcnt, phy2mlog, axis=1)
+        pack_index, rank_in_pack = cls.balanced_packing(
+            tokens_per_phy, num_gpus // num_nodes
+        )
+        phy2pphy = pack_index * phy_experts_per_gpu + rank_in_pack
+        pphy2phy = inverse(phy2pphy)
+
+        # Reorder node-local logical indices into the post-packing physical order.
+        pphy2mlog = np.take_along_axis(phy2mlog, pphy2phy, axis=1)
+        pphy2mlog = (
+            pphy2mlog.reshape(num_layers, num_nodes, -1)
+            + np.arange(
+                0,
+                num_logical_experts,
+                num_logical_experts // num_nodes,
+                dtype=np.int64,
+            )[None, :, None]
+        ).reshape(num_layers, -1)
+        # Map node-local logical indices back to global logical expert ids.
+        pphy2log = np.take_along_axis(mlog2log, pphy2mlog, axis=1)
+        return pphy2log
+
+    @classmethod
+    def preserve_intragpu_slots(
+        cls,
+        phy2log: np.ndarray,
+        num_ranks: int,
+        old_phy2log: np.ndarray,
+    ) -> np.ndarray:
+        """
+        Reorder the new mapping per GPU so that experts that remain on the same GPU
+        keep their previous slot positions when possible. Incoming experts to that GPU
+        fill any remaining available slots. This is applied only when the number of GPUs
+        is unchanged and the slots per GPU remain the same between
+        the old and new mappings.
+        """
+        num_phy_experts = phy2log.shape[1]
+        if num_ranks <= 0 or num_phy_experts % num_ranks != 0:
+            return phy2log
+
+        # Move to CPU and convert to NumPy for processing
+        slots_per_gpu = num_phy_experts // num_ranks
+        num_layers = phy2log.shape[0]
+
+        post_phy2log = phy2log.copy()
+
+        for gpu_idx in range(num_ranks):
+            start = gpu_idx * slots_per_gpu
+            end = start + slots_per_gpu
+            # Experts across all layers for this GPU
+            old_local = old_phy2log[:, start:end]  # [layers, slots]
+            new_local = phy2log[:, start:end]  # [layers, slots]
+
+            used_new_indices = np.zeros((num_layers, slots_per_gpu), dtype=bool)
+            preserved_positions = np.zeros((num_layers, slots_per_gpu), dtype=bool)
+
+            # First pass: preserve same-logical experts in their previous slots
+            for slot_idx in range(slots_per_gpu):
+                # matches: [layers, slots], True where new local experts have
+                # the same logical value as the old from 'slot_idx' and not checked yet
+                matches = (new_local == old_local[:, slot_idx][:, None]) & (
+                    ~used_new_indices
+                )
+                has_any = matches.any(axis=1)
+                if np.any(has_any):
+                    first_idx = np.argmax(matches, axis=1)
+                    layer_indices = np.nonzero(has_any)[0]
+                    matched_new_positions = first_idx[layer_indices]
+                    post_phy2log[layer_indices, start + slot_idx] = new_local[
+                        layer_indices, matched_new_positions
+                    ]
+                    used_new_indices[layer_indices, matched_new_positions] = True
+                    preserved_positions[layer_indices, slot_idx] = True
+
+            # Second pass: fill remaining slots with remaining new experts
+            remaining_mask = ~used_new_indices  # [layers, slots]
+            fill_mask = ~preserved_positions  # [layers, slots]
+            if remaining_mask.any() and fill_mask.any():
+                idx_base = np.tile(np.arange(slots_per_gpu), (num_layers, 1))
+                # Sentinel value for unavailable positions.
+                large = slots_per_gpu + 1
+                # Priorities: keep original index for available spots, set sentinel
+                # for unavailable; lower is earlier.
+                remaining_priority = np.where(remaining_mask, idx_base, large)
+                fill_priority = np.where(fill_mask, idx_base, large)
+                # Sort to get ordered indices of available src/dst positions per layer.
+                remaining_indices = np.argsort(remaining_priority, axis=1)
+                fill_indices = np.argsort(fill_priority, axis=1)
+                # Fill count per layer (cannot exceed either side).
+                remaining_counts = remaining_mask.sum(axis=1)
+                fill_counts = fill_mask.sum(axis=1)
+                take_counts = np.minimum(remaining_counts, fill_counts)
+                # Assign remaining new experts to remaining slots per layer.
+                for layer_idx in range(num_layers):
+                    k = int(take_counts[layer_idx])
+                    if k <= 0:
+                        continue
+                    src_pos = remaining_indices[layer_idx, :k]
+                    dst_pos = fill_indices[layer_idx, :k]
+                    post_phy2log[layer_idx, start + dst_pos] = new_local[
+                        layer_idx, src_pos
+                    ]
+
+        return post_phy2log
+
+    @classmethod
+    def rebalance_experts(
+        cls,
+        weight: torch.Tensor,
+        num_replicas: int,
+        num_groups: int,
+        num_nodes: int,
+        num_ranks: int,
+        old_global_expert_indices: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """
+        Entry point for expert-parallelism load balancer.
+
+        Parameters:
+            weight: [layers, num_logical_experts], the load statistics for all
+                logical experts
+            num_replicas: number of physical experts, must be a multiple of
+                `num_gpus`
+            num_groups: number of expert groups
+            num_nodes: number of server nodes, where the intra-node network
+                (e.g, NVLink) is faster
+            num_ranks: number of ranks, must be a multiple of `num_nodes`
+            old_global_expert_indices: [layers, num_logical_experts], the old global
+                expert indices. Used to avoid unnecessary weight copying
+                for experts moving within one rank.
+        Returns:
+            phy2log: [layers, num_replicas], the expert
+                index of each replica
+        """
+        weight_np = weight.float().cpu().numpy()
+        old_phy2log_np = (
+            old_global_expert_indices.cpu().numpy()
+            if old_global_expert_indices is not None
+            else None
+        )
+
+        if num_groups % num_nodes == 0:
+            # use hierarchical load-balance policy
+            phy2log_np = cls.rebalance_experts_hierarchical(
+                weight_np, num_replicas, num_groups, num_nodes, num_ranks
+            )
+        else:
+            # use global load-balance policy
+            phy2log_np = cls.rebalance_experts_hierarchical(
+                weight_np, num_replicas, 1, 1, num_ranks
+            )
+
+        # Optional postprocessing to preserve slots for experts moving
+        # within the same GPU
+        # Only apply when the number of GPUs and slots per GPU remain unchanged.
+        # Helps to avoid unnecessary weight copying when experts move
+        # within the same GPU.
+        if old_phy2log_np is not None:
+            phy2log_np = cls.preserve_intragpu_slots(
+                phy2log_np, num_ranks, old_phy2log_np
+            )
+
+        phy2log = torch.from_numpy(phy2log_np)
+        return phy2log

vllm/distributed/eplb/rebalance_execute.py

@@ -0,0 +1,706 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+The actual execution of the rearrangement.
+
+This involves the exchange of expert weights between GPUs.
+"""
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+from torch.distributed import ProcessGroup, all_gather
+
+from vllm.distributed.eplb.eplb_communicator import EplbCommunicator
+from vllm.distributed.eplb.eplb_utils import CpuGpuEvent
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class TransferMetadata:
+    """Metadata describing a completed EPLB buffer transfer."""
+
+    is_unchanged: np.ndarray
+    """Mask of (num_local_experts,) indicating experts unchanged after rebalance."""
+    is_received_locally: np.ndarray
+    """Mask of (num_local_experts,) indicating experts received from local data."""
+    recv_primary_mask: np.ndarray
+    """Mask of (num_local_experts,) indicating primary experts received."""
+    recv_count: int
+    """Number of received experts for the layer."""
+    recv_expert_ids: np.ndarray
+    """Expert ids (num_local_experts,) of remote primary experts."""
+    recv_dst_rows: np.ndarray
+    """Target expert indices (num_local_experts,) in local tensors to send."""
+
+
+@dataclass
+class AsyncEplbLayerResult:
+    """
+    The result of one completed async EPLB layer transfer.
+    """
+
+    layer_idx: int
+    """Index of the MoE layer that was transferred."""
+    new_physical_to_logical_map: torch.Tensor
+    """
+    New physical→logical mapping for layers_idx, on CPU.
+    Shape: (num_physical_experts)
+    """
+    transfer_metadata: TransferMetadata
+    """Metadata describing what was received during transfer_layer."""
+    consumed_event: CpuGpuEvent
+    """
+    Event used to synchronize access to the intermediate buffer. The async worker calls
+    wait() after it finishes transferring weights to the intermediate buffer. The main
+    thread calls record() after it finishes transferring weights out of the intermediate
+    buffer in _move_to_workspace()
+    """
+
+
+def get_ep_ranks_with_experts_batch(
+    expert_ids: np.ndarray,
+    num_local_experts: int,
+    old_indices: np.ndarray,
+    new_indices: np.ndarray,
+) -> tuple[dict[int, list[int]], dict[int, list[int]]]:
+    """
+    Get the ranks of the experts that need to be exchanged.
+
+    Args:
+        expert_ids: 1D array of expert indices to query.
+        num_local_experts: The number of local experts.
+        old_indices: The old indices of the experts.
+        new_indices: The new indices of the experts.
+
+    Returns:
+        A tuple of two dictionaries mapping expert_id to:
+        - ranks_to_send: The ranks that have this expert and need to send.
+        - ranks_to_recv: The ranks that need to receive this expert.
+    """
+    ranks_to_send_map: dict[int, list[int]] = {}
+    ranks_to_recv_map: dict[int, list[int]] = {}
+
+    # Fast path: if no experts, return empty dicts
+    if expert_ids.size == 0:
+        return ranks_to_send_map, ranks_to_recv_map
+
+    unique_experts = np.unique(expert_ids)
+    num_positions = len(old_indices)
+    position_indices = np.arange(num_positions, dtype=np.int32)
+
+    # Vectorized approach: find all positions matching any query expert in one pass
+    # Use np.isin to get boolean masks for all relevant positions at once
+    old_relevant_mask = np.isin(old_indices, unique_experts)
+    new_relevant_mask = np.isin(new_indices, unique_experts)
+
+    # Process old_indices (send ranks)
+    if np.any(old_relevant_mask):
+        old_relevant_positions = position_indices[old_relevant_mask]
+        old_relevant_experts = old_indices[old_relevant_mask]
+        old_relevant_ranks = old_relevant_positions // num_local_experts
+
+        # Sort by expert first, then by position (to maintain first-appearance order)
+        sort_order = np.lexsort((old_relevant_positions, old_relevant_experts))
+        sorted_experts = old_relevant_experts[sort_order]
+        sorted_ranks = old_relevant_ranks[sort_order]
+
+        # Find boundaries where expert changes
+        expert_boundaries = np.concatenate(
+            [[0], np.where(np.diff(sorted_experts) != 0)[0] + 1, [len(sorted_experts)]]
+        )
+
+        # For each expert, extract unique ranks in order of first appearance
+        for i in range(len(expert_boundaries) - 1):
+            start, end = expert_boundaries[i], expert_boundaries[i + 1]
+            expert = int(sorted_experts[start])
+            expert_ranks = sorted_ranks[start:end]
+
+            # Get unique ranks preserving order
+            _, unique_idx = np.unique(expert_ranks, return_index=True)
+            unique_ranks = expert_ranks[np.sort(unique_idx)]
+            ranks_to_send_map[expert] = unique_ranks.tolist()
+
+    # Process new_indices (recv ranks)
+    if np.any(new_relevant_mask):
+        new_relevant_positions = position_indices[new_relevant_mask]
+        new_relevant_experts = new_indices[new_relevant_mask]
+        new_relevant_ranks = new_relevant_positions // num_local_experts
+
+        # Sort by expert first, then by position
+        sort_order = np.lexsort((new_relevant_positions, new_relevant_experts))
+        sorted_experts = new_relevant_experts[sort_order]
+        sorted_ranks = new_relevant_ranks[sort_order]
+
+        # Find boundaries where expert changes
+        expert_boundaries = np.concatenate(
+            [[0], np.where(np.diff(sorted_experts) != 0)[0] + 1, [len(sorted_experts)]]
+        )
+
+        # For each expert, extract unique ranks and exclude local copies
+        for i in range(len(expert_boundaries) - 1):
+            start, end = expert_boundaries[i], expert_boundaries[i + 1]
+            expert = int(sorted_experts[start])
+            expert_ranks = sorted_ranks[start:end]
+
+            # Get unique ranks preserving order
+            _, unique_idx = np.unique(expert_ranks, return_index=True)
+            unique_ranks = expert_ranks[np.sort(unique_idx)]
+
+            # Remove ranks that have local copies (in send map)
+            send_ranks_set = set(ranks_to_send_map.get(expert, []))
+            recv_ranks_actual = [
+                int(r) for r in unique_ranks if r not in send_ranks_set
+            ]
+            ranks_to_recv_map[expert] = recv_ranks_actual
+
+    # Handle experts that only appear in old (send only) or new (recv only)
+    for expert in unique_experts:
+        expert = int(expert)
+        if expert not in ranks_to_send_map:
+            ranks_to_send_map[expert] = []
+        if expert not in ranks_to_recv_map:
+            ranks_to_recv_map[expert] = []
+
+    return ranks_to_send_map, ranks_to_recv_map
+
+
+def move_to_buffer(
+    num_local_experts: int,
+    old_indices: np.ndarray,
+    new_indices: np.ndarray,
+    expert_weights: Sequence[torch.Tensor],
+    expert_weights_buffers: Sequence[torch.Tensor],
+    cuda_stream: torch.cuda.Stream | None,
+    ep_rank: int,
+    communicator: EplbCommunicator,
+    layer_idx: int = 0,
+) -> TransferMetadata:
+    """
+    Rearranges expert weights during EPLB rebalancing.
+
+    Args:
+        num_local_experts: Number of local experts.
+        old_indices: (num_experts_total,) ndarray of current (old)
+            global-to-local expert assignments.
+        new_indices: (num_experts_total,) ndarray of desired (new)
+            global-to-local assignments after rebalance.
+        expert_weights: Original expert weights for the layer.
+        expert_weights_buffers: Intermediate buffers (one per tensor).
+        cuda_stream: CUDA stream for async copies (can be None for sync mode).
+        ep_rank: Rank of this process in expert parallel group.
+        communicator: EplbCommunicator instance for P2P communication.
+        layer_idx: Index of the MoE layer being transferred.
+
+    Returns:
+        TransferMetadata: Metadata needed for completing remote weight transfers.
+    """
+    assert old_indices.shape == new_indices.shape
+    recv_primary_mask = np.zeros((num_local_experts,), dtype=np.bool_)
+    send_expert_ids = np.full((num_local_experts,), -1, dtype=np.int64)
+    send_src_rows = np.full((num_local_experts,), -1, dtype=np.int32)
+    recv_expert_ids = np.full((num_local_experts,), -1, dtype=np.int64)
+    recv_dst_rows = np.full((num_local_experts,), -1, dtype=np.int32)
+
+    base = ep_rank * num_local_experts
+    local_rows = np.arange(num_local_experts, dtype=np.int32)
+    local_global = base + local_rows
+
+    old_local_expert_ids = old_indices[local_global]
+    new_local_expert_ids = new_indices[local_global]
+
+    # Unchanged mask
+    is_unchanged = old_local_expert_ids == new_local_expert_ids
+
+    # Local receive eligibility
+    new_valid = new_local_expert_ids != -1
+    can_recv_local = np.isin(
+        new_local_expert_ids, old_local_expert_ids, assume_unique=False
+    )
+    is_received_locally = np.logical_or(
+        is_unchanged, np.logical_and(new_valid, can_recv_local)
+    )
+
+    # Send map: first src row per unique expert present locally in old mapping
+    send_count = 0
+    valid_old = old_local_expert_ids != -1
+    if np.any(valid_old):
+        uniq_experts, first_idx = np.unique(
+            old_local_expert_ids[valid_old], return_index=True
+        )
+        filtered_rows = local_rows[valid_old]
+        src_rows = filtered_rows[first_idx]
+        send_count = int(uniq_experts.shape[0])
+        send_expert_ids[:send_count] = uniq_experts
+        send_src_rows[:send_count] = src_rows
+
+    # Recv map: primary dst per unique expert needed remotely
+    recv_count = 0
+    need_recv_mask = np.logical_and(~is_received_locally, new_valid)
+    if np.any(need_recv_mask):
+        desired_experts = new_local_expert_ids[need_recv_mask]
+        desired_dsts = local_rows[need_recv_mask]
+        uniq_recv_experts, uniq_indices = np.unique(desired_experts, return_index=True)
+        dst_rows = desired_dsts[uniq_indices]
+        recv_count = int(uniq_recv_experts.shape[0])
+        recv_expert_ids[:recv_count] = uniq_recv_experts
+        recv_dst_rows[:recv_count] = dst_rows
+        recv_primary_mask[dst_rows] = True
+
+    eligible_local_buffer_mask = np.logical_and(~is_unchanged, is_received_locally)
+
+    # 1. Local moves into tmp buffers
+    if bool(eligible_local_buffer_mask.any()) and send_count > 0:
+        dest_indices = np.nonzero(eligible_local_buffer_mask)[0].tolist()
+        expert_to_src_map = dict(
+            zip(send_expert_ids[:send_count], send_src_rows[:send_count])
+        )
+        for dst in dest_indices:
+            expert = new_local_expert_ids[dst]
+            src_local = expert_to_src_map.get(expert, -1)
+            if src_local != -1:
+                with torch.cuda.stream(cuda_stream):
+                    for w, b in zip(expert_weights, expert_weights_buffers):
+                        b[dst].copy_(w[src_local], non_blocking=True)
+
+    communicator.set_transfer_context(old_indices, layer_idx)
+
+    # 2. Post sends
+    if send_count > 0:
+        experts = send_expert_ids[:send_count]
+        srcs = send_src_rows[:send_count]
+        order = np.argsort(experts, kind="stable")
+        experts = experts[order]
+        srcs = srcs[order]
+
+        send_map, recv_map = get_ep_ranks_with_experts_batch(
+            experts,
+            num_local_experts,
+            old_indices,
+            new_indices,
+        )
+
+        for expert, src in zip(experts.tolist(), srcs.tolist()):
+            ranks_to_send = send_map[expert]
+            ranks_to_recv = recv_map[expert]
+            if not ranks_to_send or not ranks_to_recv:
+                continue
+            num_dst_per_sender = len(ranks_to_recv) // len(ranks_to_send)
+            sender_pos = ranks_to_send.index(ep_rank)
+            recv_begin = sender_pos * num_dst_per_sender
+            recv_end = recv_begin + num_dst_per_sender
+            recv_ranks = ranks_to_recv[recv_begin:recv_end]
+            remainder_start = len(ranks_to_send) * num_dst_per_sender
+            recver_pos = remainder_start + sender_pos
+            if recver_pos < len(ranks_to_recv):
+                recv_ranks.append(ranks_to_recv[recver_pos])
+            expert_tensors = [w[src] for w in expert_weights]
+            for dst in recv_ranks:
+                communicator.add_send(expert_tensors, dst, expert_id=int(expert))
+
+    # 3. Post recvs
+    if recv_count > 0:
+        experts = recv_expert_ids[:recv_count]
+        dsts = recv_dst_rows[:recv_count]
+        order = np.argsort(experts, kind="stable")
+        experts = experts[order]
+        dsts = dsts[order]
+
+        send_map, recv_map = get_ep_ranks_with_experts_batch(
+            experts,
+            num_local_experts,
+            old_indices,
+            new_indices,
+        )
+
+        for expert, dst in zip(experts.tolist(), dsts.tolist()):
+            ranks_to_send = send_map[expert]
+            ranks_to_recv = recv_map[expert]
+            if not ranks_to_send or not ranks_to_recv:
+                continue
+            num_dst_per_sender = len(ranks_to_recv) // len(ranks_to_send)
+            recver_pos = ranks_to_recv.index(ep_rank)
+            remainder_start = len(ranks_to_send) * num_dst_per_sender
+            if recver_pos < remainder_start:
+                src = ranks_to_send[recver_pos // num_dst_per_sender]
+            else:
+                src = ranks_to_send[recver_pos - remainder_start]
+            communicator.add_recv(
+                [b[dst] for b in expert_weights_buffers],
+                src,
+                expert_id=int(expert),
+            )
+
+    # 4. Execute transfers and wait for completion.
+    communicator.execute()
+    return TransferMetadata(
+        is_unchanged=is_unchanged,
+        is_received_locally=is_received_locally,
+        recv_primary_mask=recv_primary_mask,
+        recv_count=recv_count,
+        recv_expert_ids=recv_expert_ids,
+        recv_dst_rows=recv_dst_rows,
+    )
+
+
+def move_from_buffer(
+    expert_weights: Sequence[torch.Tensor],
+    expert_weights_buffers: list[torch.Tensor],
+    transfer_metadata: TransferMetadata,
+    new_indices: np.ndarray,
+    ep_rank: int,
+) -> None:
+    """
+    Copies expert weights from communication buffers back to the target weight tensors
+    after EPLB rebalancing.
+
+    Args:
+        expert_weights: List of the actual MoE layer weights used in the execution.
+        expert_weights_buffers: Intermediate buffers containing the experts weights
+            after the transfer is completed.
+        transfer_metadata: TransferMetadata containing transfer metadata.
+        new_indices: (num_experts_total,) mapping from local rows to desired
+            (possibly global) expert id, after rebalance.
+        ep_rank: Rank of the process in the expert parallel group.
+    """
+    is_unchanged = transfer_metadata.is_unchanged
+    is_received_locally = transfer_metadata.is_received_locally
+    recv_primary_mask = transfer_metadata.recv_primary_mask
+    recv_count = transfer_metadata.recv_count
+    recv_expert_ids = transfer_metadata.recv_expert_ids
+    recv_dst_rows = transfer_metadata.recv_dst_rows
+    num_local_experts = is_unchanged.shape[0]
+
+    # Mask for rows to copy back from buffers:
+    # copy if locally received OR remote primary recv
+    copy_mask = np.logical_or(is_received_locally, recv_primary_mask)
+    dest_mask_np = np.logical_and(~is_unchanged, copy_mask)
+    if bool(dest_mask_np.any()):
+        dest_indices = np.nonzero(dest_mask_np)[0].tolist()
+        for dst in dest_indices:
+            for w, b in zip(expert_weights, expert_weights_buffers):
+                w[dst].copy_(b[dst], non_blocking=True)
+
+    if recv_count == 0:
+        return
+
+    # Duplicate remote received rows to non-primary duplicate dsts
+    base = ep_rank * num_local_experts
+    local_experts = new_indices[base + np.arange(num_local_experts, dtype=np.int32)]
+    duplicate_mask = np.logical_and(
+        np.logical_and(~is_unchanged, ~is_received_locally),
+        np.logical_and(~recv_primary_mask, local_experts != -1),
+    )
+    # All received experts are unique in the destination, so no need to copy duplicates
+    if not bool(duplicate_mask.any()):
+        return
+
+    dup_dst_rows = np.nonzero(duplicate_mask)[0]
+    dup_experts = local_experts[dup_dst_rows]
+
+    prim_experts = recv_expert_ids[:recv_count]
+    prim_dsts = recv_dst_rows[:recv_count]
+    order = np.argsort(prim_experts, kind="stable")
+    prim_experts_sorted = prim_experts[order]
+    prim_dsts_sorted = prim_dsts[order]
+    pos = np.searchsorted(prim_experts_sorted, dup_experts)
+    valid = np.logical_and(
+        pos < prim_experts_sorted.shape[0],
+        prim_experts_sorted[np.minimum(pos, prim_experts_sorted.shape[0] - 1)]
+        == dup_experts,
+    )
+    if not bool(valid.any()):
+        return
+
+    matched_dst_rows = dup_dst_rows[valid]
+    matched_src_rows = prim_dsts_sorted[pos[valid]]
+
+    for dst, src in zip(matched_dst_rows.tolist(), matched_src_rows.tolist()):
+        for w in expert_weights:
+            w[dst].copy_(w[src], non_blocking=True)
+
+
+def transfer_layer(
+    old_layer_indices: torch.Tensor,
+    new_layer_indices: torch.Tensor,
+    expert_weights: Sequence[torch.Tensor],
+    expert_weights_buffer: Sequence[torch.Tensor],
+    ep_group: ProcessGroup,
+    communicator: EplbCommunicator,
+    is_profile: bool = False,
+    cuda_stream: torch.cuda.Stream | None = None,
+    rank_mapping: dict[int, int] | None = None,
+    layer_idx: int = 0,
+) -> TransferMetadata:
+    """
+    Rearranges the expert weights in place according to the new expert indices.
+
+    The value of the indices arguments are logical indices of the experts,
+    while keys are physical.
+
+    Args:
+        old_layer_indices: Shape (num_physical_experts,).
+        new_layer_indices: Shape (num_physical_experts,).
+        expert_weights: Iterable of weight tensors for this layer, each with shape
+            (num_local_physical_experts, hidden_size_i).
+            For example, a linear layer may have up and down projection.
+        expert_weights_buffer: Intermediate buffers (one per weight tensor).
+        ep_group: The device process group for expert parallelism.
+        communicator: EplbCommunicator instance for P2P communication.
+        is_profile (bool): If `True`, do not perform any actual weight copy.
+            This is used during profile run, where we only perform dummy
+            communications to reserve enough memory for the buffers.
+        cuda_stream: CUDA stream for async copies (can be None for sync mode).
+        rank_mapping: Optional rank mapping for elastic expert parallelism.
+        layer_idx: Index of the MoE layer being transferred.
+
+    Returns:
+        TransferMetadata: Metadata needed for completing remote weight transfers,
+            including is_unchanged and is_received_locally masks.
+    """
+    ep_size = ep_group.size()
+    if rank_mapping is not None:
+        # Add a layer dimension for compatibility with mapping functions
+        old_layer_indices_2d = old_layer_indices.unsqueeze(0)
+        new_layer_indices_2d = new_layer_indices.unsqueeze(0)
+
+        if len(rank_mapping) == ep_group.size():
+            # scale down
+            new_layer_indices_2d = _map_new_expert_indices_with_rank_mapping(
+                new_layer_indices_2d,
+                rank_mapping,
+            )
+        else:
+            # scale up
+            old_layer_indices_2d = _map_old_expert_indices_with_rank_mapping(
+                old_layer_indices_2d,
+                rank_mapping,
+                ep_group.size(),
+            )
+
+        # Remove the layer dimension
+        old_layer_indices = old_layer_indices_2d.squeeze(0)
+        new_layer_indices = new_layer_indices_2d.squeeze(0)
+
+    assert old_layer_indices.shape == new_layer_indices.shape
+    num_physical_experts = old_layer_indices.shape[0]
+    assert len(expert_weights[0]) >= 1
+    num_local_physical_experts = expert_weights[0].shape[0]
+    assert num_physical_experts == ep_size * num_local_physical_experts
+
+    old_layer_indices_np = old_layer_indices.cpu().numpy()
+    new_layer_indices_np = new_layer_indices.cpu().numpy()
+
+    return move_to_buffer(
+        num_local_experts=num_local_physical_experts,
+        old_indices=old_layer_indices_np,
+        new_indices=new_layer_indices_np,
+        expert_weights=expert_weights,
+        expert_weights_buffers=expert_weights_buffer,
+        cuda_stream=cuda_stream,
+        ep_rank=ep_group.rank(),
+        communicator=communicator,
+        layer_idx=layer_idx,
+    )
+
+
+def rearrange_expert_weights_inplace(
+    old_global_expert_indices: torch.Tensor,
+    new_global_expert_indices: torch.Tensor,
+    expert_weights: Sequence[Sequence[torch.Tensor]],
+    expert_buffer: Sequence[torch.Tensor],
+    ep_group: ProcessGroup,
+    communicator: EplbCommunicator,
+    is_profile: bool = False,
+    rank_mapping: dict[int, int] | None = None,
+) -> None:
+    """
+    Rearranges the expert weights in place according to the new expert indices.
+
+    The value of the indices arguments are logical indices of the experts,
+    while keys are physical.
+
+    Args:
+        old_global_expert_indices: Shape (num_moe_layers, num_physical_experts).
+        new_global_expert_indices: Shape (num_moe_layers, num_physical_experts).
+        expert_weights: A sequence of shape (num_moe_layers)(weight_count)
+            of tensors of shape (num_local_physical_experts, hidden_size_i).
+            For example, a linear layer may have up and down projection,
+            so weight_count = 2. Each weight's hidden size can be different.
+        expert_buffer: Pre-allocated receive buffer tensors (one per
+            weight tensor in a single layer).
+        ep_group: The device process group for expert parallelism.
+        communicator: EplbCommunicator instance for P2P communication.
+        is_profile (bool): If `True`, do not perform any actual weight copy.
+            This is used during profile run, where we only perform dummy
+            communications to reserve enough memory for the buffers.
+        rank_mapping: A dictionary mapping old rank to new rank.
+    """
+    if rank_mapping is not None:
+        if len(rank_mapping) == ep_group.size():
+            # scale down
+            new_global_expert_indices = _map_new_expert_indices_with_rank_mapping(
+                new_global_expert_indices,
+                rank_mapping,
+            )
+        else:
+            # scale up
+            old_global_expert_indices = _map_old_expert_indices_with_rank_mapping(
+                old_global_expert_indices,
+                rank_mapping,
+                ep_group.size(),
+            )
+
+    assert old_global_expert_indices.shape[1] == new_global_expert_indices.shape[1]
+
+    num_moe_layers, num_physical_experts = old_global_expert_indices.shape
+    assert len(expert_weights) == num_moe_layers
+    assert len(expert_weights[0]) >= 1
+
+    num_local_physical_experts = expert_weights[0][0].shape[0]
+    assert new_global_expert_indices.shape == (num_moe_layers, num_physical_experts)
+
+    ep_size = ep_group.size()
+    ep_rank = ep_group.rank()
+    assert num_physical_experts == ep_size * num_local_physical_experts
+
+    first_layer_weights = list(expert_weights[0])
+
+    if is_profile:
+        if communicator.needs_profile_buffer_reservation:
+            # Reserve NCCL communication buffers via a dummy all_gather.
+            # Backends that pre-allocate their own transfer buffers
+            # skip this to avoid the extra memory spike during profiling.
+            profile_buffer: list[torch.Tensor] = [
+                torch.empty_like(w) for w in first_layer_weights
+            ]
+            for weight, buffer in zip(expert_weights[0], profile_buffer):
+                dummy_recv_buffer = [buffer for _ in range(ep_size)]
+                torch.distributed.barrier()
+                all_gather(
+                    dummy_recv_buffer,
+                    weight,
+                    group=ep_group,
+                )
+        return
+
+    weights_buffer = list(expert_buffer)
+
+    old_global_expert_indices_cpu = old_global_expert_indices.cpu().numpy()
+    new_global_expert_indices_cpu = new_global_expert_indices.cpu().numpy()
+
+    for layer_idx in range(num_moe_layers):
+        transfer_metadata = move_to_buffer(
+            num_local_experts=num_local_physical_experts,
+            old_indices=old_global_expert_indices_cpu[layer_idx],
+            new_indices=new_global_expert_indices_cpu[layer_idx],
+            expert_weights=expert_weights[layer_idx],
+            expert_weights_buffers=weights_buffer,
+            cuda_stream=None,
+            ep_rank=ep_rank,
+            communicator=communicator,
+            layer_idx=layer_idx,
+        )
+
+        move_from_buffer(
+            expert_weights=expert_weights[layer_idx],
+            expert_weights_buffers=weights_buffer,
+            transfer_metadata=transfer_metadata,
+            new_indices=new_global_expert_indices_cpu[layer_idx],
+            ep_rank=ep_rank,
+        )
+
+
+def _map_old_expert_indices_with_rank_mapping(
+    old_global_expert_indices: torch.Tensor,
+    rank_mapping: dict[int, int],
+    new_ep_size: int,
+) -> torch.Tensor:
+    """
+    Map the old global expert indices to the new global expert indices.
+
+    Args:
+        old_global_expert_indices:
+            Shape (num_layers, old_ep_size * num_local_physical_experts).
+        rank_mapping: Mapping from old rank to new rank.
+        new_ep_size: New expert parallelism size.
+
+    Returns:
+        Mapped expert indices with shape
+        (num_layers, new_ep_size * num_local_physical_experts).
+    """
+    num_layers, old_num_physical_experts = old_global_expert_indices.shape
+    assert rank_mapping, "Rank mapping is required"
+
+    # Get sizes from parameters and rank_mapping
+    old_ep_size = len(rank_mapping)
+    num_local_physical_experts = old_num_physical_experts // old_ep_size
+    new_num_physical_experts = new_ep_size * num_local_physical_experts
+
+    # Create mapped tensor with new shape, initialized to -1
+    mapped_expert_indices = torch.full(
+        (num_layers, new_num_physical_experts),
+        fill_value=-1,
+        dtype=old_global_expert_indices.dtype,
+        device=old_global_expert_indices.device,
+    )
+
+    # Handle rank mapping (scale up/down with rank changes)
+    for old_rank in range(old_ep_size):
+        new_rank = rank_mapping.get(old_rank)
+        if new_rank is not None and new_rank >= 0 and new_rank < new_ep_size:
+            # This old rank exists in the new configuration
+            old_start_idx = old_rank * num_local_physical_experts
+            old_end_idx = (old_rank + 1) * num_local_physical_experts
+            new_start_idx = new_rank * num_local_physical_experts
+            new_end_idx = (new_rank + 1) * num_local_physical_experts
+
+            mapped_expert_indices[:, new_start_idx:new_end_idx] = (
+                old_global_expert_indices[:, old_start_idx:old_end_idx]
+            )
+        # If new_rank is None or >= new_ep_size, the experts remain -1
+        # (scale down case)
+
+    return mapped_expert_indices
+
+
+def _map_new_expert_indices_with_rank_mapping(
+    new_global_expert_indices: torch.Tensor,
+    rank_mapping: dict[int, int],
+) -> torch.Tensor:
+    num_layers, new_num_physical_experts = new_global_expert_indices.shape
+    assert rank_mapping, "Rank mapping is required"
+
+    # Get sizes from parameters and rank_mapping
+    old_ep_size = len(rank_mapping)
+    new_ep_size = sum(new_rank != -1 for new_rank in rank_mapping.values())
+    num_local_physical_experts = new_num_physical_experts // new_ep_size
+    old_num_physical_experts = old_ep_size * num_local_physical_experts
+
+    mapped_expert_indices = torch.full(
+        (num_layers, old_num_physical_experts),
+        fill_value=-1,
+        dtype=new_global_expert_indices.dtype,
+        device=new_global_expert_indices.device,
+    )
+
+    for old_rank in range(old_ep_size):
+        new_rank = rank_mapping[old_rank]
+        if new_rank >= 0 and new_rank < new_ep_size:
+            old_start_idx = old_rank * num_local_physical_experts
+            old_end_idx = (old_rank + 1) * num_local_physical_experts
+            new_start_idx = new_rank * num_local_physical_experts
+            new_end_idx = (new_rank + 1) * num_local_physical_experts
+
+            mapped_expert_indices[:, old_start_idx:old_end_idx] = (
+                new_global_expert_indices[:, new_start_idx:new_end_idx]
+            )
+
+    return mapped_expert_indices
+
+
+__all__ = ["transfer_layer", "move_from_buffer", "TransferMetadata"]

vllm/distributed/kv_events.py

@@ -0,0 +1,537 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import queue
+import threading
+import time
+from abc import ABC, abstractmethod
+from collections import Counter, deque
+from collections.abc import Callable
+from dataclasses import asdict
+from itertools import count
+from queue import Queue
+from typing import Any
+
+import msgspec
+import zmq
+
+from vllm.config.kv_events import KVEventsConfig
+from vllm.logger import init_logger
+from vllm.v1.core.kv_cache_utils import ExternalBlockHash
+
+logger = init_logger(__name__)
+
+
+class EventBatch(
+    msgspec.Struct,
+    array_like=True,  # type: ignore[call-arg]
+    omit_defaults=True,  # type: ignore[call-arg]
+    gc=False,  # type: ignore[call-arg]
+):
+    ts: float
+    events: list[Any]
+    data_parallel_rank: int | None = None
+
+
+class KVCacheEvent(
+    msgspec.Struct,
+    omit_defaults=True,  # type: ignore[call-arg]
+    gc=False,  # type: ignore[call-arg]
+    tag=True,
+):
+    """Base class for all KV cache-related events"""
+
+
+MEDIUM_GPU = "GPU"
+MEDIUM_CPU = "CPU"
+MEDIUM_FS = "FS"
+MEDIUM_OBJ = "OBJ"
+
+
+class BlockStored(KVCacheEvent):
+    block_hashes: list[ExternalBlockHash]
+    parent_block_hash: ExternalBlockHash | None
+    token_ids: list[int]
+    block_size: int
+
+    lora_id: int | None
+    """Deprecated: use `lora_name` for KV block key hash.
+    Retained for backward compatibility.
+    """
+
+    medium: str | None
+    lora_name: str | None
+
+    extra_keys: list[tuple[Any, ...] | None] | None = None
+    """Extra keys used in block hash computation, one entry per block in
+    block_hashes. Each entry contains MM identifiers, LoRA name, cache_salt,
+    prompt embedding hashes, etc. for that specific block. Exposed for external
+    KV cache consumers to reconstruct block hashes.
+    """
+
+    group_idx: int | None = None
+    # Store events carry cache-spec metadata so consumers can classify and
+    # filter groups as they are learned. Remove events only need group_idx+hash.
+    kv_cache_spec_kind: str | None = None
+    kv_cache_spec_sliding_window: int | None = None
+
+    def __hash__(self) -> int:
+        return hash(
+            (
+                tuple(self.block_hashes),
+                self.parent_block_hash,
+                tuple(self.token_ids),
+                self.block_size,
+                self.lora_id,
+                self.medium,
+                tuple(self.extra_keys) if self.extra_keys else None,
+                self.group_idx,
+                self.kv_cache_spec_kind,
+                self.kv_cache_spec_sliding_window,
+            )
+        )
+
+
+class BlockRemoved(KVCacheEvent):
+    block_hashes: list[ExternalBlockHash]
+    medium: str | None
+    group_idx: int | None = None
+
+    def __hash__(self) -> int:
+        return hash(
+            (
+                tuple(self.block_hashes),
+                self.medium,
+                self.group_idx,
+            )
+        )
+
+
+class AllBlocksCleared(KVCacheEvent):
+    pass
+
+
+class KVEventBatch(EventBatch):
+    events: list[BlockStored | BlockRemoved | AllBlocksCleared]
+
+
+class KVEventAggregator:
+    """
+    Aggregates KV events across multiple workers.
+    Tracks how many times each event appears and returns only those
+    that were emitted by all workers.
+    """
+
+    __slots__ = ("_event_counter", "_num_workers")
+
+    def __init__(self, num_workers: int) -> None:
+        if num_workers <= 0:
+            raise ValueError("num_workers must be greater than zero.")
+        self._event_counter: Counter[KVCacheEvent] = Counter()
+        self._num_workers: int = num_workers
+
+    def add_events(self, events: list[KVCacheEvent]) -> None:
+        """
+        Add events from a worker batch.
+
+        Args:
+            events: List of KVCacheEvent objects.
+        """
+        if not isinstance(events, list):
+            raise TypeError("events must be a list of KVCacheEvent.")
+        self._event_counter.update(events)
+
+    def get_common_events(self) -> list[KVCacheEvent]:
+        """
+        Return events that appeared in all workers.
+
+        Returns:
+            List of events present in all workers.
+        """
+        return [
+            event
+            for event, count in self._event_counter.items()
+            if count == self._num_workers
+        ]
+
+    def get_all_events(self) -> list[KVCacheEvent]:
+        """
+        Return all events for all workers.
+
+        Returns:
+            List of events for all workers.
+        """
+        return list(self._event_counter.elements())
+
+    def clear_events(self) -> None:
+        """
+        Clear all tracked events.
+        """
+        self._event_counter.clear()
+
+    def increment_workers(self, count: int = 1) -> None:
+        """
+        Increment the number of workers contributing events.
+
+        Args:
+            count: Number to increment the workers by.
+        """
+        if count <= 0:
+            raise ValueError("count must be positive.")
+        self._num_workers += count
+
+    def reset_workers(self) -> None:
+        """
+        Reset the number of workers to 1.
+        """
+        self._num_workers = 1
+
+    def get_number_of_workers(self) -> int:
+        """
+        Return the number of workers.
+
+        Returns:
+            int number of workers.
+        """
+        return self._num_workers
+
+    def __repr__(self) -> str:
+        return (
+            f"<KVEventAggregator workers={self._num_workers}, "
+            f"events={len(self._event_counter)}>"
+        )
+
+
+class KVConnectorKVEvents(ABC):
+    """
+    Abstract base class for KV events.
+    Acts as a container for KV events from the connector.
+    """
+
+    @abstractmethod
+    def add_events(self, events: list[KVCacheEvent]) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def aggregate(self) -> "KVConnectorKVEvents":
+        raise NotImplementedError
+
+    @abstractmethod
+    def increment_workers(self, count: int = 1) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_all_events(self) -> list[KVCacheEvent]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_number_of_workers(self) -> int:
+        raise NotImplementedError
+
+    @abstractmethod
+    def clear_events(self) -> None:
+        raise NotImplementedError
+
+    def merge(self, other: "KVConnectorKVEvents") -> "KVConnectorKVEvents":
+        self.add_events(other.get_all_events())
+        return self
+
+
+class EventPublisher(ABC):
+    """Lightweight publisher for EventBatch batches with data parallelism
+    support.
+
+    In data parallel setups, each DP rank runs its own EventPublisher instance
+    to avoid duplicate events and ensure proper event attribution:
+
+    - Each DP rank creates a separate publisher
+    - Publishers automatically annotate events with their data_parallel_rank
+    - This allows consumers to distinguish events from different DP ranks
+
+    The publisher is responsible for adding DP metadata since the scheduler
+    operates independently of DP topology and shouldn't need DP awareness.
+    """
+
+    def __init__(self, data_parallel_rank: int = 0) -> None:
+        self._data_parallel_rank = data_parallel_rank
+
+    @abstractmethod
+    def publish(self, events: EventBatch) -> None:
+        """Emit events in order.
+
+        Implementations should guarantee at-least-once delivery and
+        monotonic ordering (e.g., via sequence numbers).
+        """
+
+    @abstractmethod
+    def shutdown(self) -> None:
+        """Shutdown the publisher."""
+
+
+class NullEventPublisher(EventPublisher):
+    """No-op implementation (default when disabled)."""
+
+    def publish(self, events) -> None:
+        return
+
+    def shutdown(self) -> None:
+        return
+
+
+class ZmqEventPublisher(EventPublisher):
+    """Reliable PUB/ROUTER publisher with an in-memory replay buffer.
+
+    Spawns a separate thread to handle publishing from a queue.
+
+    Parameters
+    ----------
+    endpoint:
+        PUB address. Use `tcp://*:5557` to bind or `tcp://host:5557` to
+        connect.
+    replay_endpoint:
+        Optional ROUTER address for replay requests. When given, subscribers can
+        request missed batches by sending the starting sequence number as an
+        8-byte big-endian integer.
+    buffer_steps:
+        Number of past batches to keep for replay.
+    hwm:
+        ZeroMQ high-water-mark for PUB socket.
+    max_queue_size:
+        Maximum number of events to buffer in memory.
+    topic:
+        Topic to publish events to.
+    """
+
+    SHUTDOWN_TIMEOUT: float = 1.0
+    END_SEQ = (-1).to_bytes(8, "big", signed=True)
+
+    def __init__(
+        self,
+        data_parallel_rank: int,
+        endpoint: str = "tcp://*:5557",
+        replay_endpoint: str | None = None,
+        buffer_steps: int = 10_000,
+        hwm: int = 100_000,
+        max_queue_size: int = 100_000,
+        topic: str = "",
+    ) -> None:
+        # Storage
+        super().__init__(data_parallel_rank)
+        self._event_queue = Queue[EventBatch | None](maxsize=max_queue_size)
+        self._buffer = deque[tuple[int, bytes]](maxlen=buffer_steps)
+
+        # ZMQ sockets
+        self._ctx = zmq.Context.instance()
+        self._pub: zmq.Socket | None = None
+        self._replay: zmq.Socket | None = None
+        self._dp_rank = data_parallel_rank
+
+        self._endpoint = self.offset_endpoint_port(endpoint, self._dp_rank)
+        self._replay_endpoint = self.offset_endpoint_port(
+            replay_endpoint, self._dp_rank
+        )
+        self._hwm = hwm
+        self._socket_setup()
+
+        # Payload
+        self._seq_gen = count()
+        self._topic_bytes = topic.encode("utf-8")
+
+        # Thread
+        self._running = True
+        logger.info("Starting ZMQ publisher thread")
+
+        self._thread = threading.Thread(
+            target=self._publisher_thread, daemon=True, name="zmq-publisher"
+        )
+        self._thread.start()
+
+    def publish(self, events: EventBatch) -> None:
+        if not self._running:
+            raise RuntimeError("Publisher is closed")
+        if events.data_parallel_rank is None:
+            events.data_parallel_rank = self._data_parallel_rank
+        self._event_queue.put(events)
+
+    def shutdown(self) -> None:
+        """Stop the publisher thread and clean up resources."""
+        self._running = False
+        self._event_queue.put_nowait(None)
+
+        start = time.time()
+        pending_items = True
+        while pending_items and (time.time() - start < self.SHUTDOWN_TIMEOUT):
+            pending_items = not self._event_queue.empty()
+            if pending_items:
+                time.sleep(0.1)
+
+        if pending_items:
+            logger.warning(
+                "Warning: Queue still has %s items after %s seconds timeout",
+                self._event_queue.qsize(),
+                self.SHUTDOWN_TIMEOUT,
+            )
+
+        if self._thread.is_alive():
+            self._thread.join(timeout=self.SHUTDOWN_TIMEOUT)
+
+        # Clean up ZMQ resources
+        try:
+            if self._pub is not None:
+                self._pub.close(linger=0)
+            if self._replay is not None:
+                self._replay.close(linger=0)
+        finally:
+            pass  # Do not terminate context; other sockets may use it
+
+    def _socket_setup(self) -> None:
+        """Initialize sockets
+        https://pyzmq.readthedocs.io/en/v19.0.0/morethanbindings.html#thread-safety
+        """
+        if self._pub is None:
+            self._pub = self._ctx.socket(zmq.PUB)
+            self._pub.set_hwm(self._hwm)
+            # Heuristic: bind if wildcard / * present, else connect.
+            # bind stable, connect volatile convention
+            if self._endpoint is not None and (
+                "*" in self._endpoint
+                or "::" in self._endpoint
+                or self._endpoint.startswith("ipc://")
+                or self._endpoint.startswith("inproc://")
+            ):
+                self._pub.bind(self._endpoint)
+            elif self._endpoint is not None:
+                self._pub.connect(self._endpoint)
+
+        # Set up replay socket: use ROUTER
+        # 1) handles multiple REQ clients (identities)
+        # 2) lets us send back one request → many replies (streamed events)
+        # 3) works in our non‑blocking poll loop alongside PUB
+        if self._replay_endpoint is not None:
+            self._replay = self._ctx.socket(zmq.ROUTER)
+            self._replay.bind(self._replay_endpoint)
+
+    def _publisher_thread(self) -> None:
+        """Background thread that processes the event queue."""
+        self._pack = msgspec.msgpack.Encoder()
+
+        assert self._pub is not None  # narrows type for mypy
+
+        while self._running or self._event_queue.qsize() > 0:
+            # --- replay (non-critical) ---------------------------------
+            if self._replay is not None and self._replay.poll(0):
+                try:
+                    self._service_replay()
+                except Exception as e:
+                    logger.exception("Error in replay: %s", e)
+
+            # --- main queue (critical) ---------------------------------
+            try:
+                event = self._event_queue.get(timeout=0.1)
+                if event is None:
+                    break  # Sentinel received, exit thread
+            except queue.Empty:
+                continue
+
+            try:
+                seq = next(self._seq_gen)
+
+                payload = self._pack.encode(event)
+                seq_bytes = seq.to_bytes(8, "big")
+                self._pub.send_multipart((self._topic_bytes, seq_bytes, payload))
+
+                self._buffer.append((seq, payload))
+                self._event_queue.task_done()
+
+            except Exception as e:
+                # Publishing failed;  back-off a bit to avoid a tight error loop
+                logger.exception("Error in publisher thread: %s", e)
+                time.sleep(0.1)
+
+    def _service_replay(self) -> None:
+        """If a replay request is waiting, send buffered batches."""
+        assert self._replay is not None  # narrows type for mypy
+
+        frame = self._replay.recv_multipart()
+        if len(frame) != 3:
+            logger.warning("Invalid replay request: %s", frame)
+            return
+        client_id, _, start_seq_bytes = frame
+        start_seq = int.from_bytes(start_seq_bytes, "big")
+
+        for seq, buf in self._buffer:
+            if seq >= start_seq:
+                # Subscriber receives (topic, seq_bytes, payload)
+                self._replay.send_multipart(
+                    (client_id, b"", self._topic_bytes, seq.to_bytes(8, "big"), buf)
+                )
+        # Send end of sequence marker
+        self._replay.send_multipart((client_id, b"", b"", self.END_SEQ, b""))
+
+    @staticmethod
+    def offset_endpoint_port(
+        endpoint: str | None, data_parallel_rank: int
+    ) -> str | None:
+        """Helper function to offset the port in an endpoint by
+            the data parallel rank.
+
+        Args:
+            endpoint: The endpoint string
+                (e.g., "tcp://*:5557" or "inproc://cache")
+            data_parallel_rank: The data parallel rank to offset by
+
+        Returns:
+            The endpoint with the port offset by data_parallel_rank
+                or suffix appended
+        """
+        # Do nothing if input is None or data_parallel_rank is 0
+        if not endpoint or data_parallel_rank == 0:
+            return endpoint
+
+        if "inproc" in endpoint:
+            return f"{endpoint}_dp{data_parallel_rank}"
+        if "tcp" in endpoint:
+            if endpoint and ":" in endpoint:
+                # Get everything after the last colon (the port)
+                last_colon_idx = endpoint.rfind(":")
+                base_addr = endpoint[:last_colon_idx]
+                base_port = int(endpoint[last_colon_idx + 1 :])
+                new_port = base_port + data_parallel_rank
+                return f"{base_addr}:{new_port}"
+            return endpoint
+        raise ValueError("Invalid endpoint: must contain 'inproc' or 'tcp'")
+
+
+class EventPublisherFactory:
+    _registry: dict[str, Callable[..., EventPublisher]] = {
+        "null": NullEventPublisher,
+        "zmq": ZmqEventPublisher,
+    }
+
+    @classmethod
+    def register_publisher(cls, name: str, ctor: Callable[..., EventPublisher]) -> None:
+        if name in cls._registry:
+            raise KeyError(f"publisher '{name}' already registered")
+        cls._registry[name] = ctor
+
+    @classmethod
+    def create(
+        cls, config: KVEventsConfig | None, data_parallel_rank: int = 0
+    ) -> EventPublisher:
+        """Create publisher from a config mapping."""
+        if (
+            config is None
+            or not config.enable_kv_cache_events
+            or config.publisher == "null"
+        ):
+            return NullEventPublisher()
+
+        config_dict = asdict(config)
+
+        kind = config_dict.pop("publisher")
+        config_dict.pop("enable_kv_cache_events")
+        try:
+            constructor = cls._registry[kind]
+        except KeyError as exc:
+            raise ValueError(f"Unknown event publisher '{kind}'") from exc
+        return constructor(data_parallel_rank=data_parallel_rank, **config_dict)

vllm/distributed/kv_transfer/README.md

@@ -0,0 +1,29 @@
+
+# Distributed KV cache transfer
+
+This folder implements distributed KV cache transfer across vLLM instances.
+Currently the main use case is for disaggregated prefilling.
+
+## Abstractions
+
+The KV cache transfer contains three layer of abstractions:
+
+- KV pipe: a FIFO pipe for torch.tensor transmission. Key APIs: `send_tensor` and `recv_tensor`.
+- KV lookup buffer: a lookup buffer for KV caches. Key: the tokens, value: the KV caches (and/or hidden states). Key APIs: `insert` and `drop_select` (similar to SQL semantics).
+- KV connector: a connector that connects the KV pipe and KV lookup buffer to vLLM. Key APIs: `send_kv_caches_and_hidden_states` and `recv_kv_caches_and_hidden_states`.
+
+Why we need KV lookup buffer: FIFO pipe itself is not enough as prefill vLLM worker may process requests in a different order compared to decode vLLM worker. Say the QPS is really high, prefill worker may handle requests in order A -> B -> C, but the decode worker may process request C first. This is not the case that can be naturally handled by FIFO pipe, so we provide KV lookup buffer to help translate a FIFO pipe to a lookup buffer.
+
+NOTE: KV pipe layer is bypassable: you can skip this layer if your distributed
+communication service already supports key-value-based lookup (like redis or
+RDMA database).
+
+NOTE: If you want to not only transfer KV caches, but adjust the model execution flow of vLLM as well (for example, allow vLLM to receive KV caches on some tokens and do prefill on the remaining tokens), you can bypass both KV pipe layer and KV lookup buffer layer, and directly implement on KV connector layer. Bear in mind that as vLLM's model input is constantly changing, this implementation will likely be broken when vLLM has new updates.
+
+## Disaggregated prefilling
+
+The example usage is in [this file](../../../examples/disaggregated/disaggregated_prefill.sh).
+
+Here is the diagram of how we run disaggregated prefilling.
+
+![Disaggregated prefill workflow](./disagg_prefill_workflow.jpg)

vllm/distributed/kv_transfer/__init__.py

@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.distributed.kv_transfer.kv_transfer_state import (
+    KVConnectorBaseType,
+    ensure_kv_transfer_initialized,
+    ensure_kv_transfer_shutdown,
+    get_kv_transfer_group,
+    has_kv_transfer_group,
+    is_v1_kv_transfer_group,
+)
+
+__all__ = [
+    "get_kv_transfer_group",
+    "has_kv_transfer_group",
+    "is_v1_kv_transfer_group",
+    "ensure_kv_transfer_initialized",
+    "ensure_kv_transfer_shutdown",
+    "KVConnectorBaseType",
+]

vllm/distributed/kv_transfer/kv_connector/base.py

@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Defines the base type for KV cache connectors."""
+
+from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1
+
+KVConnectorBase = KVConnectorBase_V1
+KVConnectorBaseType = KVConnectorBase_V1
+
+__all__ = ["KVConnectorBase", "KVConnectorBaseType"]

vllm/distributed/kv_transfer/kv_connector/factory.py

@@ -0,0 +1,242 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import importlib
+from collections.abc import Callable
+from typing import TYPE_CHECKING, cast
+
+from vllm.config.kv_transfer import KVTransferConfig
+from vllm.distributed.kv_transfer.kv_connector.base import (
+    KVConnectorBase,
+    KVConnectorBaseType,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1 import (
+    KVConnectorRole,
+    supports_hma,
+)
+from vllm.logger import init_logger
+from vllm.utils.func_utils import supports_kw
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+
+logger = init_logger(__name__)
+
+
+class KVConnectorFactory:
+    _registry: dict[str, Callable[[], type[KVConnectorBase]]] = {}
+
+    @classmethod
+    def register_connector(cls, name: str, module_path: str, class_name: str) -> None:
+        """Register a connector with a lazy-loading module and class name."""
+        if name in cls._registry:
+            raise ValueError(f"Connector '{name}' is already registered.")
+
+        def loader() -> type[KVConnectorBase]:
+            module = importlib.import_module(module_path)
+            return getattr(module, class_name)
+
+        cls._registry[name] = loader
+
+    @classmethod
+    def create_connector(
+        cls,
+        config: "VllmConfig",
+        role: KVConnectorRole,
+        kv_cache_config: "KVCacheConfig",
+    ) -> KVConnectorBase:
+        kv_transfer_config = config.kv_transfer_config
+        if kv_transfer_config is None:
+            raise ValueError("kv_transfer_config must be set to create a connector")
+        connector_cls = cls.get_connector_class(kv_transfer_config)
+
+        # check if the connector supports HMA
+        hma_enabled = not config.scheduler_config.disable_hybrid_kv_cache_manager
+        if hma_enabled and not cls.supports_hma_config(kv_transfer_config):
+            raise ValueError(
+                f"Connector {connector_cls.__name__} does not support HMA but "
+                f"HMA is enabled. Please set `--disable-hybrid-kv-cache-manager`."
+            )
+
+        logger.info(
+            "Creating v1 connector with name: %s and engine_id: %s",
+            connector_cls.__name__,
+            kv_transfer_config.engine_id,
+        )
+        # NOTE(Kuntai): v1 connector is explicitly separated into two roles.
+        # Scheduler connector:
+        # - Co-locate with scheduler process
+        # - Should only be used inside the Scheduler class
+        # Worker connector:
+        # - Co-locate with worker process
+        # - Should only be used inside the forward context & attention layer
+        # We build separately to enforce strict separation
+        return connector_cls(config, role, kv_cache_config)
+
+    @classmethod
+    def get_connector_class_by_name(
+        cls, connector_name: str
+    ) -> type[KVConnectorBaseType]:
+        """Get a registered connector class by name.
+
+        Raises ValueError if the connector is not registered.
+
+        Args:
+            connector_name: Name of the registered connector.
+
+        Returns:
+            The connector class.
+        """
+        if connector_name not in cls._registry:
+            raise ValueError(f"Connector '{connector_name}' is not registered.")
+        return cls._registry[connector_name]()
+
+    @classmethod
+    def get_connector_class(
+        cls, kv_transfer_config: "KVTransferConfig"
+    ) -> type[KVConnectorBaseType]:
+        connector_name = kv_transfer_config.kv_connector
+        if connector_name is None:
+            raise ValueError("Connector name is not set in KVTransferConfig")
+        connector_module_path = kv_transfer_config.kv_connector_module_path
+        if connector_module_path is not None and not connector_module_path:
+            raise ValueError("kv_connector_module_path cannot be an empty string.")
+        if connector_module_path:
+            # External module path takes priority over internal registry.
+            connector_module = importlib.import_module(connector_module_path)
+            try:
+                connector_cls = getattr(connector_module, connector_name)
+            except AttributeError as e:
+                raise AttributeError(
+                    f"Class {connector_name} not found in {connector_module_path}"
+                ) from e
+            connector_cls = cast(type[KVConnectorBaseType], connector_cls)
+            if not supports_kw(connector_cls, "kv_cache_config"):
+                msg = (
+                    f"Connector {connector_cls.__name__} uses deprecated "
+                    "2-argument constructor signature. External v1 KV "
+                    "connectors must accept kv_cache_config as the third "
+                    "constructor argument and pass it to super().__init__()."
+                )
+                logger.error(msg)
+                raise ValueError(msg)
+        elif connector_name in cls._registry:
+            connector_cls = cls._registry[connector_name]()
+        else:
+            raise ValueError(f"Unsupported connector type: {connector_name}")
+        return connector_cls
+
+    @classmethod
+    def supports_hma_config(cls, kv_transfer_config: "KVTransferConfig") -> bool:
+        """Return whether this KV transfer config supports HMA.
+
+        MultiConnector is a special case: the wrapper class implements
+        SupportsHMA, but effective support depends on every configured child.
+        """
+        connector_cls = cls.get_connector_class(kv_transfer_config)
+        if kv_transfer_config.kv_connector != "MultiConnector":
+            return supports_hma(connector_cls)
+
+        from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
+            MultiConnector,
+        )
+
+        return MultiConnector.all_children_support_hma(kv_transfer_config)
+
+
+# Register various connectors here.
+# The registration should not be done in each individual file, as we want to
+# only load the files corresponding to the current connector.
+
+KVConnectorFactory.register_connector(
+    "ExampleConnector",
+    "vllm.distributed.kv_transfer.kv_connector.v1.example_connector",
+    "ExampleConnector",
+)
+
+KVConnectorFactory.register_connector(
+    "ExampleHiddenStatesConnector",
+    "vllm.distributed.kv_transfer.kv_connector.v1.example_hidden_states_connector",
+    "ExampleHiddenStatesConnector",
+)
+
+KVConnectorFactory.register_connector(
+    "LMCacheConnectorV1",
+    "vllm.distributed.kv_transfer.kv_connector.v1.lmcache_connector",
+    "LMCacheConnectorV1",
+)
+
+KVConnectorFactory.register_connector(
+    "LMCacheMPConnector",
+    "vllm.distributed.kv_transfer.kv_connector.v1.lmcache_mp_connector",
+    "LMCacheMPConnector",
+)
+
+KVConnectorFactory.register_connector(
+    "NixlConnector",
+    "vllm.distributed.kv_transfer.kv_connector.v1.nixl",
+    "NixlConnector",
+)
+
+KVConnectorFactory.register_connector(
+    "NixlPullConnector",
+    "vllm.distributed.kv_transfer.kv_connector.v1.nixl",
+    "NixlPullConnector",
+)
+
+KVConnectorFactory.register_connector(
+    "NixlPushConnector",
+    "vllm.distributed.kv_transfer.kv_connector.v1.nixl",
+    "NixlPushConnector",
+)
+
+KVConnectorFactory.register_connector(
+    "MultiConnector",
+    "vllm.distributed.kv_transfer.kv_connector.v1.multi_connector",
+    "MultiConnector",
+)
+
+KVConnectorFactory.register_connector(
+    "MoRIIOConnector",
+    "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector",
+    "MoRIIOConnector",
+)
+
+KVConnectorFactory.register_connector(
+    "OffloadingConnector",
+    "vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector",
+    "OffloadingConnector",
+)
+
+KVConnectorFactory.register_connector(
+    "DecodeBenchConnector",
+    "vllm.distributed.kv_transfer.kv_connector.v1.decode_bench_connector",
+    "DecodeBenchConnector",
+)
+
+KVConnectorFactory.register_connector(
+    "MooncakeConnector",
+    "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector",
+    "MooncakeConnector",
+)
+KVConnectorFactory.register_connector(
+    "MooncakeStoreConnector",
+    "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.connector",
+    "MooncakeStoreConnector",
+)
+KVConnectorFactory.register_connector(
+    "FlexKVConnectorV1",
+    "vllm.distributed.kv_transfer.kv_connector.v1.flexkv_connector",
+    "FlexKVConnectorV1",
+)
+KVConnectorFactory.register_connector(
+    "SimpleCPUOffloadConnector",
+    "vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector",
+    "SimpleCPUOffloadConnector",
+)
+KVConnectorFactory.register_connector(
+    "HF3FSKVConnector",
+    "vllm.distributed.kv_transfer.kv_connector.v1.hf3fs.hf3fs_connector",
+    "HF3FSKVConnector",
+)

vllm/distributed/kv_transfer/kv_connector/utils.py

@@ -0,0 +1,635 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+KV cache helper for store.
+"""
+
+from collections.abc import Iterator
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Literal, cast
+
+import torch
+
+from vllm.config import (
+    VllmConfig,
+    get_current_vllm_config,
+    get_layers_from_vllm_config,
+    set_current_vllm_config,
+)
+from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
+from vllm.logger import init_logger
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+from vllm.platforms import current_platform
+from vllm.v1.attention.backend import AttentionBackend
+from vllm.v1.kv_cache_interface import MambaSpec
+from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
+
+if TYPE_CHECKING:
+    from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
+    from vllm.v1.kv_cache_interface import KVCacheSpec
+
+logger = init_logger(__name__)
+
+EngineId = str
+# block ids as returned by the hybrid KV cache manager. list[list[int]] are allow
+# mutability and are for connector internal use only.
+BlockIds = tuple[list[int], ...] | list[list[int]]
+
+
+def get_kv_connector_cache_layout():
+    # NOTE (NickLucche) When running disaggregated PD with NIXL, HND layout is
+    # used for faster transfer.
+    vllm_config = get_current_vllm_config()
+    kv_config = vllm_config.kv_transfer_config
+    if kv_config is not None:
+        connector_cls = KVConnectorFactory.get_connector_class(kv_config)
+        required_kvcache_layout = connector_cls.get_required_kvcache_layout(vllm_config)
+        if required_kvcache_layout is not None:
+            return required_kvcache_layout
+        logger.info_once(
+            "Connectors do not specify a kv cache layout, defaulting to NHD."
+        )
+    return "NHD"
+
+
+class KVOutputAggregator:
+    """Utility class to aggregate the output of all workers into a single
+    output corresponding to Rank 0 for scheduler."""
+
+    def __init__(self, expected_finished_count: int):
+        # Complete transfer tracker. Used to track finished requests
+        # [req_id -> n_remaining_workers]
+        self._recv_remaining_count = dict[str, int]()
+        self._send_remaining_count = dict[str, int]()
+        self._expected_finished_count = expected_finished_count
+
+    @classmethod
+    def from_connector(cls, connector: "KVConnectorBase", world_size: int):
+        return cls(connector.get_finished_count() or world_size)
+
+    def aggregate(
+        self, outputs: list[ModelRunnerOutput | None], output_rank: int = 0
+    ) -> ModelRunnerOutput | None:
+        if not outputs[output_rank]:
+            return None
+
+        # Aggregate kv_connector_output from all workers
+
+        def update_finished_set(
+            req_ids: set[str] | None,
+            remaining_count_dict: dict[str, int],
+            finished_set: set[str],
+        ) -> None:
+            for req_id in req_ids or ():
+                remaining_count = remaining_count_dict.get(
+                    req_id, self._expected_finished_count
+                )
+                remaining_count_dict[req_id] = remaining_count - 1
+                if remaining_count_dict[req_id] == 0:
+                    finished_set.add(req_id)
+                    del remaining_count_dict[req_id]
+
+        finished_sending = set[str]()
+        finished_recving = set[str]()
+        aggregated_kv_connector_stats = None
+        aggregated_kv_connector_worker_meta = None
+        combined_kv_cache_events = None
+        invalid_block_ids = set[int]()
+        for model_runner_output in outputs:
+            assert model_runner_output is not None
+            kv_output = model_runner_output.kv_connector_output
+            if not kv_output:
+                continue
+            # Allow the worker to dynamically update the expected number of
+            # finished sending/recving for new requests.
+            if (
+                kv_output.expected_finished_count > 0
+                and kv_output.expected_finished_count != self._expected_finished_count
+            ):
+                logger.debug(
+                    "Expected finished requests updated from %d to %d",
+                    self._expected_finished_count,
+                    kv_output.expected_finished_count,
+                )
+                self._expected_finished_count = kv_output.expected_finished_count
+
+            update_finished_set(
+                kv_output.finished_sending, self._send_remaining_count, finished_sending
+            )
+            update_finished_set(
+                kv_output.finished_recving, self._recv_remaining_count, finished_recving
+            )
+
+            # Aggregate kv_connector_stats from all workers.
+            if aggregated_kv_connector_stats is None:
+                # Use the first worker's kv_connector_stats as accumulator.
+                aggregated_kv_connector_stats = kv_output.kv_connector_stats
+            elif kv_connector_stats := kv_output.kv_connector_stats:
+                assert isinstance(
+                    aggregated_kv_connector_stats, type(kv_connector_stats)
+                )
+                aggregated_kv_connector_stats = aggregated_kv_connector_stats.aggregate(
+                    kv_connector_stats
+                )
+
+            # Aggregate kv_connector_worker_meta from all workers.
+            if aggregated_kv_connector_worker_meta is None:
+                # Use the first worker's kv_connector_worker_meta as accumulator.
+                aggregated_kv_connector_worker_meta = kv_output.kv_connector_worker_meta
+            elif kv_connector_worker_meta := kv_output.kv_connector_worker_meta:
+                aggregated_kv_connector_worker_meta = (
+                    aggregated_kv_connector_worker_meta.aggregate(
+                        kv_connector_worker_meta
+                    )
+                )
+
+            # Combine kv_cache_events from all workers.
+            if combined_kv_cache_events is None:
+                # Use the first worker's kv_cache events as start event list.
+                combined_kv_cache_events = kv_output.kv_cache_events
+            elif kv_cache_events := kv_output.kv_cache_events:
+                assert isinstance(
+                    combined_kv_cache_events,
+                    type(kv_cache_events),
+                )
+                worker_kv_cache_events = kv_cache_events.get_all_events()
+                combined_kv_cache_events.add_events(worker_kv_cache_events)
+                combined_kv_cache_events.increment_workers(1)
+
+            invalid_block_ids |= kv_output.invalid_block_ids
+
+        # select output of the worker specified by output_rank
+        output = outputs[output_rank]
+
+        assert output is not None
+        output.kv_connector_output = KVConnectorOutput(
+            finished_sending=finished_sending or None,
+            finished_recving=finished_recving or None,
+            kv_connector_stats=aggregated_kv_connector_stats or None,
+            kv_cache_events=combined_kv_cache_events or None,
+            kv_connector_worker_meta=aggregated_kv_connector_worker_meta or None,
+            invalid_block_ids=invalid_block_ids,
+            expected_finished_count=self._expected_finished_count,
+        )
+
+        return output
+
+
+def _make_src_and_dst_indices(
+    src_block_ids: list[int],
+    dst_block_ids: list[int],
+    src_device: torch.device | str,
+    dst_device: torch.device | str,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    src_indices = torch.tensor(src_block_ids, device=src_device, dtype=torch.int64)
+    dst_indices = torch.tensor(dst_block_ids, device=dst_device, dtype=torch.int64)
+    return src_indices, dst_indices
+
+
+def copy_kv_blocks(
+    src_kv_caches: dict[str, torch.Tensor],
+    dst_kv_caches: dict[str, torch.Tensor],
+    src_block_ids: list[int],
+    dst_block_ids: list[int],
+    direction: Literal["h2d", "d2h"],
+) -> None:
+    """Copy kv blocks between different buffers."""
+    if (
+        not src_kv_caches
+        or not dst_kv_caches
+        or not src_block_ids
+        or not dst_block_ids
+        or len(src_block_ids) != len(dst_block_ids)
+    ):
+        return
+
+    src_device = next(iter(src_kv_caches.values())).device
+    dst_device = next(iter(dst_kv_caches.values())).device
+
+    src_indices, dst_indices = _make_src_and_dst_indices(
+        src_block_ids=src_block_ids,
+        dst_block_ids=dst_block_ids,
+        src_device=src_device,
+        dst_device=dst_device,
+    )
+
+    if direction == "h2d":
+        copy_fn = current_platform.insert_blocks_to_device
+    else:
+        copy_fn = current_platform.swap_out_blocks_to_host
+    for layer_name in src_kv_caches:
+        src_tensor = src_kv_caches[layer_name]
+        dst_tensor = dst_kv_caches[layer_name]
+        copy_fn(src_tensor, dst_tensor, src_indices, dst_indices)
+
+
+def kv_postprocess_blksize_on_receive(cache, indices, block_size_ratio):
+    """
+    Transforms the layout of received KV cache blocks to the local block_size.
+    (Only works for local blocksize > remote blocksize)
+
+    example:
+    local blocksize = 16 tokens, remote blocksize = 4 tokens
+    local block[0] = remote block[0, 1, 2, 3]
+    remote is |h0-b0|h1-b0|h2-b0|h3-b0|h0-b1|h1-b1|h2-b1|h3-b1|...
+    local is  |h0-b0..................|h1-b0..................|...
+    permute is to:
+    1. view => view remote as n_blocks * remote_shape(H,remoteN,D)
+    2. permute => (H, nblocks, remoteN, D)
+    3. flatten => (H, localN, D)
+    """
+    blocks_to_update = cache.index_select(0, indices)
+    # use physical order
+    blocks_to_update = blocks_to_update.permute(0, 2, 1, 3)
+    n_kv_heads, block_size, head_size = blocks_to_update.shape[1:]
+    remote_block_size = block_size // block_size_ratio
+    n_blocks = block_size_ratio
+
+    permuted_blocks = (
+        blocks_to_update.reshape(-1, n_blocks, n_kv_heads, remote_block_size, head_size)
+        .permute(0, 2, 1, 3, 4)
+        .flatten(2, 3)
+    )
+    permuted_blocks = permuted_blocks.permute(0, 2, 1, 3)
+    cache.index_copy_(0, indices, permuted_blocks)
+
+
+def kv_postprocess_layout_on_receive(cache, indices):
+    """Transforms the layout of received KV cache blocks to the local format.
+
+    This method corrects layout mismatches from direct memory copies by
+    permuting the tensor dimensions.
+
+    - **Source Layout:** `[num_blocks, n_kv_head, block_size, head_dim]`
+    - **Target Layout:** `[num_blocks, block_size, n_kv_head, head_dim]`
+
+    Implementation:
+    - x = blocks_to_update.reshape(src_shape) # view local kv with sender layout
+    - permuted_blocks = x.permute(*inv_order) # transpose n_kv_heads, block_size
+    - cache.index_copy_(0, indices, permuted_blocks) # copy permuted kv back
+
+    """
+    blocks_to_update = cache.index_select(0, indices)
+    target_shape = list(blocks_to_update.shape)
+    target_shape[0] = -1
+    inv_order = [0, 2, 1, 3]
+    src_shape = tuple(target_shape[i] for i in inv_order)
+    blocks_to_update = cache.index_select(0, indices)
+    permuted_blocks = blocks_to_update.reshape(src_shape).permute(*inv_order)
+    cache.index_copy_(0, indices, permuted_blocks)
+
+
+def kv_postprocess_blksize_and_layout_on_receive(cache, indices, block_size_ratio):
+    """
+    Transforms the layout of received KV cache to the local block_size and HND.
+    (Only works for local blocksize > remote blocksize)
+
+    prefill is HND, smaller block_size
+    decode(local) is NHD, larger block_size
+    """
+    blocks_to_update = cache.index_select(0, indices)
+
+    block_size, n_kv_heads, head_size = blocks_to_update.shape[1:]
+    remote_block_size = block_size // block_size_ratio
+    n_blocks = block_size_ratio
+
+    permuted_blocks = (
+        blocks_to_update.reshape(-1, n_blocks, n_kv_heads, remote_block_size, head_size)
+        .permute(0, 1, 3, 2, 4)
+        .flatten(1, 2)
+    )
+    cache.index_copy_(0, indices, permuted_blocks)
+
+
+def yield_req_data(
+    scheduler_output,
+) -> Iterator[tuple[str, tuple[list[int], ...] | None, bool]]:
+    """
+    Yields:
+        (req_id, new_block_id_groups, preempted)
+    """
+    # new requests
+    for req_data in scheduler_output.scheduled_new_reqs:
+        yield req_data.req_id, req_data.block_ids, False
+
+    # cached requests
+    cached_reqs = scheduler_output.scheduled_cached_reqs
+    yield from zip(
+        cached_reqs.req_ids,
+        cached_reqs.new_block_ids,
+        (req_id in cached_reqs.resumed_req_ids for req_id in cached_reqs.req_ids),
+    )
+
+
+def get_current_attn_backends(
+    vllm_config: VllmConfig, layer_names: list[str] | None = None
+) -> list[type[AttentionBackend]]:
+    """Get all distinct attention backends for the given layers.
+
+    Args:
+        vllm_config: The current vLLM configuration.
+        layer_names: Optional list of layer names to scope the lookup.
+            When None, all attention layers are considered.
+
+    Returns:
+        Deduplicated list of attention backend classes.
+    """
+    layer_type = cast(type[Any], AttentionLayerBase)
+    layers = get_layers_from_vllm_config(vllm_config, layer_type, layer_names)
+    if layers:
+        seen: dict[str, type[AttentionBackend]] = {}
+        for layer in layers.values():
+            backend = layer.get_attn_backend()
+            seen[backend.full_cls_name()] = backend
+        return list(seen.values())
+
+    # Fallback for tests, when static_forward_context is empty.
+    logger.debug(
+        "No layers found in the vLLM config. Falling back to default attention backend."
+    )
+    from vllm.v1.attention.selector import get_attn_backend
+
+    with set_current_vllm_config(vllm_config):
+        return [
+            get_attn_backend(
+                head_size=vllm_config.model_config.get_head_size(),
+                dtype=vllm_config.model_config.dtype,
+                kv_cache_dtype=vllm_config.cache_config.cache_dtype,
+                use_mla=vllm_config.model_config.use_mla,
+            )
+        ]
+
+
+def get_current_attn_backend(
+    vllm_config: VllmConfig, layer_names: list[str] | None = None
+) -> type[AttentionBackend]:
+    """Get the first attention backend for the given layers."""
+    return get_current_attn_backends(vllm_config, layer_names)[0]
+
+
+# ---- Per-engine transfer info ----
+
+
+@dataclass(frozen=True)
+class EngineTransferInfo:
+    """Common per-remote-engine transfer state, computed at handshake.
+
+    Stored per ``(engine_id, pp_rank)`` inside ``TransferTopology._engines``.
+    """
+
+    remote_tp_size: int
+
+    remote_block_len: int
+    """Block length (bytes)"""
+
+    remote_block_size: int
+    """Tokens per block."""
+
+    remote_physical_blocks_per_logical: int
+    """Physical blocks per logical block."""
+
+    remote_pp_rank: int = 0
+    """Remote producer PP rank for this engine."""
+
+    start_layer: int = 0
+    """Global index of the first layer owned by this PP rank."""
+
+    end_layer: int = 0
+    """Exclusive global index after the last layer owned by this PP rank."""
+
+
+# ---- Transfer topology ----
+
+
+@dataclass
+class TransferTopology:
+    """Single source of truth for local TP identity and per-engine remote info."""
+
+    tp_rank: int
+    tp_size: int
+    block_size: int
+    engine_id: EngineId
+    is_mla: bool
+    is_mamba: bool
+    total_num_kv_heads: int
+    attn_backends: list[type[AttentionBackend]]
+    tensor_shape: torch.Size | None = None
+
+    def __post_init__(self):
+        self.local_physical_heads = max(1, self.total_num_kv_heads // self.tp_size)
+
+        self._engines: dict[tuple[EngineId, int], EngineTransferInfo] = {}
+
+        # Figure out whether the first dimension of the cache is K/V
+        # or num_blocks.
+        attn_backend = self.attn_backends[0]
+        if not self.is_mamba:
+            _MOCK_BLOCK_SIZE = 16
+            kv_cache_shape: tuple[int, ...] = attn_backend.get_kv_cache_shape(
+                num_blocks=1,
+                block_size=_MOCK_BLOCK_SIZE,
+                num_kv_heads=1,
+                head_size=1,
+            )
+            logger.debug("Test kv_cache_shape: %s", kv_cache_shape)
+            assert kv_cache_shape[0] == 1, (
+                "KV cache layout must be blocks-first; expected mocked "
+                f"num_blocks=1 in leading dim, got shape {kv_cache_shape}."
+            )
+            if not self.is_mla:
+                assert len(kv_cache_shape) == 4, (
+                    "Attention KV cache layout must be standardized as "
+                    "[num_blocks, num_kv_heads, block_size, content_size], "
+                    f"got shape {kv_cache_shape}."
+                )
+
+        self._cross_layers_blocks = False
+        if self.tensor_shape is not None:
+            self._cross_layers_blocks = (
+                len(self.tensor_shape) == len(kv_cache_shape) + 1
+            )
+
+        if self._cross_layers_blocks:
+            logger.debug("Using cross-layer KV cache")
+            _MOCK_NUM_LAYERS = 80
+            kv_cache_shape = (_MOCK_NUM_LAYERS,) + kv_cache_shape
+            try:
+                kv_cache_stride_order = attn_backend.get_kv_cache_stride_order(
+                    include_num_layers_dimension=self._cross_layers_blocks
+                )
+            except (AttributeError, NotImplementedError):
+                assert self.tensor_shape is not None
+                kv_cache_stride_order = tuple(range(len(self.tensor_shape)))
+            kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
+
+    # ============================================================
+    # Engine registration
+    # ============================================================
+
+    def register_remote_engine(
+        self,
+        remote_engine_id: EngineId,
+        info: EngineTransferInfo,
+    ) -> EngineTransferInfo:
+        """Register a remote engine, unifying worker dicts state.
+
+        The caller (worker) is responsible for computing the info via
+        the transfer policy.  This method only stores and deduplicates.
+        """
+        assert remote_engine_id != self.engine_id, (
+            f"Cannot register local engine {self.engine_id} as remote. "
+            f"Local identity is set via __init__ params."
+        )
+        engine_key = (remote_engine_id, info.remote_pp_rank)
+        if engine_key in self._engines:
+            return self._engines[engine_key]
+        self._engines[engine_key] = info
+        return info
+
+    def get_engine_info(
+        self, remote_engine_id: EngineId, remote_pp_rank: int = 0
+    ) -> EngineTransferInfo:
+        return self._engines[(remote_engine_id, remote_pp_rank)]
+
+    def unregister_remote_engine(self, remote_engine_id: EngineId) -> None:
+        # Remove all pp_rank entries for the remote engine.
+        for key in [k for k in self._engines if k[0] == remote_engine_id]:
+            del self._engines[key]
+
+    # ============================================================
+    # Layout properties
+    # ============================================================
+
+    @property
+    def cross_layers_blocks(self) -> bool:
+        return self._cross_layers_blocks
+
+    @property
+    def virtually_split_kv_in_blocks(self) -> bool:
+        # Whether to logically split each block into two separately-indexable
+        # sub-regions. With K and V packed into the content dim, an attention
+        # block transfers as a single unit — no K/V sub-split is needed. Only
+        # Mamba still needs this, to index its two state regions (conv/ssm)
+        # separately. Not applicable to cross-layer blocks (per-layer
+        # interleaving means a simple half-split does not separate the parts).
+        return self.is_mamba and not self._cross_layers_blocks
+
+    # ============================================================
+    # Common methods
+    # ============================================================
+
+    def tp_ratio(self, remote_tp_size: int) -> int:
+        """Calculate the tensor parallel ratio between local and remote TP.
+
+        Positive when local_tp >= remote_tp (local workers read from the
+        same remote worker in groups of size ``tp_ratio``).  Negative when
+        remote_tp > local_tp (ratio is flipped).
+        """
+        if self.tp_size >= remote_tp_size:
+            assert self.tp_size % remote_tp_size == 0, (
+                f"Local tensor parallel size {self.tp_size} is not divisible "
+                f"by remote tensor parallel size {remote_tp_size}."
+            )
+            return self.tp_size // remote_tp_size
+        assert remote_tp_size % self.tp_size == 0, (
+            f"Remote tensor parallel size {remote_tp_size} is not divisible "
+            f"by local tensor parallel size {self.tp_size}."
+        )
+        return -(remote_tp_size // self.tp_size)
+
+    def block_size_ratio(self, remote_block_size: int) -> int:
+        """Calculate the block size ratio between local and remote."""
+        assert self.block_size % remote_block_size == 0, (
+            f"Local block size {self.block_size} is not divisible "
+            f"by remote block size {remote_block_size} or vice versa."
+        )
+        return self.block_size // remote_block_size
+
+    def is_kv_replicated(
+        self, remote_engine_id: EngineId, remote_pp_rank: int = 0
+    ) -> bool:
+        """Whether the KV cache is replicated across TP workers due to the
+        number of TP workers being greater than the number of KV heads.
+        """
+        return (
+            self._engines[(remote_engine_id, remote_pp_rank)].remote_tp_size
+            > self.total_num_kv_heads
+        )
+
+    def replicates_kv_cache(
+        self, remote_engine_id: EngineId, remote_pp_rank: int = 0
+    ) -> bool:
+        # MLA is always replicated as the hidden dim can't be split.
+        return self.is_mla or self.is_kv_replicated(remote_engine_id, remote_pp_rank)
+
+    @property
+    def local_replicates_kv_cache(self) -> bool:
+        """Whether the local engine's KV cache is replicated."""
+        return self.is_mla or self.tp_size > self.total_num_kv_heads
+
+    def handshake_target_ranks(self, remote_tp_size: int) -> list[int]:
+        """Pre-registration: compute which remote TP ranks to handshake with.
+
+        Pure math based on local/remote TP sizes — does not require
+        the remote engine to be registered yet.
+        """
+        tp_ratio = self.tp_ratio(remote_tp_size)
+        if tp_ratio > 0:
+            return [self.tp_rank // tp_ratio]
+        abs_ratio = -tp_ratio
+        return [self.tp_rank * abs_ratio + i for i in range(abs_ratio)]
+
+    def target_remote_ranks(
+        self, remote_engine_id: EngineId, remote_pp_rank: int = 0
+    ) -> list[int]:
+        """Get the remote TP rank(s) that the current local TP rank will
+        read from.  When remote tp_size > local tp_size, reads from
+        multiple remote ranks.
+        """
+        info = self._engines[(remote_engine_id, remote_pp_rank)]
+        tp_ratio = self.tp_ratio(info.remote_tp_size)
+        if tp_ratio > 0:
+            return [self.tp_rank // tp_ratio]
+        # remote TP > local TP: read from |tp_ratio| remote workers
+        abs_ratio = -tp_ratio
+        return [self.tp_rank * abs_ratio + i for i in range(abs_ratio)]
+
+    def get_transfer_cache_regions(
+        self, cache: torch.Tensor, layer_spec: "KVCacheSpec"
+    ) -> list[torch.Tensor] | torch.Tensor:
+        """Return the cache tensor(s) to register as NIXL memory regions,
+        also accounting for hybrid SSM models specificities.
+        """
+        if isinstance(layer_spec, MambaSpec):
+            # Register the whole kv cache shared tensor, including
+            # SSM/Conv.
+            conv, ssm = cache
+            return [conv]
+
+        # Check may be hacky but it's matching
+        # `_update_hybrid_attention_mamba_layout`.
+        if self.is_mamba and cache.shape[0] == 2:
+            # When MAMBA is present, all backends are blocks first, so
+            # that blocks can be shared between attention layers and mamba
+            # layers.  Runner already adjusted strides for FlashAttn-like
+            # backends so its num_blocks first.
+            # Swap [2<>num_blocks] dims for hybrid SSM layout.
+            cache = cache.transpose(0, 1)
+
+        # K and V are packed into one tensor (content dim), so each layer
+        # registers as a single region.
+        return [cache]
+
+    def describe(self, remote_engine_id: EngineId, remote_pp_rank: int = 0) -> str:
+        """One-line summary of transfer config for logging."""
+        info = self._engines[(remote_engine_id, remote_pp_rank)]
+        return (
+            f"TransferTopology("
+            f"tp_ratio={self.tp_ratio(info.remote_tp_size)}, "
+            f"num_kv_heads={self.total_num_kv_heads if not self.is_mla else 1}, "
+            f"local_tp={self.tp_size}, "
+            f"remote_tp={info.remote_tp_size}, "
+            f"remote_pp={remote_pp_rank}, "
+            f"local_rank={self.tp_rank}, "
+            f"remote_block_len={info.remote_block_len})"
+        )

vllm/distributed/kv_transfer/kv_connector/v1/__init__.py

@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorBase_V1,
+    KVConnectorRole,
+    SupportsHMA,
+    supports_hma,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.decode_bench_connector import (  # noqa: E501
+    DecodeBenchConnector,
+)
+
+__all__ = [
+    "KVConnectorRole",
+    "KVConnectorBase_V1",
+    "supports_hma",
+    "SupportsHMA",
+    "DecodeBenchConnector",
+]

vllm/distributed/kv_transfer/kv_connector/v1/base.py

@@ -0,0 +1,708 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+KVConnectorBase_V1 Class for Distributed KV Cache & Hidden State
+communication in vLLM v1
+
+The class provides the following primitives:
+    Scheduler-side: runs in the scheduler, binds metadata, which
+    is used by the worker-side to load/save KV cache.
+        get_num_new_matched_tokens() - get number of new tokens
+            that exist in the remote KV cache. Might be called multiple
+            times for a given request and should be side-effect free.
+        update_state_after_alloc() - update KVConnector state after
+            temporary buffer alloc by the CacheManager.
+        update_connector_output() - update KVConnector state after
+            output is received from worker-side connectors.
+        request_finished() - called once when a request is finished,
+            with the computed kv cache blocks for the request.
+            Returns whether KV cache should be freed now or if the
+            connector now assumes responsibility for freeing the
+            the blocks asynchronously. Also optionally returns KV
+            transfer params.
+        take_events() - returns new KV events that were collected
+            by the connector since the last call.
+
+    Worker-side: runs in each worker, loads/saves KV cache to/from
+    the Connector based on the metadata.
+        handle_preemptions() - called for handling preempted requests
+            or request evicted blocks before they are overwritten
+
+        start_load_kv() - starts loading all KVs (maybe async)
+        wait_for_layer_load() - blocks until layer i load is done
+
+        save_kv_layer() - starts saving KV for layer i (maybe async)
+        wait_for_save() - blocks until all saves are done
+
+        get_finished() - called with ids of finished requests, returns
+            ids of requests that have completed async sending/recving.
+        build_connector_worker_meta() - builds metadata to be sent
+            back to the scheduler-side connector
+"""
+
+import enum
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Iterable
+from typing import TYPE_CHECKING, Any, Literal
+
+import torch
+
+from vllm.logger import init_logger
+from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.outputs import KVConnectorOutput
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.distributed.kv_events import KVCacheEvent, KVConnectorKVEvents
+    from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
+        KVConnectorPromMetrics,
+        KVConnectorStats,
+        PromMetric,
+        PromMetricT,
+    )
+    from vllm.forward_context import ForwardContext
+    from vllm.v1.core.block_pool import BlockPool
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+    from vllm.v1.request import Request
+
+# s_tensor_list, d_tensor_list, s_indices, d_indices, direction
+CopyBlocksOp = Callable[
+    [
+        dict[str, torch.Tensor],
+        dict[str, torch.Tensor],
+        list[int],
+        list[int],
+        Literal["h2d", "d2h"],
+    ],
+    None,
+]
+
+logger = init_logger(__name__)
+
+
+class SupportsHMA(ABC):
+    """
+    The class that indicates the corresponding connector supports hybrid memory
+    allocator (HMA).
+    This is required to use the connector together with hybrid memory allocator.
+    """
+
+    @abstractmethod
+    def request_finished_all_groups(
+        self,
+        request: "Request",
+        block_ids: tuple[list[int], ...],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        """
+        Called exactly once when a request has finished for all kv cache groups,
+        before its blocks are freed for each group.
+
+        NOTE(Kuntai): This function is only supported by connectors that support HMA.
+
+        The connector may assumes responsibility for freeing the blocks
+        asynchronously by returning True.
+
+        Returns:
+            True if the request is being saved/sent asynchronously and blocks
+            should not be freed until the request_id is returned from
+            get_finished().
+            Optional KVTransferParams to be included in the request outputs
+            returned by the engine.
+        """
+        raise NotImplementedError
+
+
+def supports_hma(connector: Any) -> bool:
+    if isinstance(connector, type):
+        return issubclass(connector, SupportsHMA)
+    else:
+        return isinstance(connector, SupportsHMA)
+
+
+class KVConnectorRole(enum.Enum):
+    # Connector running in the scheduler process
+    SCHEDULER = 0
+
+    # Connector running in the worker process
+    WORKER = 1
+
+
+class KVConnectorHandshakeMetadata(ABC):  # noqa: B024
+    """
+    Metadata used for out of band connector handshake between
+    P/D workers. This needs to serializable.
+    """
+
+    pass
+
+
+class KVConnectorMetadata(ABC):  # noqa: B024
+    """
+    Abstract Metadata used to communicate
+    Scheduler KVConnector -> Worker KVConnector.
+    """
+
+    pass
+
+
+class KVConnectorWorkerMetadata(ABC):
+    """
+    Abstract Metadata used to communicate back
+    Worker KVConnector -> Scheduler KVConnector.
+
+    Each worker can output its own metadata.
+    For a single engine step, all metadata objects returned by workers
+    will be aggregated using the `aggregate` method below, before
+    being passed to the Scheduler KVConnector.
+    """
+
+    @abstractmethod
+    def aggregate(
+        self, other: "KVConnectorWorkerMetadata"
+    ) -> "KVConnectorWorkerMetadata":
+        """
+        Aggregate metadata with another `KVConnectorWorkerMetadata` object.
+        """
+        pass
+
+
+class KVConnectorBase_V1(ABC):
+    """
+    Base class for KV connectors.
+    """
+
+    @property
+    def prefer_cross_layer_blocks(self) -> bool:
+        """
+        Indicates whether this connector prefers KV blocks that hold KV data for all
+        layers, which can speed up KV data transfers. Defaults to False.
+        """
+        return False
+
+    def __init__(
+        self,
+        vllm_config: "VllmConfig",
+        role: KVConnectorRole,
+        kv_cache_config: "KVCacheConfig",
+    ):
+        logger.warning(
+            "Initializing KVConnectorBase_V1. This API is experimental and "
+            "subject to change in the future as we iterate the design."
+        )
+        self._connector_metadata: KVConnectorMetadata | None = None
+        self._vllm_config = vllm_config
+        if vllm_config.kv_transfer_config is not None:
+            self._kv_transfer_config = vllm_config.kv_transfer_config
+        else:
+            raise ValueError("kv_transfer_config must be set for KVConnectorBase_V1")
+        self._kv_cache_config = kv_cache_config
+        self._role = role
+
+    @property
+    def role(self) -> KVConnectorRole:
+        return self._role
+
+    # ==============================
+    # Worker-side methods
+    # ==============================
+
+    def bind_connector_metadata(self, connector_metadata: KVConnectorMetadata) -> None:
+        """Set the connector metadata from the scheduler.
+
+        This function should be called by the model runner every time
+        before the model execution. The metadata will be used for runtime
+        KV cache loading and saving.
+
+        Args:
+            connector_metadata (dict): the connector metadata.
+        """
+        self._connector_metadata = connector_metadata
+
+    def clear_connector_metadata(self) -> None:
+        """Clear the connector metadata.
+
+        This function should be called by the model runner every time
+        after the model execution.
+        """
+        self._connector_metadata = None
+
+    def _get_connector_metadata(self) -> KVConnectorMetadata:
+        """Get the connector metadata.
+
+        This function should only be called inside the connector.
+
+        Returns:
+            ConnectorMetadata: the connector metadata.
+        """
+        # Should only be called while set to valid metadata.
+        assert self._connector_metadata is not None
+        return self._connector_metadata
+
+    def has_connector_metadata(self) -> bool:
+        """Check whether the connector metadata is currently set.
+
+        Returns:
+            bool: True if connector metadata exists, False otherwise.
+        """
+        return self._connector_metadata is not None
+
+    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
+        """
+        Initialize with the KV caches. Useful for pre-registering the
+        KV Caches in the KVConnector (e.g. for NIXL).
+
+        Args:
+            kv_caches: dictionary of layer names, kv cache
+        """
+        return
+
+    def register_cross_layers_kv_cache(
+        self, kv_cache: torch.Tensor, attn_backend: type["AttentionBackend"]
+    ):
+        """
+        Initialize with a single KV cache tensor used by all layers.
+        The first dimension should be num_layers.
+        This function will only be called for models with uniform layers,
+        and only if the prefers_cross_layer_blocks is set to True.
+        Only one of the functions
+        {register_kv_caches, register_cross_layers_kv_cache} will be called.
+
+        Args:
+            kv_cache: a cross-layers kv cache tensor
+            attn_backend: The attention backend that corresponds to all layers
+        """
+        return
+
+    def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
+        """
+        Set the xPU-specific ops for copying KV between host and device.
+        Needed when host buffer is used for kv transfer (e.g., in NixlConnector)
+        """
+        return
+
+    def handle_preemptions(self, kv_connector_metadata: KVConnectorMetadata):
+        """
+        Handle preempted requests or evicted blocks BEFORE they are overwritten.
+        Needed for connectors which use async saves (e.g., OffloadingConnector)
+        """
+        return
+
+    @abstractmethod
+    def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None:
+        """
+        Start loading the KV cache from the connector to vLLM's paged
+        KV buffer. This is called from the forward context before the
+        forward pass to enable async loading during model execution.
+
+        Args:
+            forward_context (ForwardContext): the forward context.
+            **kwargs: additional arguments for the load operation
+
+        Note:
+            The number of elements in kv_caches and layer_names should be
+            the same.
+
+        """
+        pass
+
+    @abstractmethod
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        """
+        Block until the KV for a specific layer is loaded into vLLM's
+        paged buffer. This is called from within attention layer to ensure
+        async copying from start_load_kv is complete.
+
+        This interface will be useful for layer-by-layer pipelining.
+
+        Args:
+            layer_name: the name of that layer
+        """
+        pass
+
+    @abstractmethod
+    def save_kv_layer(
+        self,
+        layer_name: str,
+        kv_layer: torch.Tensor,
+        attn_metadata: "AttentionMetadata",
+        **kwargs: Any,
+    ) -> None:
+        """
+        Start saving a layer of KV cache from vLLM's paged buffer
+        to the connector. This is called from within attention layer to
+        enable async copying during execution.
+
+        Args:
+            layer_name (str): the name of the layer.
+            kv_layer (torch.Tensor): the paged KV buffer of the current
+                layer in vLLM.
+            attn_metadata (AttentionMetadata): the attention metadata.
+            **kwargs: additional arguments for the save operation.
+        """
+        pass
+
+    @abstractmethod
+    def wait_for_save(self):
+        """
+        Block until all the save operations is done. This is called
+        as the forward context exits to ensure that the async saving
+        from save_kv_layer is complete before finishing the forward.
+
+        This prevents overwrites of paged KV buffer before saving done.
+        """
+        pass
+
+    def get_finished(
+        self, finished_req_ids: set[str]
+    ) -> tuple[set[str] | None, set[str] | None]:
+        """
+        Notifies worker-side connector ids of requests that have
+        finished generating tokens on the worker.
+        The scheduler process (via the Executors) will use this output
+        to track which workers are done.
+
+        Returns:
+            ids of requests that have finished asynchronous transfer
+            (requests that previously returned True from request_finished()),
+            tuple of (sending/saving ids, recving/loading ids).
+            The finished saves/sends req ids must belong to a set provided in a
+            call to this method (this call or a prior one).
+        """
+        return None, None
+
+    def get_block_ids_with_load_errors(self) -> set[int]:
+        """
+        Get the set of block IDs that failed to load.
+
+        Returns:
+            Set of block IDs that encountered load errors.
+            Empty set if no load errors occurred.
+
+        Notes:
+            - Applies to both sync- and async-loading requests.
+            - Async loading: failed blocks may be reported in any forward pass
+              up to and including the pass where the request ID is returned by
+              `get_finished()`. Even if failures occur, the request must still
+              be reported via `get_finished()`, and the failed block IDs must
+              appear here no later than that same pass.
+            - Sync loading: failed blocks should be reported in the forward
+              pass in which they are detected.
+        """
+        return set()
+
+    def shutdown(self):
+        """
+        Shutdown the connector. This is called when the worker process
+        is shutting down to ensure that all the async operations are
+        completed and the connector is cleaned up properly.
+        """
+        return None
+
+    def get_kv_connector_stats(self) -> "KVConnectorStats | None":
+        """
+        Get the KV connector stats collected during the last interval.
+        """
+        return None
+
+    def get_kv_connector_kv_cache_events(self) -> "KVConnectorKVEvents | None":
+        """
+        Get the KV connector kv cache events collected during the last interval.
+        This function should be called by the model runner every time after the
+        model execution and before cleanup.
+        """
+        return None
+
+    def get_handshake_metadata(self) -> KVConnectorHandshakeMetadata | None:
+        """
+        Get the KVConnector handshake metadata for this connector.
+        This metadata is used for out-of-band connector handshake
+        between P/D workers.
+
+        Returns:
+            KVConnectorHandshakeMetadata: the handshake metadata.
+            None if no handshake metadata is available.
+        """
+        return None
+
+    def build_connector_worker_meta(self) -> KVConnectorWorkerMetadata | None:
+        """
+        Build the KVConnector worker metadata for this engine step.
+
+        Returns:
+            KVConnectorWorkerMetadata: the worker metadata.
+            None if no worker metadata is available.
+        """
+        return None
+
+    # ==============================
+    # Scheduler-side methods
+    # ==============================
+
+    def bind_gpu_block_pool(self, gpu_block_pool: "BlockPool") -> None:
+        """
+        Bind the GPU block pool to the connector for per-GPU block status tracking.
+        For example, inc/dec ref counts, or iterate over the prefix cache blocks.
+
+        Args:
+            gpu_block_pool: the GPU block pool.
+        """
+        return
+
+    @abstractmethod
+    def get_num_new_matched_tokens(
+        self,
+        request: "Request",
+        num_computed_tokens: int,
+    ) -> tuple[int | None, bool]:
+        """
+        Get number of new tokens that can be loaded from the
+        external KV cache beyond the num_computed_tokens.
+
+        Args:
+            request (Request): the request object.
+            num_computed_tokens (int): the number of locally
+                computed tokens for this request
+
+        Returns:
+            A tuple with the following elements:
+                - An optional number of tokens that can be loaded from the
+                  external KV cache beyond what is already computed.
+                  If None, it means that the connector needs more time to
+                  determine the number of matched tokens, and the scheduler
+                  should query for this request again later.
+                - `True` if external KV cache tokens will be loaded
+                  asynchronously (between scheduler steps). Must be
+                  'False' if the first element is 0.
+
+        Notes:
+            The connector should only consider the largest prefix of prompt-
+            tokens for which KV cache is actually available at the time of the
+            call. If the cache cannot be loaded for some tokens (e.g., due to
+            connectivity issues or eviction), those tokens must not be taken
+            into account.
+        """
+        pass
+
+    @abstractmethod
+    def update_state_after_alloc(
+        self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
+    ):
+        """
+        Update KVConnector state after block allocation.
+
+        If get_num_new_matched_tokens previously returned True for a
+        request, this function may be called twice for that same request -
+        first when blocks are allocated for the connector tokens to be
+        asynchronously loaded into, and second when any additional blocks
+        are allocated, after the load/transfer is complete.
+
+        Decide whether to load based on ``num_external_tokens``, not on
+        whether ``blocks`` is empty: ``blocks`` may be non-empty even when
+        ``num_external_tokens == 0`` (e.g. a non-chosen sub-connector of
+        MultiConnector still receives the request's real blocks).
+
+        Args:
+            request (Request): the request object.
+            blocks (KVCacheBlocks): the blocks allocated for the request.
+            num_external_tokens (int): the number of tokens to load from the
+                external KV cache. 0 means nothing should be loaded.
+        """
+        pass
+
+    @abstractmethod
+    def build_connector_meta(
+        self, scheduler_output: SchedulerOutput
+    ) -> KVConnectorMetadata:
+        """
+        Build the connector metadata for this step.
+
+        This function should NOT modify fields in the scheduler_output.
+        Also, calling this function will reset the state of the connector.
+
+        Args:
+            scheduler_output (SchedulerOutput): the scheduler output object.
+        """
+        pass
+
+    def on_new_request(self, request: "Request") -> None:
+        """Called by the scheduler when a new request is added.
+
+        Connectors can override this to inspect the request and perform
+        bookkeeping. The default implementation is a no-op.
+        """
+        return
+
+    def update_connector_output(self, connector_output: KVConnectorOutput):
+        """
+        Update KVConnector state from worker-side connectors output.
+
+        Args:
+            connector_output (KVConnectorOutput): the worker-side
+                connectors output.
+        """
+        return
+
+    def request_finished(
+        self,
+        request: "Request",
+        block_ids: list[int],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        """
+        Called exactly once when a request has finished, before its blocks are
+        freed.
+
+        The connector may assumes responsibility for freeing the blocks
+        asynchronously by returning True.
+
+        Returns:
+            True if the request is being saved/sent asynchronously and blocks
+            should not be freed until the request_id is returned from
+            get_finished().
+            Optional KVTransferParams to be included in the request outputs
+            returned by the engine.
+        """
+        return False, None
+
+    def take_events(self) -> Iterable["KVCacheEvent"]:
+        """
+        Take the KV cache events from the connector.
+
+        Yields:
+            New KV cache events since the last call.
+        """
+        return ()
+
+    def has_pending_push_work(self) -> bool:
+        """Return True if the connector has push-mode work that requires
+        the engine main loop to keep stepping (e.g. a P-side request whose
+        KV blocks are waiting to be WRITTEN to a D node).
+
+        Connectors that don't implement push-based KV transfer should
+        leave this as False.
+        """
+        # TODO: replace with a more general connector hook for keeping the
+        # scheduler alive (e.g. extend has_unfinished_requests).
+        return False
+
+    @classmethod
+    def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None:
+        """
+        Get the required KV cache layout for this connector.
+        Args:
+            vllm_config (VllmConfig): the vllm config.
+
+        Returns:
+            str: the required KV cache layout. e.g. HND, or NHD.
+            None if the connector does not require a specific layout.
+        """
+
+        if cls is KVConnectorBase_V1:
+            raise TypeError(
+                "get_required_kvcache_layout should not be called "
+                "on the abstract base class"
+            )
+        return None
+
+    @classmethod
+    def requires_piecewise_for_cudagraph(cls, extra_config: dict[str, Any]) -> bool:
+        """
+        Check if this connector requires PIECEWISE CUDA graph mode.
+
+        Connectors that use asynchronous layer-by-layer operations
+        (wait_for_layer_load/save_kv_layer) should override this method
+        to return True when those operations are enabled. These operations
+        cannot be captured in CUDA graphs and will be skipped during replay,
+        causing data races. PIECEWISE mode allows Python code to execute
+        between graph pieces, ensuring proper synchronization.
+
+        Args:
+            extra_config: The kv_connector_extra_config dict from
+                KVTransferConfig.
+
+        Returns:
+            True if this connector requires PIECEWISE CUDA graph mode,
+            False otherwise.
+        """
+        return False
+
+    def get_finished_count(self) -> int | None:
+        """
+        Get the count of requests expected to complete send/receive operations
+        via this connector. This method is used to initialize the
+        KVOutputAggregator, overwriting the default world_size.
+
+        Returns:
+            int: expected sending or receiving completion count.
+        """
+
+        return None
+
+    @classmethod
+    def build_kv_connector_stats(
+        cls, data: dict[str, Any] | None = None
+    ) -> "KVConnectorStats | None":
+        """
+        KVConnectorStats resolution method. This method allows dynamically
+        registered connectors to return their own KVConnectorStats object,
+        which can implement custom aggregation logic on the data dict.
+        """
+        return None
+
+    def set_xfer_handshake_metadata(
+        self, metadata: dict[int, KVConnectorHandshakeMetadata]
+    ) -> None:
+        """
+        Set the KV connector handshake metadata for this connector.
+
+        Args:
+            metadata (KVConnectorHandshakeMetadata): the handshake metadata to set.
+        """
+        return None
+
+    def set_xfer_handshake_metadata_pp_aware(
+        self, metadata: dict[tuple[int, int], KVConnectorHandshakeMetadata]
+    ) -> None:
+        """
+        Set handshake metadata keyed by (pp_rank, tp_rank).
+        - Default implementation assumes pp_rank is always 0
+        - PP-aware connectors override this to consume all PP producer shards.
+        """
+        if any(pp_rank != 0 for pp_rank, _ in metadata):
+            raise ValueError(
+                f"{type(self).__name__} received pp_rank > 0 handshake metadata "
+                "but does not support PP-disaggregated KV transfer."
+            )
+        self.set_xfer_handshake_metadata(
+            {tp_rank: meta for (_, tp_rank), meta in metadata.items()}
+        )
+
+    @classmethod
+    def build_prom_metrics(
+        cls,
+        vllm_config: "VllmConfig",
+        metric_types: dict[type["PromMetric"], type["PromMetricT"]],
+        labelnames: list[str],
+        per_engine_labelvalues: dict[int, list[object]],
+    ) -> "KVConnectorPromMetrics | None":
+        """
+        Create a KVConnectorPromMetrics subclass which should register
+        per-connector Prometheus metrics and implement observe() to
+        expose connector transfer stats via Prometheus.
+        """
+        return None
+
+    def reset_cache(self) -> bool | None:
+        """
+        Reset the connector's internal cache.
+
+        Returns:
+            bool: True if the cache was successfully reset, False otherwise.
+        """
+        logger.debug(
+            "Connector cache reset requested, but %s does not implement reset_cache().",
+            type(self).__name__,
+        )
+
+        return None

vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py

@@ -0,0 +1,480 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+DecodeBenchConnector: A KV Connector for decode instance performance testing.
+
+This connector emulates a prefill-decode disaggregated setting by filling
+the KV cache with dummy values, allowing measurement of decoder performance
+under larger input sequence lengths (ISL) in resource-limited environments.
+
+Usage:
+    To use this connector for benchmarking, configure it in the kv_transfer_config:
+
+    Example:
+        vllm serve <model> --kv-transfer-config '{
+            "kv_connector": "DecodeBenchConnector",
+            "kv_role": "kv_both",
+            "kv_connector_extra_config": {
+                "fill_mean": 0.015,
+                "fill_std": 0.0
+            }
+        }'
+
+    Then run your benchmark with desired input/output lengths:
+        vllm bench serve --base-url http://127.0.0.1:8000 --model <model> \\
+            --dataset-name random --random-input-len 40000 \\
+            --random-output-len 100 --max-concurrency 10
+
+    Configuration options (via kv_connector_extra_config):
+        - fill_mean (float): Mean value for random normal fill (default: 0.015)
+        - fill_std (float): Standard deviation for random fill (default: 0.0)
+          Set to 0 for constant values, >0 for random sampling
+"""
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from vllm.distributed.kv_transfer.kv_connector.v1 import (
+    KVConnectorBase_V1,
+    KVConnectorRole,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorMetadata,
+    SupportsHMA,
+)
+from vllm.logger import init_logger
+from vllm.utils.math_utils import cdiv
+from vllm.v1.attention.backend import AttentionMetadata
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.forward_context import ForwardContext
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+    from vllm.v1.core.sched.output import SchedulerOutput
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+    from vllm.v1.request import Request
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class DecodeBenchConnectorMetadata(KVConnectorMetadata):
+    """Metadata for DecodeBenchConnector.
+
+    Contains information about which requests need their KV cache filled
+    with dummy values for benchmarking purposes.
+    """
+
+    # request_id -> (block_ids_per_group, num_tokens_to_fill)
+    # block_ids_per_group is a tuple of lists, one per KV cache group
+    # For standard attention: single group, e.g., ([1, 2, 3],)
+    # For MLA: multiple groups, e.g., ([1, 2], [1, 2])
+    reqs_to_fill: dict[str, tuple[tuple[list[int], ...], int]]
+
+
+class DecodeBenchConnector(KVConnectorBase_V1, SupportsHMA):
+    """
+    A KV Connector for decode instance performance testing.
+
+    This connector fills the KV cache with dummy (non-zero) values to
+    emulate a prefill-decode disaggregated setting, enabling performance
+    testing of the decoder with larger input sequence lengths.
+    """
+
+    def __init__(
+        self,
+        vllm_config: "VllmConfig",
+        role: KVConnectorRole,
+        kv_cache_config: "KVCacheConfig",
+    ):
+        super().__init__(vllm_config, role, kv_cache_config)
+
+        self.connector_scheduler: DecodeBenchConnectorScheduler | None = None
+        self.connector_worker: DecodeBenchConnectorWorker | None = None
+
+        if role == KVConnectorRole.SCHEDULER:
+            self.connector_scheduler = DecodeBenchConnectorScheduler(vllm_config)
+        elif role == KVConnectorRole.WORKER:
+            self.connector_worker = DecodeBenchConnectorWorker(vllm_config)
+
+    # ==============================
+    # Worker-side methods
+    # ==============================
+
+    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
+        assert self.connector_worker is not None
+        self.connector_worker.register_kv_caches(kv_caches)
+
+    def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None:
+        assert self.connector_worker is not None
+        assert isinstance(self._connector_metadata, DecodeBenchConnectorMetadata)
+        self.connector_worker.start_fill_kv(self._connector_metadata)
+
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        # All operations are synchronous, so nothing to wait for
+        pass
+
+    def save_kv_layer(
+        self,
+        layer_name: str,
+        kv_layer: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        **kwargs: Any,
+    ) -> None:
+        # This connector doesn't save KV cache (benchmarking only)
+        pass
+
+    def wait_for_save(self):
+        # This connector doesn't save KV cache (benchmarking only)
+        pass
+
+    # ==============================
+    # Scheduler-side methods
+    # ==============================
+
+    def get_num_new_matched_tokens(
+        self,
+        request: "Request",
+        num_computed_tokens: int,
+    ) -> tuple[int | None, bool]:
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.get_num_new_matched_tokens(
+            request, num_computed_tokens
+        )
+
+    def update_state_after_alloc(
+        self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
+    ):
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.update_state_after_alloc(
+            request, blocks, num_external_tokens
+        )
+
+    def build_connector_meta(
+        self, scheduler_output: "SchedulerOutput"
+    ) -> KVConnectorMetadata:
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.build_connector_meta(scheduler_output)
+
+    def request_finished(
+        self,
+        request: "Request",
+        block_ids: list[int],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        assert self.connector_scheduler is not None
+        self.connector_scheduler.request_finished(request)
+        return False, None
+
+    def request_finished_all_groups(
+        self,
+        request: "Request",
+        block_ids: tuple[list[int], ...],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        # HMA-enabled path: same cleanup as the single-group variant since
+        # this connector owns no external state per block.
+        assert self.connector_scheduler is not None
+        self.connector_scheduler.request_finished(request)
+        return False, None
+
+
+class DecodeBenchConnectorScheduler:
+    """Scheduler-side implementation for DecodeBenchConnector."""
+
+    def __init__(self, vllm_config: "VllmConfig"):
+        self.vllm_config = vllm_config
+        self.block_size = vllm_config.cache_config.block_size
+
+        # Track which requests have already been filled
+        self._filled_requests: set[str] = set()
+
+        # Track pending fills for the current scheduler step
+        # request_id -> (block_ids_per_group, num_tokens_to_fill)
+        # Note: _pending_fills doesn't need explicit cleanup - it's cleared
+        # after build_connector_meta() is called in the same scheduler step
+        self._pending_fills: dict[str, tuple[tuple[list[int], ...], int]] = {}
+
+    def get_num_new_matched_tokens(
+        self,
+        request: "Request",
+        num_computed_tokens: int,
+    ) -> tuple[int, bool]:
+        """
+        For new requests, return the number of tokens that should be filled
+        with dummy KV cache values.
+
+        Returns:
+            (num_tokens_to_fill, is_async)
+            - num_tokens_to_fill: number of uncomputed tokens minus 1
+                (we fill everything except the last token for decode)
+            - is_async: False (synchronous filling)
+        """
+        req_id = request.request_id
+
+        # Only fill once per request on first scheduling
+        if req_id in self._filled_requests:
+            return 0, False
+
+        # Calculate how many tokens we need to fill
+        # Fill all uncomputed tokens except the last one (which will be decoded)
+        # This simulates having processed a long prefill
+        num_uncomputed_tokens = request.num_tokens - num_computed_tokens
+        num_tokens_to_fill = max(0, num_uncomputed_tokens - 1)
+
+        if num_tokens_to_fill == 0:
+            return 0, False
+
+        # Return False for synchronous operation - the fill is fast enough
+        # that async overhead isn't worth it
+        return num_tokens_to_fill, False
+
+    def update_state_after_alloc(
+        self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
+    ):
+        """
+        Called after blocks are allocated. Store the block IDs so we can
+        fill them with dummy values.
+
+        Supports both standard attention (single KV cache group) and MLA
+        (multiple KV cache groups).
+        """
+        req_id = request.request_id
+
+        if num_external_tokens == 0:
+            return
+
+        # Get the block IDs that were allocated
+        # block_groups is a tuple of lists, one per KV cache group
+        # For standard attention: 1 group
+        # For MLA: multiple groups (one per attention type)
+        block_groups = blocks.get_block_ids()
+
+        # Calculate how many blocks we need to fill
+        # num_external_tokens are the tokens we said we'd provide
+        num_blocks_to_fill = cdiv(num_external_tokens, self.block_size)
+
+        # Extract the first num_blocks_to_fill blocks from each group
+        # All groups should have the same block IDs for the same request
+        block_ids_per_group = tuple(
+            group_blocks[:num_blocks_to_fill] for group_blocks in block_groups
+        )
+
+        # Store the blocks to fill for all group. _pending_fills doesn't need cleanup
+        # as it's cleared after build_connector_meta
+        self._pending_fills[req_id] = (
+            block_ids_per_group,
+            num_external_tokens,
+        )
+        self._filled_requests.add(req_id)
+
+        logger.debug(
+            "DecodeBenchConnector: Allocated %d blocks across %d KV cache groups "
+            "for request %s",
+            num_blocks_to_fill,
+            len(block_groups),
+            req_id,
+        )
+
+    def build_connector_meta(
+        self, scheduler_output: "SchedulerOutput"
+    ) -> KVConnectorMetadata:
+        """
+        Build metadata containing information about which blocks to fill
+        with dummy KV values.
+        """
+        meta = DecodeBenchConnectorMetadata(reqs_to_fill=self._pending_fills.copy())
+
+        # Clear pending fills after building metadata
+        self._pending_fills.clear()
+
+        return meta
+
+    def request_finished(self, request: "Request"):
+        """
+        Called when a request has finished. Clean up any state.
+        """
+        self._filled_requests.discard(request.request_id)
+
+
+class DecodeBenchConnectorWorker:
+    """Worker-side implementation for DecodeBenchConnector."""
+
+    def __init__(self, vllm_config: "VllmConfig"):
+        self.vllm_config = vllm_config
+        self.block_size = vllm_config.cache_config.block_size
+
+        # Get fill parameters from extra config
+        kv_transfer_config = vllm_config.kv_transfer_config
+        assert kv_transfer_config is not None
+        self.fill_mean = kv_transfer_config.get_from_extra_config("fill_mean", 0.015)
+        self.fill_std = kv_transfer_config.get_from_extra_config("fill_std", 0.0)
+
+        # Will be populated via register_kv_caches
+        self.kv_caches: dict[str, torch.Tensor] | None = None
+
+        # Mapping from KV cache group index to list of layer names in that group
+        self.group_to_layers: dict[int, list[str]] | None = None
+
+    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
+        """Store references to the KV cache tensors and build group mapping."""
+        self.kv_caches = kv_caches
+
+        # For simplicity, assume all layers belong to group 0 (standard attention)
+        # For MLA models with multiple groups, the metadata will handle the mapping
+        # We just need to fill the blocks specified in the metadata
+        self.group_to_layers = {0: list(kv_caches.keys())}
+
+        logger.debug(
+            "DecodeBenchConnector: Registered %d KV cache layers",
+            len(kv_caches),
+        )
+
+    def start_fill_kv(self, metadata: DecodeBenchConnectorMetadata):
+        """
+        Fill the allocated KV cache blocks with dummy (non-zero) values.
+
+        This simulates having a populated KV cache from a prefill phase,
+        allowing decode performance testing with larger context sizes.
+
+        Supports both standard attention (single group) and MLA (multiple groups).
+        """
+        if not metadata.reqs_to_fill:
+            return
+
+        assert self.kv_caches is not None, "KV caches must be registered before filling"
+        assert self.group_to_layers is not None, "Group mapping must be initialized"
+
+        for req_id, (block_ids_per_group, num_tokens) in metadata.reqs_to_fill.items():
+            # Fill blocks for each KV cache group
+            for group_idx, block_ids in enumerate(block_ids_per_group):
+                self._fill_blocks(group_idx, block_ids, num_tokens)
+
+            logger.debug(
+                "DecodeBenchConnector: Filled %d blocks (%d tokens) across %d groups "
+                "for request %s",
+                len(block_ids_per_group[0]) if block_ids_per_group else 0,
+                num_tokens,
+                len(block_ids_per_group),
+                req_id,
+            )
+
+    def _fill_blocks(self, group_idx: int, block_ids: list[int], num_tokens: int):
+        """
+        Fill specified blocks with dummy non-zero values for a specific KV cache group.
+
+        Args:
+            group_idx: The KV cache group index to fill
+            block_ids: List of block IDs to fill in this group
+            num_tokens: Total number of tokens to fill across these blocks
+        """
+        if not block_ids:
+            return
+
+        assert self.kv_caches is not None
+        assert self.group_to_layers is not None
+
+        # Get the layers that belong to this group
+        layer_names = self.group_to_layers.get(group_idx, [])
+
+        # Fill only the layers in this group
+        for layer_name in layer_names:
+            if layer_name not in self.kv_caches:
+                logger.warning(
+                    "DecodeBenchConnector: Layer %s not found in KV caches", layer_name
+                )
+                continue
+
+            kv_cache = self.kv_caches[layer_name]
+
+            # Attention layers store KV as a single block-indexed tensor whose
+            # first dim is num_blocks; fill the requested block rows. Hybrid /
+            # linear-attention layers (e.g. Mamba, Kimi Delta Attention) store
+            # their state as a list/tuple of tensors that are NOT block-indexed
+            # — each tensor is a single state buffer with no num_blocks
+            # dimension — so fill each tensor in its entirety with the same
+            # dummy values.
+            if isinstance(kv_cache, torch.Tensor):
+                self._fill_block_tensor(kv_cache, block_ids)
+            elif isinstance(kv_cache, (list, tuple)) and all(
+                isinstance(t, torch.Tensor) for t in kv_cache
+            ):
+                for state_tensor in kv_cache:
+                    self._fill_state_tensor(state_tensor)
+            else:
+                logger.warning_once(
+                    "DecodeBenchConnector: skipping fill for layer %s whose KV "
+                    "cache is %s, not a tensor or a list/tuple of tensors.",
+                    layer_name,
+                    type(kv_cache).__name__,
+                )
+                continue
+
+        logger.debug(
+            "DecodeBenchConnector: Filled %d blocks in group %d with %s values "
+            "(mean=%.3f, std=%.3f)",
+            len(block_ids),
+            group_idx,
+            "random" if self.fill_std > 0 else "constant",
+            self.fill_mean,
+            self.fill_std,
+        )
+
+    def _fill_block_tensor(self, kv_cache: torch.Tensor, block_ids: list[int]):
+        """Fill the requested block rows of a block-indexed KV cache tensor.
+
+        Args:
+            kv_cache: A KV cache tensor whose first dim is num_blocks.
+            block_ids: Block IDs to fill. IDs that are out of range for this
+                tensor's first dim are ignored.
+        """
+        # Convert block_ids to tensor on device
+        block_ids_tensor = torch.tensor(
+            block_ids, dtype=torch.long, device=kv_cache.device
+        )
+
+        # Filter invalid block IDs
+        valid_mask = block_ids_tensor < kv_cache.shape[0]
+        valid_block_ids = block_ids_tensor[valid_mask]
+
+        if len(valid_block_ids) == 0:
+            return
+
+        # Create fill values - either constant or random
+        block_shape = kv_cache.shape[1:]
+        if self.fill_std > 0:
+            # Random normal sampling
+            fill_values = torch.normal(
+                mean=self.fill_mean,
+                std=self.fill_std,
+                size=(len(valid_block_ids),) + block_shape,
+                dtype=kv_cache.dtype,
+                device=kv_cache.device,
+            )
+        else:
+            # Constant fill value
+            fill_values = torch.full(
+                (len(valid_block_ids),) + block_shape,
+                self.fill_mean,
+                dtype=kv_cache.dtype,
+                device=kv_cache.device,
+            )
+
+        # Batch fill operation
+        kv_cache[valid_block_ids] = fill_values
+
+    def _fill_state_tensor(self, kv_cache: torch.Tensor):
+        """Fill an entire non-block-indexed state tensor with dummy values.
+
+        Hybrid / linear-attention layers (e.g. Mamba, Kimi Delta Attention)
+        store their per-layer state as tensors with no num_blocks dimension,
+        so the whole tensor is filled with the same constant or random values
+        used for block fills, rather than selected block rows.
+
+        Args:
+            kv_cache: A state tensor to fill in its entirety.
+        """
+        if self.fill_std > 0:
+            kv_cache.normal_(mean=self.fill_mean, std=self.fill_std)
+        else:
+            kv_cache.fill_(self.fill_mean)

vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py

@@ -0,0 +1,444 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import os
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+import safetensors
+import torch
+
+from vllm.config import VllmConfig
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorBase_V1,
+    KVConnectorMetadata,
+    KVConnectorRole,
+)
+from vllm.logger import init_logger
+from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata
+from vllm.utils.hashing import safe_hash
+from vllm.v1.attention.backend import AttentionMetadata
+from vllm.v1.core.sched.output import SchedulerOutput
+
+if TYPE_CHECKING:
+    from vllm.forward_context import ForwardContext
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+    from vllm.v1.request import Request
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class ReqMeta:
+    # Request tokens
+    token_ids: torch.Tensor
+    # Slot mappings, should have the same length as token_ids
+    slot_mapping: torch.Tensor
+    # Is store or load
+    is_store: bool
+    mm_hashes: list[str]
+
+    @staticmethod
+    def make_meta(
+        token_ids: list[int],
+        block_ids: list[int],
+        block_size: int,
+        is_store: bool,
+        mm_hashes: list[str],
+    ) -> "ReqMeta":
+        valid_num_tokens = align_to_block_size(len(token_ids), block_size)
+        token_ids_tensor = torch.tensor(token_ids)[:valid_num_tokens]
+        block_ids_tensor = torch.tensor(block_ids)
+        num_blocks = block_ids_tensor.shape[0]
+        block_offsets = torch.arange(0, block_size)
+        slot_mapping = (
+            block_offsets.reshape((1, block_size))
+            + block_ids_tensor.reshape((num_blocks, 1)) * block_size
+        )
+        slot_mapping = slot_mapping.flatten()[:valid_num_tokens]
+        return ReqMeta(
+            token_ids=token_ids_tensor,
+            slot_mapping=slot_mapping,
+            is_store=is_store,
+            mm_hashes=mm_hashes,
+        )
+
+
+@dataclass
+class ExampleConnectorMetadata(KVConnectorMetadata):
+    requests: list[ReqMeta] = field(default_factory=list)
+
+    def add_request(
+        self,
+        token_ids: list[int],
+        block_ids: list[int],
+        block_size: int,
+        is_store: bool,
+        mm_hashes: list[str],
+    ) -> None:
+        self.requests.append(
+            ReqMeta.make_meta(token_ids, block_ids, block_size, is_store, mm_hashes)
+        )
+
+
+class ExampleConnector(KVConnectorBase_V1):
+    # NOTE: This is Simple debug implementation of the KV connector.
+    # It save / load the KV cache to / from the disk.
+    # It does extra work which will overwrite the existing prefix-cache in GPU
+    # - to remove the overhead, need to add some "mask" in the ReqMeta class
+
+    def __init__(
+        self,
+        vllm_config: "VllmConfig",
+        role: KVConnectorRole,
+        kv_cache_config: "KVCacheConfig",
+    ):
+        super().__init__(
+            vllm_config=vllm_config,
+            role=role,
+            kv_cache_config=kv_cache_config,
+        )
+        self._block_size = vllm_config.cache_config.block_size
+        self._requests_need_load: dict[str, Request] = {}
+        self._storage_path = self._kv_transfer_config.get_from_extra_config(
+            "shared_storage_path", "/tmp"
+        )
+        logger.info(self._kv_transfer_config)
+        logger.info("Shared storage path is %s", self._storage_path)
+
+    def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None:
+        """Start loading the KV cache from the connector buffer to vLLM's
+        paged KV buffer.
+
+        Args:
+            forward_context (ForwardContext): the forward context.
+            **kwargs: additional arguments for the load operation
+
+        Note:
+            The number of elements in kv_caches and layer_names should be
+            the same.
+        """
+
+        def inject_kv_into_layer(
+            dst_kv_cache_layer: torch.Tensor,
+            src_kv_cache: torch.Tensor,
+            slot_mapping: torch.Tensor,
+            attn_metadata: AttentionMetadata,
+        ) -> None:
+            """Inject the KV cache into the layer.
+
+            Args:
+                dst_kv_cache_layer (torch.Tensor): the destination KV cache
+                    layer. In shape [num_pages, page_size, xxx] for MLA,
+                    [num_pages, 2, page_size, xxx] otherwise.
+                src_kv_cache (torch.Tensor): the source KV cache.
+                slot_mapping (torch.Tensor): the slot mapping. In shape
+                    [num_tokens].
+            """
+            if isinstance(attn_metadata, MLACommonMetadata):
+                dst_kv_cache_layer_shape = dst_kv_cache_layer.shape
+                num_pages = dst_kv_cache_layer_shape[0]
+                page_size = dst_kv_cache_layer_shape[1]
+                dst_kv_cache_layer = dst_kv_cache_layer.reshape(
+                    num_pages * page_size, -1
+                )
+                dst_kv_cache_layer[slot_mapping, ...] = src_kv_cache
+            else:
+                block_idxs = slot_mapping // self._block_size
+                offsets = slot_mapping % self._block_size
+                dst_kv_cache_layer[block_idxs, :, offsets] = src_kv_cache
+
+        # Get the metadata
+        metadata: KVConnectorMetadata = self._get_connector_metadata()
+        assert isinstance(metadata, ExampleConnectorMetadata)
+
+        attn_metadata = forward_context.attn_metadata
+        if attn_metadata is None:
+            logger.warning("In connector.start_load_kv, but the attn_metadata is None")
+            return
+
+        # Load the KV for each request each layer
+        for request in metadata.requests:
+            if request.is_store:
+                continue
+            logger.info(
+                "Inject KV cache of %d tokens to the paged memory",
+                len(request.slot_mapping),
+            )
+            for layer_name in forward_context.no_compile_layers:
+                layer = forward_context.no_compile_layers[layer_name]
+
+                # Only process layers that have kv_cache
+                # attribute (attention layers) Skip non-attention
+                # layers like FusedMoE/MLP etc.
+                kv_cache_layer = getattr(layer, "kv_cache", None)
+                if kv_cache_layer is None:
+                    continue
+
+                filename = self._generate_filename_debug(
+                    layer_name, request.token_ids, request.mm_hashes
+                )
+                kv_cache = safetensors.torch.load_file(
+                    filename, device=str(kv_cache_layer.device)
+                )["kv_cache"]
+                if isinstance(attn_metadata, dict):
+                    inject_kv_into_layer(
+                        kv_cache_layer,
+                        kv_cache,
+                        request.slot_mapping,
+                        attn_metadata[layer_name],
+                    )
+
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        """Blocking until the KV for a specific layer is loaded into vLLM's
+        paged buffer.
+
+        This interface will be useful for layer-by-layer pipelining.
+
+        Args:
+            layer_name: the name of that layer
+        """
+        return
+
+    def save_kv_layer(
+        self,
+        layer_name: str,
+        kv_layer: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        **kwargs: Any,
+    ) -> None:
+        """Start saving the KV cache of the layer from vLLM's paged buffer
+        to the connector.
+
+        Args:
+            layer_name (str): the name of the layer.
+            kv_layer (torch.Tensor): the paged KV buffer of the current
+                layer in vLLM.
+            attn_metadata (AttentionMetadata): the attention metadata.
+            **kwargs: additional arguments for the save operation.
+        """
+
+        def extract_kv_from_layer(
+            layer: torch.Tensor,
+            slot_mapping: torch.Tensor,
+        ) -> torch.Tensor:
+            """Extract the KV cache from the layer.
+
+            Assume the shape of the layer is (num_pages, page_size, xxx)
+            for MLA, and (num_pages, 2, page_size, xxx) otherwise.
+            """
+            if isinstance(attn_metadata, MLACommonMetadata):
+                num_pages, page_size = layer.shape[0], layer.shape[1]
+                return layer.reshape(num_pages * page_size, -1)[slot_mapping, ...]
+            block_idxs = slot_mapping // self._block_size
+            offsets = slot_mapping % self._block_size
+            return layer[block_idxs, :, offsets]
+
+        connector_metadata = self._get_connector_metadata()
+        assert isinstance(connector_metadata, ExampleConnectorMetadata)
+        for request in connector_metadata.requests:
+            if request.is_store:
+                filename = self._generate_filename_debug(
+                    layer_name, request.token_ids, request.mm_hashes
+                )
+                kv_cache = extract_kv_from_layer(kv_layer, request.slot_mapping)
+                tensors = {"kv_cache": kv_cache.detach().cpu()}
+                safetensors.torch.save_file(tensors, filename)
+
+    def wait_for_save(self):
+        return
+
+    def get_num_new_matched_tokens(
+        self,
+        request: "Request",
+        num_computed_tokens: int,
+    ) -> tuple[int | None, bool]:
+        """
+        Get number of new tokens that can be loaded from the
+        external KV cache beyond the num_computed_tokens.
+
+        Args:
+            request (Request): the request object.
+            num_computed_tokens (int): the number of locally
+                computed tokens for this request
+
+        Returns:
+            the number of tokens that can be loaded from the
+            external KV cache beyond what is already computed.
+        """
+        # NOTE: in this debug implementation, we assume that the prompt is
+        # cached_prompt + newly_generated_single_token
+        # Therefore, we use prompt_token_ids[:-1] to determine the folder name
+
+        # NOTE: in current v1 scheduler, the num_computed_tokens is aligned
+        # with the block granularity. And it expects the returned blocks and
+        # num_computed_tokens to also be aligned with the block granularity.
+        if not self._found_match_for_request(request):
+            return 0, False
+
+        logger.info("External Cache Hit!")
+
+        # Now, first num_tokens_to_check tokens are hit, we need to prepare
+        # the metadata for the worker connector to correctly load the KV
+        token_ids = request.prompt_token_ids or []
+        num_tokens_to_check = align_to_block_size(len(token_ids) - 1, self._block_size)
+
+        return num_tokens_to_check - num_computed_tokens, False
+
+    def update_state_after_alloc(
+        self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
+    ):
+        """
+        Update KVConnector state after block allocation.
+
+        If blocks were allocated, add to _requests_need_load,
+        such that we load the KVs in the next forward pass.
+        """
+        if num_external_tokens > 0:
+            self._requests_need_load[request.request_id] = request
+
+    def build_connector_meta(
+        self,
+        scheduler_output: SchedulerOutput,
+    ) -> KVConnectorMetadata:
+        """Build the connector metadata for this step.
+
+        This function should NOT modify any fields in the scheduler_output.
+        Also, calling this function will reset the state of the connector.
+
+        Args:
+            scheduler_output (SchedulerOutput): the scheduler output object.
+        """
+        meta = ExampleConnectorMetadata()
+
+        total_need_load = 0
+        for new_req in scheduler_output.scheduled_new_reqs:
+            token_ids = new_req.prompt_token_ids or []
+            mm_hashes = [f.identifier for f in new_req.mm_features]
+            if new_req.req_id in self._requests_need_load:
+                meta.add_request(
+                    token_ids=token_ids,
+                    block_ids=new_req.block_ids[0],
+                    block_size=self._block_size,
+                    is_store=False,
+                    mm_hashes=mm_hashes,
+                )
+                total_need_load += 1
+            else:
+                # NOTE: here, we set the store and load being exclusive,
+                # but a single request can have both store and load.
+                # NOTE(rob): for this debug implementation, we only cache
+                # the original prompt tokens.
+                if not self._found_match_for_prompt(token_ids, mm_hashes):
+                    meta.add_request(
+                        token_ids=token_ids,
+                        block_ids=new_req.block_ids[0],
+                        block_size=self._block_size,
+                        is_store=True,
+                        mm_hashes=mm_hashes,
+                    )
+
+        cached_reqs = scheduler_output.scheduled_cached_reqs
+        for i, req_id in enumerate(cached_reqs.req_ids):
+            resumed_from_preemption = req_id in cached_reqs.resumed_req_ids
+            if not resumed_from_preemption or req_id not in self._requests_need_load:
+                continue
+
+            num_computed_tokens = cached_reqs.num_computed_tokens[i]
+            num_new_tokens = scheduler_output.num_scheduled_tokens[req_id]
+            new_block_ids = cached_reqs.new_block_ids[i]
+
+            # NOTE(rob): cached_req_data does not have the full
+            # list of token ids (only new tokens). So we look it
+            # up in the actual request object.
+            request = self._requests_need_load[req_id]
+            total_tokens = num_computed_tokens + num_new_tokens
+            token_ids = request.all_token_ids[:total_tokens]
+
+            # NOTE(rob): For resumed req, new_block_ids is all
+            # of the block_ids for the request.
+            assert new_block_ids is not None
+            block_ids = new_block_ids[0]
+
+            meta.add_request(
+                token_ids=token_ids,
+                block_ids=block_ids,
+                block_size=self._block_size,
+                is_store=False,
+                mm_hashes=[f.identifier for f in request.mm_features],
+            )
+            total_need_load += 1
+
+        assert total_need_load == len(self._requests_need_load)
+        self._requests_need_load.clear()
+        return meta
+
+    # ==============================
+    # Helper functions
+    # ==============================
+
+    def _found_match_for_request(
+        self,
+        request: "Request",
+    ) -> bool:
+        """Check if the cache is hit for the request."""
+        return self._found_match_for_prompt(
+            list(request.prompt_token_ids or []),
+            [f.identifier for f in request.mm_features],
+        )
+
+    def _found_match_for_prompt(
+        self,
+        prompt_token_ids: list[int],
+        mm_hashes: list[str],
+    ) -> bool:
+        num_tokens_to_check = align_to_block_size(
+            len(prompt_token_ids) - 1, self._block_size
+        )
+        foldername = self._generate_foldername_debug(
+            torch.tensor(prompt_token_ids)[:num_tokens_to_check],
+            mm_hashes,
+            create_folder=False,
+        )
+        return os.path.exists(foldername)
+
+    def _generate_foldername_debug(
+        self,
+        token_ids: torch.Tensor,
+        mm_hashes: list[str],
+        create_folder=False,
+    ) -> str:
+        """Generate a folder name based on the hash of the bytes of the input
+        ids.
+        """
+        token_bytes = token_ids.numpy().tobytes()
+        # Add mm_hashes to the bytes being hashed to avoid path traversal and
+        # to create a canonical key.
+        if mm_hashes:
+            mm_str = "-".join(mm_hashes)
+            token_bytes += mm_str.encode("utf-8")
+        input_ids_hash = safe_hash(token_bytes, usedforsecurity=False).hexdigest()
+
+        foldername = os.path.join(self._storage_path, input_ids_hash)
+        if create_folder:
+            os.makedirs(foldername, exist_ok=True)
+        return foldername
+
+    def _generate_filename_debug(
+        self,
+        layer_name: str,
+        token_ids: torch.Tensor,
+        mm_hashes: list[str],
+    ) -> str:
+        """Generate a file name based on the layer name and the hash
+        of the bytes of the input ids.
+        """
+        foldername = self._generate_foldername_debug(
+            token_ids, mm_hashes=mm_hashes, create_folder=True
+        )
+        return os.path.join(foldername, f"{layer_name}.safetensors")
+
+
+def align_to_block_size(num_tokens: int, block_size) -> int:
+    """Align the number of tokens to the block size."""
+    return (num_tokens - 1) // block_size * block_size

vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py

@@ -0,0 +1,613 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import fcntl
+import os
+from concurrent.futures import Future, ThreadPoolExecutor
+from dataclasses import dataclass, field
+from functools import partial
+from importlib.metadata import version
+from typing import TYPE_CHECKING, Any
+
+import torch
+from packaging.version import Version
+from safetensors.torch import load_file, save_file
+
+from vllm.config import VllmConfig, get_layers_from_vllm_config
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorBase_V1,
+    KVConnectorMetadata,
+    KVConnectorRole,
+    SupportsHMA,
+)
+from vllm.distributed.parallel_state import get_tensor_model_parallel_rank
+from vllm.logger import init_logger
+from vllm.v1.attention.backend import AttentionMetadata
+from vllm.v1.core.sched.output import SchedulerOutput
+
+if TYPE_CHECKING:
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+    from vllm.v1.request import Request
+
+logger = init_logger(__name__)
+
+
+def extract_from_kv_cache(
+    kv_cache: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    num_tokens: int,
+) -> torch.Tensor:
+    """Extract data from KV cache."""
+    block_size = kv_cache.shape[1]
+    return kv_cache[slot_mapping // block_size, slot_mapping % block_size][:num_tokens]
+
+
+def load_hidden_states(path: str) -> dict[str, torch.Tensor]:
+    """Load hidden states written by ExampleHiddenStatesConnector.
+
+    Blocks (without polling) until the async write is complete by
+    acquiring a shared flock on the companion lock file.  The kernel
+    puts the caller to sleep until the writer releases its exclusive lock.
+
+    Args:
+        path: The file path returned in kv_transfer_params["hidden_states_path"].
+
+    Returns:
+        Dict with "hidden_states" and "token_ids" tensors.
+    """
+    lock_path = path + ".lock"
+    with open(lock_path) as lf:
+        fcntl.flock(lf, fcntl.LOCK_SH)  # sleeps until writer releases LOCK_EX
+        data = load_file(path, device="cpu")
+    return data
+
+
+def cleanup_hidden_states(path: str, keep_hidden_states: bool = False) -> None:
+    """Clean up hidden states file and lock file after loading.
+
+    If keep_hidden_states is True, only removes the lock file
+    and keeps the hidden states file.
+    """
+    lock_path = path + ".lock"
+    if os.path.exists(lock_path):
+        os.remove(lock_path)
+    if not keep_hidden_states and os.path.exists(path):
+        os.remove(path)
+
+
+@dataclass
+class PendingSave:
+    req_id: str
+    filename: str
+    token_ids: torch.Tensor
+    block_ids: list[int]
+
+
+@dataclass
+class ExampleHiddenStatesConnectorMetadata(KVConnectorMetadata):
+    pending_saves: list[PendingSave] = field(default_factory=list)
+    # req_id → filename for newly scheduled requests — the worker pre-creates
+    # lock files for these so the lock exists before the client receives the
+    # output path.
+    new_req_filenames: dict[str, str] = field(default_factory=dict)
+
+
+class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
+    """
+    Simple debug implementation of a HiddenStatesConnector.
+
+    Simply extracts the hidden states from the kv cache and stores them to disk.
+    Must be used in conjunction with the `extract_hidden_states` spec decoding method.
+    """
+
+    @property
+    def prefer_cross_layer_blocks(self) -> bool:
+        """
+        Indicates whether this connector prefers KV blocks that hold KV data for all
+        layers, which can speed up KV data transfers. Defaults to False.
+        """
+        # Must be False so that drafter kv cache isn't merged with verifier's
+        return False
+
+    @classmethod
+    def _find_cache_kv_group_id(cls, kv_cache_config: "KVCacheConfig | None") -> int:
+        """Index of the KV cache group holding the extracted hidden states.
+
+        Located by spec type so it resolves on both scheduler and worker side.
+        """
+        if kv_cache_config is None:
+            return 0
+
+        from vllm.v1.kv_cache_interface import HiddenStateCacheSpec
+
+        groups = kv_cache_config.kv_cache_groups
+        group_ids = [
+            gid
+            for gid, group in enumerate(groups)
+            if isinstance(group.kv_cache_spec, HiddenStateCacheSpec)
+        ]
+        if len(group_ids) == 1:
+            return group_ids[0]
+        if not group_ids and len(groups) == 1:
+            return 0
+        raise ValueError(
+            "Could not uniquely identify the extract-hidden-states KV cache "
+            f"group among {len(groups)} groups; the hidden-states layer must be "
+            "isolated in its own group (MLA verifiers are unsupported)."
+        )
+
+    @staticmethod
+    def _get_cache_block_size(
+        vllm_config: "VllmConfig",
+        kv_cache_config: "KVCacheConfig | None",
+        cache_kv_group_id: int,
+    ) -> int:
+        """Block size of the hidden-states group, read from its own spec.
+
+        cache_config.block_size is bumped to a common multiple for hybrid
+        verifiers; the page-aligned hidden-states group keeps a smaller one.
+        """
+        if kv_cache_config is None:
+            return vllm_config.cache_config.block_size
+        cache_group = kv_cache_config.kv_cache_groups[cache_kv_group_id]
+        return cache_group.kv_cache_spec.block_size
+
+    def __init__(
+        self,
+        vllm_config: "VllmConfig",
+        role: KVConnectorRole,
+        kv_cache_config: "KVCacheConfig",
+    ):
+        super().__init__(
+            vllm_config=vllm_config,
+            role=role,
+            kv_cache_config=kv_cache_config,
+        )
+        # Read the hidden-states group and its block size from the group spec;
+        # cache_config.block_size is bumped (wrong) for hybrid verifiers.
+        self._cache_kv_group_id = self._find_cache_kv_group_id(kv_cache_config)
+        self._block_size = self._get_cache_block_size(
+            vllm_config, kv_cache_config, self._cache_kv_group_id
+        )
+        self._storage_path = self._kv_transfer_config.get_from_extra_config(
+            "shared_storage_path", "/tmp"
+        )
+        self.cache_layers: list[str] = []  # set by self.register_kv_caches
+        logger.info(self._kv_transfer_config)
+        logger.info("Shared storage path is %s", self._storage_path)
+
+        if Version(version("safetensors")) < Version("0.8.0"):
+            logger.warning(
+                "safetensors < 0.8.0 holds the GIL during save_file, which "
+                "serializes the writer thread pool and hurts throughput. "
+                "Upgrade to safetensors >= 0.8.0 for better performance."
+            )
+
+        assert self._vllm_config.speculative_config is not None, (
+            "ExampleHiddenStatesConnector only works when using "
+            "'extract_hidden_states' speculative method"
+        )
+        spec_config = self._vllm_config.speculative_config.draft_model_config.hf_config
+        self.num_hidden_states = len(
+            getattr(spec_config, "eagle_aux_hidden_state_layer_ids", [])
+        )
+
+        # Scheduler-side state
+        self._pending_saves: dict[str, PendingSave] = {}
+        self._request_filenames: dict[str, str] = {}
+
+        # Worker-side state (set by register_kv_caches).
+        self._kv_cache: torch.Tensor | None = None
+
+        # Only TP rank 0 writes hidden states to disk; other TP ranks no-op.
+        # Set in register_kv_caches (after distributed init).
+        self._is_tp_rank_zero: bool = True
+
+        # Async write infrastructure (worker-side).
+        # Dedicated CUDA stream for DtoH copies so they don't block
+        # the default stream (model forward). Thread pool for disk writes.
+        self._copy_stream: torch.cuda.Stream | None = None  # lazy init
+        self._executor = ThreadPoolExecutor(
+            max_workers=self._kv_transfer_config.get_from_extra_config(
+                "num_writer_threads", 8
+            ),
+            thread_name_prefix="vllm-hs-save",
+        )
+        # Whether to use a filesystem lock when writing files to shared storage.
+        # This is necessary for online transfer clients to avoid incomplete reads,
+        # but can be disabled for offline tasks that run tasks in batches to completion
+        self.allow_custom_save_path = self._kv_transfer_config.get_from_extra_config(
+            "allow_custom_save_path", False
+        )
+        if self.allow_custom_save_path:
+            logger.warning(
+                "allow_custom_save_path is enabled. API clients can write "
+                "hidden states to arbitrary paths on the server filesystem. "
+                "Only enable this with trusted clients."
+            )
+        self.use_lock = self._kv_transfer_config.get_from_extra_config(
+            "use_synchronization_lock", True
+        )
+        # req_id → open fd on the .lock file with LOCK_EX held.
+        # Pre-created in wait_for_save when a request first arrives,
+        # consumed by _submit_async_write which passes the fd to the
+        # thread pool worker for release after writing.
+        self._lock_fds: dict[str, int] = {}
+        # req_id → in-flight disk-write Future for that req_id.
+        self._req_futures: dict[str, Future] = {}
+        # req_id → CUDA event marking completion of the DtoH copy. Once
+        # this event is complete the request is considered "done sending"
+        # by get_finished; clients block on the per-file flock to wait for
+        # the disk write itself.
+        self._req_copy_events: dict[str, torch.cuda.Event] = {}
+        # req_ids reported as finished-generating by the scheduler,
+        # accumulated across get_finished calls.
+        self._accumulated_finished_req_ids: set[str] = set()
+
+    def _get_copy_stream(self) -> torch.cuda.Stream:
+        """Lazily create the copy stream (CUDA must be initialized)."""
+        if self._copy_stream is None:
+            self._copy_stream = torch.cuda.Stream()
+        return self._copy_stream
+
+    # ==============================
+    # Worker-side methods
+    # ==============================
+    def start_load_kv(self, *args, **kwargs: Any) -> None:
+        pass  # Store-only connector — nothing to load
+
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        pass  # Store-only connector — nothing to load
+
+    def wait_for_save(self) -> None:
+        """Pre-create lock files for newly arrived requests.
+
+        This runs on the worker BEFORE the scheduler returns the output
+        path to the client, guaranteeing that the lock file exists (and
+        LOCK_EX is held) by the time the client tries to open it.
+        """
+        if not self._is_tp_rank_zero:
+            return
+        if not self.use_lock or not self.has_connector_metadata():
+            return
+        metadata = self._get_connector_metadata()
+        if not isinstance(metadata, ExampleHiddenStatesConnectorMetadata):
+            return
+        for req_id, filename in metadata.new_req_filenames.items():
+            if req_id in self._lock_fds:
+                continue
+            lock_path = filename + ".lock"
+            os.makedirs(os.path.dirname(lock_path), exist_ok=True)
+            lock_fd = os.open(lock_path, os.O_CREAT | os.O_WRONLY | os.O_TRUNC, 0o644)
+            fcntl.flock(lock_fd, fcntl.LOCK_EX)
+            self._lock_fds[req_id] = lock_fd
+
+    def _on_write_done(self, req_id: str, future: Future) -> None:
+        """Surface any exception from the disk-write thread and drop the
+        completed future from the in-flight tracking dict."""
+        self._req_futures.pop(req_id, None)
+        exc = future.exception()
+        if exc is not None:
+            logger.error("Hidden-states write failed for req_id=%s: %r", req_id, exc)
+
+    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
+        # Delay tp rank0 initialization until after distributed init
+        self._is_tp_rank_zero = get_tensor_model_parallel_rank() == 0
+
+        from vllm.model_executor.models.extract_hidden_states import (
+            CacheOnlyAttentionLayer,
+        )
+
+        # Filter layers to only include CacheOnlyAttentionLayers
+        layers = get_layers_from_vllm_config(
+            self._vllm_config, CacheOnlyAttentionLayer, list(kv_caches.keys())
+        )
+        self.cache_layers = list(layers.keys())
+        assert len(self.cache_layers) == 1, (
+            f"Expected 1 CacheOnlyAttentionLayer, got {len(self.cache_layers)}"
+        )
+        self._kv_cache = kv_caches[self.cache_layers[0]]
+
+        # Block size must match the indexed buffer, else reads hit the wrong
+        # slots. Raise (not assert) so the check survives `python -O`.
+        if self._block_size != self._kv_cache.shape[1]:
+            raise ValueError(
+                f"Hidden-states block-size mismatch: derived {self._block_size} "
+                f"but buffer block size is {self._kv_cache.shape[1]}; read slots "
+                "would be wrong (likely a hybrid block-size resolution bug)."
+            )
+
+    @staticmethod
+    def _write_tensors(
+        tensors: dict[str, torch.Tensor],
+        event: torch.cuda.Event,
+        filename: str,
+        lock_fd: int | None,
+    ) -> None:
+        """Thread worker: wait for async DtoH copy, write to disk, release lock.
+
+        ``lock_fd`` is an open file descriptor on the companion ``.lock``
+        file with ``LOCK_EX`` already held.  Closing it releases the lock,
+        which unblocks any client sleeping on ``LOCK_SH``.
+        """
+        try:
+            event.synchronize()
+            save_file(tensors, filename)
+        finally:
+            if lock_fd is not None:
+                os.close(lock_fd)  # releases LOCK_EX
+
+    def save_kv_layer(
+        self,
+        layer_name: str,
+        kv_layer: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        **kwargs: Any,
+    ) -> None:
+        # Hidden states are already cached by CacheOnlyAttentionLayer during
+        # forward. Extraction happens in get_finished once all tokens are done.
+        pass
+
+    def _submit_async_write(
+        self,
+        pending: PendingSave,
+    ) -> None:
+        """Extract hidden states from KV cache and submit async DtoH + disk write.
+
+        Called from get_finished for each request that has finished generating.
+        """
+        if not self._is_tp_rank_zero:
+            return
+        assert self._kv_cache is not None
+
+        # Compute slot mapping from block_ids
+        block_ids_t = torch.tensor(pending.block_ids, dtype=torch.long)
+        num_blocks = block_ids_t.shape[0]
+        block_offsets = torch.arange(0, self._block_size, dtype=torch.long)
+        slot_mapping = (
+            block_offsets.reshape((1, self._block_size))
+            + block_ids_t.reshape((num_blocks, 1)) * self._block_size
+        )
+        slot_mapping = slot_mapping.flatten()
+
+        num_tokens = pending.token_ids.shape[0]
+
+        copy_stream = self._get_copy_stream()
+
+        # Ensure the copy stream sees all prior writes on the default stream.
+        ready_event = torch.cuda.Event()
+        ready_event.record()
+        copy_stream.wait_event(ready_event)
+
+        with torch.cuda.stream(copy_stream):
+            # Move the CPU slot_mapping to GPU on the copy stream so the
+            # implicit H2D inside fancy indexing doesn't sync the default
+            # stream.
+            slot_mapping_gpu = slot_mapping.to(
+                device=self._kv_cache.device, non_blocking=True
+            )
+            hidden_states_gpu = extract_from_kv_cache(
+                self._kv_cache, slot_mapping_gpu, num_tokens
+            )
+            # Async DtoH copy into pinned host memory.
+            pinned_hs = torch.empty_like(
+                hidden_states_gpu, device="cpu", pin_memory=True
+            )
+            pinned_hs.copy_(hidden_states_gpu, non_blocking=True)
+
+        # Record completion of this copy on the copy stream.
+        copy_done = torch.cuda.Event()
+        copy_done.record(copy_stream)
+
+        # token_ids is already on CPU (created in request_finished).
+        assert not pending.token_ids.is_cuda, (
+            "Expected token_ids on CPU, got CUDA tensor"
+        )
+        tensors = {
+            "hidden_states": pinned_hs,
+            "token_ids": pending.token_ids.clone(),
+        }
+
+        # Submit to thread pool for disk write.
+        prior = self._req_futures.get(pending.req_id)
+        assert prior is None, "Found another KV transfer request with same req_id!"
+
+        os.makedirs(os.path.dirname(pending.filename), exist_ok=True)
+
+        # Use the pre-created lock fd from wait_for_save (already holds
+        # LOCK_EX). Falls back to creating one here if use_lock is True
+        # but no pre-created fd exists (shouldn't happen in normal flow).
+        lock_fd = self._lock_fds.pop(pending.req_id, None)
+        if lock_fd is None and self.use_lock:
+            lock_path = pending.filename + ".lock"
+            lock_fd = os.open(lock_path, os.O_CREAT | os.O_WRONLY | os.O_TRUNC, 0o644)
+            fcntl.flock(lock_fd, fcntl.LOCK_EX)
+
+        future = self._executor.submit(
+            self._write_tensors, tensors, copy_done, pending.filename, lock_fd
+        )
+        self._req_copy_events[pending.req_id] = copy_done
+        self._req_futures[pending.req_id] = future
+        future.add_done_callback(partial(self._on_write_done, pending.req_id))
+
+    # ==============================
+    # Scheduler-side methods
+    # ==============================
+
+    def get_num_new_matched_tokens(
+        self,
+        request: "Request",
+        num_computed_tokens: int,
+    ) -> tuple[int | None, bool]:
+        """
+        Get number of new tokens that can be loaded from the
+        external KV cache beyond the num_computed_tokens.
+
+        Args:
+            request (Request): the request object.
+            num_computed_tokens (int): the number of locally
+                computed tokens for this request
+
+        Returns:
+            the number of tokens that can be loaded from the
+            external KV cache beyond what is already computed.
+        """
+        # This connector is store-only, so we don't need to load any tokens
+        return 0, False
+
+    def update_state_after_alloc(
+        self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
+    ):
+        # Usually used to handle allocation of new blocks for requests that are loading
+        # tokens from connector's external kv cache. We never load from external cache
+        # so this is a no-op.
+        assert num_external_tokens == 0, "This connector is store-only"
+
+    def build_connector_meta(
+        self,
+        scheduler_output: SchedulerOutput,
+    ) -> KVConnectorMetadata:
+        """Build the connector metadata for this step.
+
+        This function should NOT modify any fields in the scheduler_output.
+        Also, calling this function will reset the state of the connector.
+
+        Args:
+            scheduler_output (SchedulerOutput): the scheduler output object.
+        """
+        meta = ExampleHiddenStatesConnectorMetadata()
+
+        # Transfer pending saves into metadata (scheduler → worker bridge)
+        meta.pending_saves = list(self._pending_saves.values())
+        self._pending_saves.clear()
+
+        # Resolve save paths for new requests and tell the worker so it can
+        # pre-create lock files before the client receives the output path.
+        for new_req in scheduler_output.scheduled_new_reqs:
+            default_path = os.path.join(
+                self._storage_path, f"{new_req.req_id}.safetensors"
+            )
+            kv_params = (
+                new_req.sampling_params.extra_args.get("kv_transfer_params")
+                if new_req.sampling_params and new_req.sampling_params.extra_args
+                else None
+            ) or {}
+            custom_path = kv_params.get("hidden_states_path")
+            if custom_path is not None and not self.allow_custom_save_path:
+                logger.warning(
+                    "Request %s provided hidden_states_path but "
+                    "allow_custom_save_path is disabled. Ignoring "
+                    "custom path and using default.",
+                    new_req.req_id,
+                )
+                custom_path = None
+            filename = custom_path or default_path
+            self._request_filenames[new_req.req_id] = filename
+            meta.new_req_filenames[new_req.req_id] = filename
+
+        return meta
+
+    def request_finished(
+        self,
+        request: "Request",
+        block_ids: list[int],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        """
+        Called exactly once when a request has finished, before its blocks are
+        freed.
+
+        Returns True to delay block freeing until get_finished extracts
+        the hidden states from the KV cache.
+        """
+        req_id = request.request_id
+        filename = self._request_filenames.pop(req_id)
+        kv_params = request.kv_transfer_params or {}
+        if kv_params.get("include_output_tokens", False):
+            # Exclude the final token — it was the model's output, never an
+            # input to a forward pass, so its hidden state is not in the cache.
+            token_ids = torch.tensor(list(request.all_token_ids)[:-1])
+        elif request.prompt_token_ids is not None:
+            token_ids = torch.tensor(request.prompt_token_ids)
+        else:
+            logger.warning(
+                "Request %s has no prompt_token_ids (prompt_embeds only). "
+                "Saved token_ids will be empty.",
+                req_id,
+            )
+            token_ids = torch.tensor([], dtype=torch.long)
+        self._pending_saves[req_id] = PendingSave(
+            req_id=req_id,
+            filename=filename,
+            token_ids=token_ids,
+            block_ids=list(block_ids),
+        )
+        return True, {"hidden_states_path": filename}
+
+    def get_finished(
+        self, finished_req_ids: set[str]
+    ) -> tuple[set[str] | None, set[str] | None]:
+        """Extract hidden states and poll DtoH-copy completion.
+
+        On the worker side, connector metadata carries pending saves from the
+        scheduler. For each one we extract from the KV cache and launch an
+        async DtoH copy + thread-pool disk write.
+
+        We then poll accumulated finished req_ids: a request is "done sending"
+        once its DtoH copy event is complete. The subsequent disk write may
+        still be in flight; clients block on the per-file flock to wait for it.
+        """
+        # Extract and submit async writes for newly finished requests
+        if self.has_connector_metadata():
+            connector_metadata = self._get_connector_metadata()
+            if isinstance(connector_metadata, ExampleHiddenStatesConnectorMetadata):
+                for pending in connector_metadata.pending_saves:
+                    self._submit_async_write(pending)
+
+        # Poll for completed DtoH copies
+        self._accumulated_finished_req_ids.update(finished_req_ids)
+
+        done_sending: set[str] = set()
+        for req_id in list(self._accumulated_finished_req_ids):
+            event = self._req_copy_events.get(req_id)
+            if event is None or event.query():
+                self._req_copy_events.pop(req_id, None)
+                done_sending.add(req_id)
+                self._accumulated_finished_req_ids.discard(req_id)
+                # Clean up any leftover lock fds (e.g. aborted requests
+                # that never went through _submit_async_write).
+                lock_fd = self._lock_fds.pop(req_id, None)
+                if lock_fd is not None:
+                    os.close(lock_fd)
+
+        return done_sending or None, None
+
+    def request_finished_all_groups(
+        self,
+        request: "Request",
+        block_ids: tuple[list[int], ...],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        return self.request_finished(request, block_ids[self._cache_kv_group_id])
+
+    @classmethod
+    def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None:
+        """
+        Get the required KV cache layout for this connector.
+        Args:
+            vllm_config (VllmConfig): the vllm config.
+
+        Returns:
+            str: the required KV cache layout. e.g. HND, or NHD.
+            None if the connector does not require a specific layout.
+        """
+
+        if cls is KVConnectorBase_V1:
+            raise TypeError(
+                "get_required_kvcache_layout should not be called "
+                "on the abstract base class"
+            )
+        # NHD means we have (num_tokens, num_heads)
+        # HND means we have (num_heads, num_tokens)
+        # For now, we only support NHD layout since this keeps the
+        # hidden states for each token together in memory.
+        # HND is primarily used when sharding heads across devices.
+        return "NHD"

vllm/distributed/kv_transfer/kv_connector/v1/flexkv_connector.py

@@ -0,0 +1,260 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorBase_V1,
+    KVConnectorMetadata,
+    KVConnectorRole,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
+from vllm.logger import init_logger
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.outputs import KVConnectorOutput
+
+if TYPE_CHECKING:
+    from vllm.distributed.kv_events import KVCacheEvent
+    from vllm.forward_context import ForwardContext
+    from vllm.v1.attention.backend import AttentionMetadata
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+    from vllm.v1.request import Request
+
+logger = init_logger(__name__)
+
+
+# FlexKV is a distributed KV Store and multi-level cache management system for
+# ultra-large-scale LLM inference.
+# GitHub: https://github.com/taco-project/FlexKV
+# Install: git clone git@github.com:taco-project/FlexKV.git \
+#          && cd FlexKV && bash build.sh
+class FlexKVConnectorV1(KVConnectorBase_V1):
+    """KV Connector that offloads KV cache to FlexKV.
+
+    FlexKV is a distributed KV Store and multi-level cache management system
+    designed for ultra-large-scale LLM inference. It supports offloading KV
+    cache to CPU memory, SSD, and remote storage.
+
+    Installation:
+        See https://github.com/taco-project/FlexKV for installation instructions.
+        Quick start::
+
+            git clone git@github.com:taco-project/FlexKV.git
+            cd FlexKV && bash build.sh
+
+    Configuration:
+        Pass ``kv_connector="FlexKVConnectorV1"`` via ``--kv-transfer-config``::
+
+            --kv-transfer-config \
+            '{"kv_connector":"FlexKVConnectorV1","kv_role":"kv_both"}'
+    """
+
+    def __init__(
+        self,
+        vllm_config: "VllmConfig",
+        role: KVConnectorRole,
+        kv_cache_config: "KVCacheConfig",
+    ):
+        super().__init__(
+            vllm_config=vllm_config, role=role, kv_cache_config=kv_cache_config
+        )
+        try:
+            from flexkv.integration.vllm.vllm_v1_adapter import FlexKVConnectorV1Impl
+        except ImportError as e:
+            raise ImportError(
+                "FlexKV is not installed. Please install it to use "
+                "FlexKVConnectorV1. See https://github.com/taco-project/FlexKV "
+                "for installation instructions."
+            ) from e
+
+        self._flexkv_connector = FlexKVConnectorV1Impl(vllm_config, role)
+
+    def shutdown(self):
+        self._flexkv_connector.shutdown()
+
+    # ==============================
+    # Worker-side methods
+    # ==============================
+    def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
+        """No-op for FlexKV (currently).
+
+        FlexKV manages all KV transfers on the **scheduler side** via
+        ``build_connector_meta`` (which calls ``launch_tasks``) and
+        ``update_connector_output`` (which polls ``query_finished_task``).
+        KV blocks are transferred directly between the FlexKV server and
+        vLLM's GPU memory without worker-side intervention during the
+        forward pass — similar to how NIXL operates.
+
+        These worker-side hooks are kept (rather than omitted) to satisfy
+        the ``KVConnectorBase_V1`` interface contract and to serve as
+        extension points for a future worker-side layer-pipelining path.
+
+        Args:
+            forward_context (ForwardContext): the forward context.
+            **kwargs (Any): additional arguments (unused).
+        """
+        self._flexkv_connector.start_load_kv(forward_context, **kwargs)
+
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        """No-op for FlexKV (currently).
+
+        FlexKV manages all KV transfers on the scheduler side.
+        This hook is retained for ``KVConnectorBase_V1`` API compatibility.
+
+        Args:
+            layer_name: the name of the layer (unused).
+        """
+        self._flexkv_connector.wait_for_layer_load(layer_name)
+
+    def save_kv_layer(
+        self,
+        layer_name: str,
+        kv_layer: torch.Tensor,
+        attn_metadata: "AttentionMetadata",
+        **kwargs,
+    ) -> None:
+        """No-op for FlexKV (currently).
+
+        FlexKV offloads KV cache asynchronously from the scheduler side
+        after a request finishes (see ``request_finished``).  It does not
+        intercept individual layer tensors during the forward pass.
+
+        This hook is retained to satisfy ``KVConnectorBase_V1`` and as an
+        extension point for future per-layer async offload support.
+
+        Args:
+            layer_name (str): the name of the layer (unused).
+            kv_layer (torch.Tensor): the paged KV buffer (unused).
+            attn_metadata (AttentionMetadata): the attention metadata (unused).
+            **kwargs (Any): additional arguments (unused).
+        """
+        self._flexkv_connector.save_kv_layer(
+            layer_name, kv_layer, attn_metadata, **kwargs
+        )
+
+    def wait_for_save(self):
+        """No-op for FlexKV (currently).
+
+        KV offload tasks are tracked asynchronously by the scheduler
+        connector via ``request_finished`` / ``query_finished_task``.
+        There is no pending worker-side save to wait for at
+        forward-context exit.
+
+        Retained to satisfy ``KVConnectorBase_V1`` and as an extension
+        point for future worker-side save-completion signalling.
+        """
+        self._flexkv_connector.wait_for_save()
+
+    def get_finished(
+        self, finished_req_ids: set[str]
+    ) -> tuple[set[str] | None, set[str] | None]:
+        """Notify worker-side connector of requests that have finished
+        generating tokens.
+
+        Returns:
+            Tuple of (sending/saving ids, recving/loading ids) for requests
+            that have finished asynchronous transfer. The finished saves/sends
+            req ids must belong to a set provided in a call to this method
+            (this call or a prior one).
+        """
+        return self._flexkv_connector.get_finished(finished_req_ids)
+
+    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
+        """Initialize with the KV caches. Useful for pre-registering the
+        KV caches in the KVConnector (e.g. for NIXL).
+
+        Args:
+            kv_caches: dictionary of layer names to kv cache tensors.
+        """
+        self._flexkv_connector.register_kv_caches(kv_caches)
+
+    # ==============================
+    # Scheduler-side methods
+    # ==============================
+    def get_num_new_matched_tokens(
+        self,
+        request: "Request",
+        num_computed_tokens: int,
+    ) -> tuple[int, bool]:
+        """Get the number of new tokens that can be loaded from the
+        external KV cache beyond ``num_computed_tokens``.
+
+        Args:
+            request (Request): the request object.
+            num_computed_tokens (int): the number of locally computed
+                tokens for this request.
+
+        Returns:
+            Tuple of (num_external_tokens, is_ready) where
+            num_external_tokens is the number of additional tokens that
+            can be loaded from the external KV cache.
+        """
+        return self._flexkv_connector.get_num_new_matched_tokens(
+            request, num_computed_tokens
+        )
+
+    def update_state_after_alloc(
+        self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
+    ):
+        """Update KVConnector state after block allocation."""
+        self._flexkv_connector.update_state_after_alloc(
+            request, blocks, num_external_tokens
+        )
+
+    def build_connector_meta(
+        self, scheduler_output: SchedulerOutput
+    ) -> KVConnectorMetadata:
+        """Build the connector metadata for this step.
+
+        This function should NOT modify fields in the scheduler_output.
+        Also, calling this function will reset the state of the connector.
+
+        Args:
+            scheduler_output (SchedulerOutput): the scheduler output object.
+        """
+        return self._flexkv_connector.build_connector_meta(scheduler_output)
+
+    def update_connector_output(self, connector_output: KVConnectorOutput):
+        """Update KVConnector state from worker-side connectors output.
+
+        Args:
+            connector_output (KVConnectorOutput): the worker-side
+                connectors output.
+        """
+        self._flexkv_connector.update_connector_output(connector_output)
+
+    def request_finished(
+        self,
+        request: "Request",
+        block_ids: list[int],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        """Called when a request has finished, before its blocks are freed.
+
+        Returns:
+            Tuple of (async_save, kv_transfer_params) where async_save is
+            True if the request is being saved/sent asynchronously and blocks
+            should not be freed until the request_id is returned from
+            :meth:`get_finished`. kv_transfer_params is an optional dict of
+            KVTransferParams to be included in the request outputs.
+        """
+        return self._flexkv_connector.request_finished(request, block_ids)
+
+    def take_events(self) -> Iterable["KVCacheEvent"]:
+        """Collect buffered KV cache events.
+
+        Returns:
+            New KV cache events since the last call.
+        """
+        return self._flexkv_connector.take_events()
+
+    def get_kv_connector_stats(self) -> KVConnectorStats | None:
+        """Get the KV connector stats collected during the last interval."""
+        return self._flexkv_connector.get_kv_connector_stats()
+
+    def get_block_ids_with_load_errors(self) -> set[int]:
+        """Get the block ids that have failed to load."""
+        return self._flexkv_connector.get_block_ids_with_load_errors()

vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_client.py

@@ -0,0 +1,298 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import logging
+import multiprocessing
+import os
+import threading
+from functools import wraps
+from pathlib import Path
+
+import torch
+import torch.utils.cpp_extension
+from torch.utils.cpp_extension import load
+
+root = Path(__file__).parent.resolve()
+cuda_include_path = os.path.join(torch.utils.cpp_extension.CUDA_HOME, "include")
+hf3fs_utils = load(
+    name="hf3fs_utils",
+    sources=[f"{root}/utils/hf3fs_utils.cpp"],
+    extra_include_paths=[cuda_include_path],
+)
+
+logger = logging.getLogger(__name__)
+
+HF3FS_AVAILABLE = True
+try:
+    from hf3fs_fuse.io import (
+        deregister_fd,
+        extract_mount_point,
+        make_ioring,
+        make_iovec,
+        register_fd,
+    )
+except ImportError:
+    HF3FS_AVAILABLE = False
+
+
+def rsynchronized():
+    def _decorator(func):
+        @wraps(func)
+        def wrapper(self, *args, **kwargs):
+            with self.rlock:
+                return func(self, *args, **kwargs)
+
+        return wrapper
+
+    return _decorator
+
+
+def wsynchronized():
+    def _decorator(func):
+        @wraps(func)
+        def wrapper(self, *args, **kwargs):
+            with self.wlock:
+                return func(self, *args, **kwargs)
+
+        return wrapper
+
+    return _decorator
+
+
+class Hf3fsClient:
+    def __init__(self, path: str, size: int, bytes_per_page: int, entries: int):
+        """Initialize the HF3FS client with hf3fs_fuse.
+
+        Args:
+            path: Path to the file used for storage
+            size: Total size of the storage file in bytes
+            bytes_per_page: Size of each page in bytes
+            entries: Maximum number of concurrent operations
+        """
+        if not HF3FS_AVAILABLE:
+            raise ImportError(
+                "hf3fs_fuse.io is not available. Please install the hf3fs_fuse package."
+            )
+
+        self.path = path
+        self.size = size
+        self.bytes_per_page = bytes_per_page
+        self.entries = entries
+
+        self._closed = False
+
+        self.file = None
+        self.shm_r = None
+        self.shm_w = None
+        self.ior_r = None
+        self.ior_w = None
+        self.iov_r = None
+        self.iov_w = None
+        try:
+            # Create the file if it doesn't exist and set its size
+            self.file = os.open(self.path, os.O_RDWR | os.O_CREAT)
+            os.ftruncate(self.file, size)
+            register_fd(self.file)
+
+            self.hf3fs_mount_point = extract_mount_point(path)
+            self.bs = self.bytes_per_page
+            self.shm_r = multiprocessing.shared_memory.SharedMemory(
+                size=self.bs * self.entries, create=True
+            )
+            self.shm_w = multiprocessing.shared_memory.SharedMemory(
+                size=self.bs * self.entries, create=True
+            )
+
+            self.shm_r_tensor = torch.frombuffer(self.shm_r.buf, dtype=torch.uint8)
+            self.shm_w_tensor = torch.frombuffer(self.shm_w.buf, dtype=torch.uint8)
+
+            numel = self.bs * self.entries
+            self.r_pinned = torch.empty(
+                numel,
+                dtype=torch.uint8,
+                device="cpu",
+                pin_memory=True,
+            )
+            self.w_pinned = torch.empty(
+                numel,
+                dtype=torch.uint8,
+                device="cpu",
+                pin_memory=True,
+            )
+
+            self.numa = -1
+            self.ior_r = make_ioring(
+                self.hf3fs_mount_point,
+                self.entries,
+                for_read=True,
+                timeout=1,
+                numa=self.numa,
+            )
+            self.ior_w = make_ioring(
+                self.hf3fs_mount_point,
+                self.entries,
+                for_read=False,
+                timeout=1,
+                numa=self.numa,
+            )
+            self.iov_r = make_iovec(self.shm_r, self.hf3fs_mount_point)
+            self.iov_w = make_iovec(self.shm_w, self.hf3fs_mount_point)
+            self.shm_r.unlink()
+            self.shm_w.unlink()
+
+            self.rlock = threading.RLock()
+            self.wlock = threading.RLock()
+
+            self.stream = torch.cuda.Stream()
+            self.stream_ptr_int = self.stream.cuda_stream
+
+        except Exception:
+            self._release_resources()
+            raise
+
+        logger.debug(
+            "Initialized HF3FS client with file: %s, size: %s bytes", path, size
+        )
+
+    def _release_resources(self) -> None:
+        """Release all acquired resources safely"""
+        # iov must be released before ioring and shm
+        for attr in ("iov_r", "iov_w", "ior_r", "ior_w"):
+            obj = getattr(self, attr, None)
+            if obj is not None:
+                del obj
+                setattr(self, attr, None)
+
+        for attr in ("shm_r", "shm_w"):
+            shm = getattr(self, attr, None)
+            if shm is not None:
+                try:
+                    shm.close()
+                except Exception as e:
+                    logger.warning("Failed to close %s: %s", attr, e)
+                setattr(self, attr, None)
+
+        if self.file is not None:
+            try:
+                deregister_fd(self.file)
+            except Exception as e:
+                logger.warning("deregister_fd failed: %s", e)
+            try:
+                os.close(self.file)
+            except OSError as e:
+                logger.warning("os.close failed: %s", e)
+            self.file = None
+
+    @rsynchronized()
+    def batch_read(self, offsets: list[int], tensors: list[torch.Tensor]) -> list[int]:
+        """Read data from the file at specified offsets into tensors.
+
+        Args:
+            offsets: List of byte offsets to read from
+            tensors: List of tensors to read data into
+
+        Returns:
+            List of operation results (0 for success, non-zero for error)
+        """
+        self.check(offsets, tensors)
+        assert self.ior_r is not None
+        assert self.iov_r is not None
+
+        # prepare
+        current = 0
+        for offset, tensor in zip(offsets, tensors):
+            size = tensor.numel() * tensor.itemsize
+            self.ior_r.prepare(
+                self.iov_r[current : current + size], True, self.file, offset
+            )
+            current += size
+
+        # submit
+        ionum = len(offsets)
+        resv = self.ior_r.submit().wait(min_results=ionum)
+
+        # results
+        with torch.cuda.stream(self.stream):
+            hf3fs_utils.read_shm(
+                self.shm_r_tensor, self.r_pinned, tensors, self.stream_ptr_int
+            )
+        results = [res.result for res in resv]
+
+        return results
+
+    @wsynchronized()
+    def batch_write(
+        self, offsets: list[int], tensors: list[torch.Tensor], event: torch.cuda.Event
+    ) -> list[int]:
+        """Write data from tensors to the file at specified offsets.
+
+        Args:
+            offsets: List of byte offsets to write to
+            tensors: List of tensors containing data to write
+
+        Returns:
+            List of operation results (0 for success, non-zero for error)
+        """
+
+        self.check(offsets, tensors)
+        assert self.ior_w is not None
+        assert self.iov_w is not None
+
+        # prepare
+        with torch.cuda.stream(self.stream):
+            self.stream.wait_event(event)
+            hf3fs_utils.write_shm(
+                tensors, self.shm_w_tensor, self.w_pinned, self.stream_ptr_int
+            )
+
+        current = 0
+        for offset, tensor in zip(offsets, tensors):
+            size = tensor.numel() * tensor.itemsize
+            self.ior_w.prepare(
+                self.iov_w[current : current + size], False, self.file, offset
+            )
+            current += size
+
+        # submit
+        ionum = len(offsets)
+        resv = self.ior_w.submit().wait(min_results=ionum)
+
+        # results
+        results = [res.result for res in resv]
+
+        return results
+
+    def check(self, offsets: list[int], tensors: list[torch.Tensor]) -> None:
+        sizes = [t.numel() * t.itemsize for t in tensors]
+        if any(
+            [
+                len(offsets) > self.entries,
+                len(offsets) != len(sizes),
+                any(
+                    offset < 0 or offset + size > self.size
+                    for offset, size in zip(offsets, sizes)
+                ),
+                any(size > self.bytes_per_page for size in sizes),
+            ]
+        ):
+            self.close()
+            raise ValueError("Hf3fsClient.check Failed")
+
+    def get_size(self) -> int:
+        """Get the total size of the storage file.
+
+        Returns:
+            Size of the file in bytes
+        """
+        return self.size
+
+    def close(self) -> None:
+        """Close the client and clean up resources."""
+        if self._closed:
+            return
+        self._closed = True
+        self._release_resources()
+
+    def flush(self) -> None:
+        """Flush any pending writes to disk."""
+        if not self._closed and self.file is not None:
+            os.fsync(self.file)

vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_metadata_server.py

@@ -0,0 +1,530 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+HF3FS Metadata Server with key-based organization.
+"""
+
+import argparse
+import logging
+import threading
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+try:
+    import orjson
+
+    HAS_ORJSON = True
+except ImportError:
+    import json as orjson  # type: ignore
+
+    HAS_ORJSON = False
+
+import requests
+from fastapi import FastAPI, HTTPException, Request, Response
+from fastapi.responses import ORJSONResponse
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class RankFileMetadata:
+    """Manages file page allocation for a single rank."""
+
+    rank_id: int
+    num_pages: int
+    free_pages: list[int]
+
+    def allocate_pages(self, num_pages: int) -> list[int]:
+        """Allocate specified number of free pages."""
+        if len(self.free_pages) < num_pages:
+            return []
+
+        allocated = self.free_pages[:num_pages]
+        self.free_pages = self.free_pages[num_pages:]
+        return allocated
+
+    def release_pages(self, page_indices: list[int]) -> None:
+        """Release pages back to free pool."""
+        for page_idx in page_indices:
+            if page_idx not in self.free_pages:
+                self.free_pages.append(page_idx)
+
+    def get_free_page_count(self) -> int:
+        """Get current number of free pages."""
+        return len(self.free_pages)
+
+
+@dataclass
+class KeyMetadata:
+    """Manages metadata for a single key across multiple ranks."""
+
+    key: str
+    rank_to_page: dict[int, int]  # rank -> allocated page index
+    tp_world_size: int
+
+    def add_rank_page(self, rank: int, page_index: int) -> None:
+        """Add page allocation for a specific rank."""
+        self.rank_to_page[rank] = page_index
+
+    def get_all_pages(self) -> list[tuple[int, int]]:
+        """Get all (rank, page) pairs for this key."""
+        return [(rank, page) for rank, page in self.rank_to_page.items()]
+
+    def get_rank_page(self, rank: int) -> int | None:
+        """Get page index for a specific rank."""
+        return self.rank_to_page.get(rank)
+
+    def is_complete(self) -> bool:
+        """Check if all ranks in the TP world have allocated pages."""
+        return len(self.rank_to_page) == self.tp_world_size
+
+
+class GlobalMetadataState:
+    """Manages global metadata state across all ranks and keys."""
+
+    def __init__(self):
+        self.global_lock = threading.RLock()
+        self.rank_metadata: dict[int, RankFileMetadata] = {}
+        self.key_metadata: dict[str, KeyMetadata] = {}
+
+    def clear(self) -> None:
+        """Clear all metadata state."""
+        with self.global_lock:
+            self.rank_metadata.clear()
+            self.key_metadata.clear()
+            logger.info("Cleared all metadata state")
+
+    def initialize_rank(self, rank: int, num_pages: int) -> None:
+        """Initialize a new rank with specified number of pages."""
+        with self.global_lock:
+            if rank not in self.rank_metadata:
+                self.rank_metadata[rank] = RankFileMetadata(
+                    rank, num_pages, list(range(num_pages))
+                )
+                logger.info("Initialized rank %s with %s pages", rank, num_pages)
+
+    def allocate_pages_for_keys(
+        self, rank: int, keys: list[tuple[str, str]]
+    ) -> dict[str, int]:
+        """Allocate one page for each key on the specified rank.
+
+        Args:
+            rank: Rank ID to allocate pages on
+            keys: List of keys to allocate pages for
+
+        Returns:
+            Dictionary mapping key -> allocated page index
+        """
+        with self.global_lock:
+            if rank not in self.rank_metadata:
+                raise ValueError(f"Rank {rank} not initialized")
+
+            # Batch allocate pages for all keys
+            num_pages_needed = len(keys)
+            allocated_pages = self.rank_metadata[rank].allocate_pages(num_pages_needed)
+
+            if len(allocated_pages) < num_pages_needed:
+                logger.warning(
+                    "Rank %s only allocated %s pages for %s keys",
+                    rank,
+                    len(allocated_pages),
+                    num_pages_needed,
+                )
+
+            allocation_results = {}
+            for i, (key, prefix_key) in enumerate(keys):
+                if key in self.key_metadata:
+                    key_meta = self.key_metadata[key]
+                    if key_meta.is_complete() and rank in key_meta.rank_to_page:
+                        # key is already fully written, reuse the existing page
+                        # and release the allocated pages back to the free pool.
+                        if i < len(allocated_pages):
+                            self.rank_metadata[rank].release_pages([allocated_pages[i]])
+                        allocation_results[key] = key_meta.rank_to_page[rank]
+                        continue
+
+                if i < len(allocated_pages):
+                    allocation_results[key] = allocated_pages[i]
+                else:
+                    allocation_results[key] = -1  # No pages available
+
+            return allocation_results
+
+    def confirm_write_for_keys(
+        self,
+        rank: int,
+        key_confirmations: list[tuple[str, int]],
+        pages_to_release: list[int] | None = None,
+    ) -> None:
+        """Confirm write operations for keys and update metadata.
+
+        Args:
+            rank: Rank ID that confirmed the writes
+            key_confirmations: List of (key, page_index) tuples
+            pages_to_release: List of page indices to release back to free pool
+        """
+        with self.global_lock:
+            # Confirm successful writes
+            for key, page_index in key_confirmations:
+                if key not in self.key_metadata:
+                    # Need to determine tp_world_size from rank_metadata
+                    tp_world_size = len(self.rank_metadata)
+                    self.key_metadata[key] = KeyMetadata(key, {}, tp_world_size)
+
+                # Add confirmed page to key metadata
+                self.key_metadata[key].add_rank_page(rank, page_index)
+
+            # Release specified pages back to free pool
+            if pages_to_release:
+                self.rank_metadata[rank].release_pages(pages_to_release)
+                logger.debug(
+                    "Released %s pages on rank %s: %s",
+                    len(pages_to_release),
+                    rank,
+                    pages_to_release,
+                )
+
+    def batch_key_exists(self, keys: list[str]) -> list[bool]:
+        """Check if keys exist in metadata and all ranks have confirmed writes.
+
+        Args:
+            keys: List of keys to check
+
+        Returns:
+            List of boolean values indicating key existence and completion
+        """
+        with self.global_lock:
+            results = []
+            for key in keys:
+                if key not in self.key_metadata:
+                    results.append(False)
+                else:
+                    # Check if all ranks in the TP world have confirmed writes
+                    key_meta = self.key_metadata[key]
+                    results.append(key_meta.is_complete())
+            return results
+
+    def get_key_locations(self, rank: int, keys: list[str]) -> list[int | None]:
+        """Get page indices for keys on a specific rank.
+
+        Args:
+            rank: Rank ID to query
+            keys: List of keys to look up
+
+        Returns:
+            List of page indices in the same order as input keys (None if key not found)
+        """
+        with self.global_lock:
+            if rank not in self.rank_metadata:
+                raise ValueError(f"Rank {rank} not initialized")
+
+            results = []
+            for key in keys:
+                if key in self.key_metadata:
+                    key_meta = self.key_metadata[key]
+                    if key_meta.is_complete():
+                        page_index = key_meta.get_rank_page(rank)
+                    else:
+                        page_index = None
+
+                    results.append(page_index)
+                else:
+                    results.append(None)
+
+            return results
+
+
+class Hf3fsMetadataServer:
+    """HF3FS Metadata Server with improved key-based organization."""
+
+    def __init__(self, persistence_path: str | None = None, save_interval: int = 60):
+        self.state = GlobalMetadataState()
+        if HAS_ORJSON:
+            self.app = FastAPI(default_response_class=ORJSONResponse)
+        else:
+            self.app = FastAPI()
+        self._setup_routes()
+
+    async def _read_json(self, request: Request) -> dict:
+        """Parse request JSON using orjson if available."""
+        body = await request.body()
+        return orjson.loads(body)
+
+    def _json_response(self, content: dict):
+        """Return ORJSONResponse when available to bypass jsonable_encoder."""
+        if HAS_ORJSON:
+            return ORJSONResponse(content)
+        else:
+            return content
+
+    def _setup_routes(self):
+        """Setup FastAPI routes for new API design."""
+        self.app.post("/rank/{rank}/initialize")(self.initialize_rank)
+        self.app.post("/keys/batch_allocate")(self.batch_allocate_pages_for_keys)
+        self.app.post("/keys/confirm_write")(self.confirm_write_for_keys)
+        self.app.post("/keys/batch_exists")(self.batch_key_exists)
+        self.app.post("/keys/get_locations")(self.get_key_locations)
+        self.app.post("/clear")(self.clear)
+
+    async def initialize_rank(self, rank: int, request: Request):
+        """Initialize a rank with specified number of pages."""
+        data = await self._read_json(request)
+        role = data.get("role", "worker")
+        num_pages = data.get("num_pages", 0)
+
+        if role == "scheduler":
+            return self._json_response(
+                {"message": "Scheduler role does not require initialization"}
+            )
+
+        if role == "worker" and num_pages > 0:
+            self.state.initialize_rank(rank, num_pages)
+            return self._json_response(
+                {"message": f"Rank {rank} initialized with {num_pages} pages"}
+            )
+        else:
+            raise HTTPException(
+                status_code=400, detail="Invalid initialization parameters"
+            )
+
+    async def batch_allocate_pages_for_keys(self, request: Request):
+        """Allocate one page for each key on a specific rank."""
+        data = await self._read_json(request)
+        rank = data.get("rank")
+        keys = data.get("keys", [])
+
+        # Validate input format
+        if rank is None or not isinstance(keys, list):
+            raise HTTPException(
+                status_code=400, detail="Invalid request format: need 'rank' and 'keys'"
+            )
+
+        try:
+            # Perform allocation
+            results = self.state.allocate_pages_for_keys(rank, keys)
+
+            # Convert results to response format
+            response = {"rank": rank, "results": list(results.items())}
+            return self._json_response(response)
+        except Exception as e:
+            raise HTTPException(
+                status_code=500, detail=f"Allocation failed: {str(e)}"
+            ) from e
+
+    async def confirm_write_for_keys(self, request: Request):
+        """Confirm write operations for keys."""
+        data = await self._read_json(request)
+        rank = data.get("rank")
+        confirmations = data.get("confirmations", [])
+        pages_to_release = data.get("pages_to_release", [])
+
+        # Validate input format
+        if rank is None or not isinstance(confirmations, list):
+            raise HTTPException(
+                status_code=400,
+                detail="Invalid request format: need 'rank' and 'confirmations'",
+            )
+
+        try:
+            self.state.confirm_write_for_keys(rank, confirmations, pages_to_release)
+
+            return Response(status_code=204)
+
+        except Exception as e:
+            logger.error("Confirm write for keys failed: %s", e)
+            raise HTTPException(
+                status_code=500, detail=f"Confirmation failed: {str(e)}"
+            ) from e
+
+    async def batch_key_exists(self, request: Request):
+        """Check if multiple keys exist in metadata."""
+        data = await self._read_json(request)
+        keys = data.get("keys", [])
+
+        if not isinstance(keys, list):
+            raise HTTPException(status_code=400, detail="Invalid keys format")
+
+        try:
+            exists_results = self.state.batch_key_exists(keys)
+            return self._json_response({"exists": exists_results})
+        except Exception as e:
+            raise HTTPException(
+                status_code=500, detail=f"Key existence check failed: {str(e)}"
+            ) from e
+
+    async def get_key_locations(self, request: Request):
+        """Get page indices for keys on a specific rank."""
+        data = await self._read_json(request)
+        rank = data.get("rank")
+        keys = data.get("keys", [])
+
+        # Validate input format
+        if rank is None or not isinstance(keys, list):
+            raise HTTPException(
+                status_code=400, detail="Invalid request format: need 'rank' and 'keys'"
+            )
+
+        try:
+            # Get key locations
+            locations = self.state.get_key_locations(rank, keys)
+            return self._json_response({"locations": locations})
+        except Exception as e:
+            raise HTTPException(
+                status_code=500, detail=f"Failed to get key locations: {str(e)}"
+            ) from e
+
+    async def clear(self, request: Request):
+        """Clear the metadata server."""
+        self.state.clear()
+        return Response(status_code=204)
+
+    def run(self, host: str = "0.0.0.0", port: int = 18000):
+        """Run the metadata server."""
+        import uvicorn
+
+        logger.info("Starting improved metadata server on http://%s:%s", host, port)
+        uvicorn.run(self.app, host=host, port=port)
+
+
+# --- Client implementation ---
+class Hf3fsMetadataInterface(ABC):
+    """Interface for HF3FS metadata operations."""
+
+    @abstractmethod
+    def initialize(self, rank: int, num_pages: int = 0, role: str = "worker") -> None:
+        """Initialize the metadata service with specified number of pages."""
+        pass
+
+    @abstractmethod
+    def allocate_pages_for_keys(
+        self, rank: int, keys: list[tuple[str, str]]
+    ) -> list[tuple[str, int]]:
+        """Allocate one page for each key on the specified rank."""
+        pass
+
+    @abstractmethod
+    def confirm_write_for_keys(
+        self,
+        rank: int,
+        key_confirmations: list[tuple[str, int]],
+        pages_to_release: list[int] | None = None,
+    ) -> None:
+        """Confirm write operations for keys and optionally release pages."""
+        pass
+
+    @abstractmethod
+    def batch_key_exists(self, keys: list[str]) -> list[bool]:
+        """Check if keys exist and are complete across all ranks."""
+        pass
+
+    @abstractmethod
+    def get_key_locations(self, rank: int, keys: list[str]) -> list[int]:
+        """Get page indices for keys on a specific rank."""
+        pass
+
+
+class Hf3fsGlobalMetadataClient(Hf3fsMetadataInterface):
+    """Global HTTP metadata client for HF3FS."""
+
+    def __init__(self, base_url: str = "http://localhost:18000", max_retries: int = 3):
+        self.base_url = base_url.rstrip("/")
+        self._session = requests.Session()
+
+        retry_strategy = Retry(
+            total=max_retries,
+            backoff_factor=0.3,
+            status_forcelist=[500, 502, 503, 504],
+            allowed_methods=["GET", "POST"],
+        )
+        adapter = HTTPAdapter(max_retries=retry_strategy)
+        self._session.mount("http://", adapter)
+
+    def _post(self, endpoint: str, json_data: dict) -> dict:
+        """Make POST request to metadata server."""
+        try:
+            url = f"{self.base_url}/{endpoint}"
+            headers = {"Content-Type": "application/json"}
+            if HAS_ORJSON:
+                payload = orjson.dumps(json_data)
+            else:
+                import json
+
+                payload = json.dumps(json_data).encode("utf-8")
+            response = self._session.post(url, data=payload, headers=headers)
+            response.raise_for_status()
+
+            if response.status_code == 204 or not response.content:
+                return {}
+            if HAS_ORJSON:
+                return orjson.loads(response.content)
+            else:
+                return response.json()
+        except requests.exceptions.RequestException as e:
+            logger.error("Failed to POST to %s after retries: %s", endpoint, e)
+            raise RuntimeError(f"Failed to connect to metadata server: {e}") from e
+
+    def initialize(self, rank: int, num_pages: int = 0, role: str = "worker") -> None:
+        """Initialize a rank with specified number of pages."""
+        self._post(f"rank/{rank}/initialize", {"num_pages": num_pages, "role": role})
+
+    def allocate_pages_for_keys(
+        self, rank: int, keys: list[tuple[str, str]]
+    ) -> list[tuple[str, int]]:
+        """Allocate pages for keys on the specified rank."""
+        response = self._post("keys/batch_allocate", {"rank": rank, "keys": keys})
+
+        # Convert response to expected format
+        return response.get("results", {})
+
+    def confirm_write_for_keys(
+        self,
+        rank: int,
+        key_confirmations: list[tuple[str, int]],
+        pages_to_release: list[int] | None = None,
+    ) -> None:
+        """Confirm write operations for keys and optionally release pages."""
+        payload = {
+            "rank": rank,
+            "confirmations": key_confirmations,
+            "pages_to_release": pages_to_release or [],
+        }
+
+        self._post("keys/confirm_write", payload)
+
+    def batch_key_exists(self, keys: list[str]) -> list[bool]:
+        """Check if keys exist and are complete across all ranks."""
+        response = self._post("keys/batch_exists", {"keys": keys})
+        return response.get("exists", [])
+
+    def get_key_locations(self, rank: int, keys: list[str]) -> list[int]:
+        """Get page indices for keys on a specific rank."""
+        response = self._post("keys/get_locations", {"rank": rank, "keys": keys})
+        return response.get("locations", [])
+
+
+def run_metadata_server(
+    host: str = "0.0.0.0",
+    port: int = 18000,
+):
+    """Run the improved HF3FS metadata server."""
+    server = Hf3fsMetadataServer()
+    server.run(host=host, port=port)
+
+
+# --- Main Execution ---
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Improved HF3FS Metadata Server")
+    parser.add_argument(
+        "--host", type=str, default="0.0.0.0", help="Host to bind the server to."
+    )
+    parser.add_argument(
+        "--port", type=int, default=18000, help="Port to run the server on."
+    )
+    args = parser.parse_args()
+
+    run_metadata_server(args.host, args.port)

vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/common.py

@@ -0,0 +1,139 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import threading
+from dataclasses import dataclass, field
+from typing import Optional
+
+from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
+from vllm.v1.request import Request
+
+
+class AtomicCounter:
+    """Thread-safe atomic counter for round-robin operations."""
+
+    def __init__(self, n: int):
+        assert n > 0, "Counter size must be positive"
+        self._n = n
+        self._value = 0
+        self._lock = threading.Lock()
+
+    def next(self) -> int:
+        """Get next value in round-robin fashion."""
+        with self._lock:
+            current = self._value
+            self._value = (current + 1) % self._n
+            return current
+
+
+@dataclass
+class LoadBlockInfo:
+    """Operation for loading blocks from external storage."""
+
+    num_computed_blocks: int
+    num_blocks_to_load: int
+    need_fetch_block_ids: list[int]
+
+
+@dataclass
+class SaveBlockInfo:
+    """Operation for saving blocks to external storage."""
+
+    skip_leading_blocks: int
+
+
+@dataclass
+class RequestSchedulingState:
+    """Unified request scheduling state management."""
+
+    request_id: str
+    request: Request | None = None
+
+    # Token and block tracking
+    token_ids: list[int] = field(default_factory=list)
+    allocated_block_ids: list[int] = field(default_factory=list)
+    num_saved_blocks: int = 0
+
+    # Load operation info
+    load_op: LoadBlockInfo | None = None
+
+    # Scheduling phase
+    phase: str = "NEW"  # NEW -> WAITING_TO_LOAD -> ACTIVE -> FINISHED
+
+    def needs_loading(self) -> bool:
+        """Check if request needs loading."""
+        return self.load_op is not None and self.load_op.num_blocks_to_load > 0
+
+    def is_ready_to_load(self) -> bool:
+        """Check if request is ready for loading."""
+        return self.phase == "WAITING_TO_LOAD" and self.needs_loading()
+
+    def update_tokens_and_blocks(self, new_token_ids: list[int], new_block_ids) -> None:
+        """Update with new tokens and blocks."""
+        if new_token_ids:
+            self.token_ids.extend(new_token_ids)
+
+        if new_block_ids is not None:
+            normalized_block_ids = self._normalize_block_ids(new_block_ids)
+            self.allocated_block_ids.extend(normalized_block_ids)
+
+    def _normalize_block_ids(self, block_ids) -> list[int]:
+        """Normalize block_ids to list format."""
+        if not block_ids:
+            return []
+        if isinstance(block_ids, tuple):
+            return block_ids[0] if block_ids else []
+        if isinstance(block_ids, list):
+            return block_ids
+        return []
+
+
+@dataclass
+class HF3FSRequestMetadata:
+    """Metadata for a single request in HF3FS connector."""
+
+    request_id: str
+    token_ids: list[int]
+    block_ids: list[int]
+    load_block_op: LoadBlockInfo | None = None
+    save_block_op: SaveBlockInfo | None = None
+
+    @staticmethod
+    def from_scheduling_state(
+        state: "RequestSchedulingState",
+        block_size: int,
+        load_op: LoadBlockInfo | None = None,
+        skip_leading_blocks: int | None = None,
+    ) -> Optional["HF3FSRequestMetadata"]:
+        """Create request metadata from scheduling state."""
+        token_count = len(state.token_ids)
+        total_blocks = token_count // block_size
+
+        skip_blocks = (
+            state.num_saved_blocks
+            if skip_leading_blocks is None
+            else skip_leading_blocks
+        )
+
+        new_blocks_to_save = total_blocks - state.num_saved_blocks
+        if new_blocks_to_save <= 0 and load_op is None:
+            return None
+
+        state.num_saved_blocks = total_blocks
+        return HF3FSRequestMetadata(
+            request_id=state.request_id,
+            token_ids=state.token_ids,
+            block_ids=state.allocated_block_ids,
+            load_block_op=load_op,
+            save_block_op=SaveBlockInfo(skip_leading_blocks=skip_blocks),
+        )
+
+
+class HF3FSConnectorMetadata(KVConnectorMetadata):
+    """Container for HF3FS connector metadata."""
+
+    def __init__(self):
+        self.requests: list[HF3FSRequestMetadata] = []
+
+    def add_request(self, request_metadata: HF3FSRequestMetadata) -> None:
+        """Add request to metadata."""
+        self.requests.append(request_metadata)

vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/gather_scatter_helper.py

@@ -0,0 +1,288 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+
+from vllm.logger import init_logger
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def kv_cache_scatter_kernel(
+    kv_cache_ptrs_ptr,
+    source_ptr,
+    token_indices_ptr,
+    num_tokens_in_block,
+    hidden_size,
+    total_token_in_kvcache,
+    num_layers,
+    is_mla,
+    BLOCK_SIZE: tl.constexpr,
+):
+    layer_idx = tl.program_id(0)
+    token_pos = tl.program_id(1)
+
+    if layer_idx >= num_layers or token_pos >= num_tokens_in_block:
+        return
+
+    token_idx = tl.load(token_indices_ptr + token_pos)
+    kv_cache_ptr = tl.cast(tl.load(kv_cache_ptrs_ptr + layer_idx), source_ptr.dtype)
+
+    if token_idx >= total_token_in_kvcache:
+        return
+
+    if is_mla:
+        # MLA format: source [num_layers, num_tokens_in_block, hidden_size]
+        # MLA format: target [total_token_in_kvcache, hidden_size] (per layer)
+        source_offset = (layer_idx * num_tokens_in_block + token_pos) * hidden_size
+        target_offset = token_idx * hidden_size
+
+        for i in range(0, hidden_size, BLOCK_SIZE):
+            offset = i + tl.arange(0, BLOCK_SIZE)
+            mask = offset < hidden_size
+            val = tl.load(source_ptr + source_offset + offset, mask=mask)
+            tl.store(kv_cache_ptr + target_offset + offset, val, mask=mask)
+    else:
+        # MHA format: source [num_layers, 2, num_tokens_in_block, hidden_size]
+        # MHA format: target [2, total_token_in_kvcache, hidden_size]
+        source_offset_k = (
+            layer_idx * num_tokens_in_block * 2 + token_pos
+        ) * hidden_size
+        source_offset_v = (
+            layer_idx * num_tokens_in_block * 2 + num_tokens_in_block + token_pos
+        ) * hidden_size
+
+        target_offset_k = token_idx * hidden_size
+        target_offset_v = (total_token_in_kvcache + token_idx) * hidden_size
+
+        for i in range(0, hidden_size, BLOCK_SIZE):
+            offset = i + tl.arange(0, BLOCK_SIZE)
+            mask = offset < hidden_size
+
+            val_k = tl.load(source_ptr + source_offset_k + offset, mask=mask)
+            val_v = tl.load(source_ptr + source_offset_v + offset, mask=mask)
+
+            tl.store(kv_cache_ptr + target_offset_k + offset, val_k, mask=mask)
+            tl.store(kv_cache_ptr + target_offset_v + offset, val_v, mask=mask)
+
+
+@triton.jit
+def kv_cache_gather_kernel(
+    kv_cache_ptrs_ptr,
+    dst_ptr,
+    token_indices_ptr,
+    num_tokens_in_block,
+    hidden_size,
+    total_token_in_kvcache,
+    num_layers,
+    is_mla,
+    BLOCK_SIZE: tl.constexpr,
+):
+    layer_idx = tl.program_id(0)
+    token_pos = tl.program_id(1)
+
+    if layer_idx >= num_layers or token_pos >= num_tokens_in_block:
+        return
+
+    token_idx = tl.load(token_indices_ptr + token_pos)
+    kv_cache_ptr = tl.cast(tl.load(kv_cache_ptrs_ptr + layer_idx), dst_ptr.dtype)
+
+    if token_idx >= total_token_in_kvcache:
+        return
+
+    if is_mla:
+        # MLA format: source [total_token_in_kvcache, hidden_size] (per layer)
+        # MLA format: dst [num_layers, num_tokens_in_block, hidden_size]
+        kvcache_offset = token_idx * hidden_size
+        dst_offset = (layer_idx * num_tokens_in_block + token_pos) * hidden_size
+
+        for i in range(0, hidden_size, BLOCK_SIZE):
+            offset = i + tl.arange(0, BLOCK_SIZE)
+            mask = offset < hidden_size
+            val = tl.load(kv_cache_ptr + kvcache_offset + offset, mask=mask)
+            tl.store(dst_ptr + dst_offset + offset, val, mask=mask)
+    else:
+        # MHA format: source [2, total_token_in_kvcache, hidden_size]
+        # MHA format: dst [num_layers, 2, num_tokens_in_block, hidden_size]
+        dst_offset_k = (layer_idx * num_tokens_in_block * 2 + token_pos) * hidden_size
+        dst_offset_v = (
+            layer_idx * num_tokens_in_block * 2 + num_tokens_in_block + token_pos
+        ) * hidden_size
+
+        kvcache_offset_k = token_idx * hidden_size
+        kvcache_offset_v = (total_token_in_kvcache + token_idx) * hidden_size
+
+        for i in range(0, hidden_size, BLOCK_SIZE):
+            offset = i + tl.arange(0, BLOCK_SIZE)
+            mask = offset < hidden_size
+
+            val_k = tl.load(kv_cache_ptr + kvcache_offset_k + offset, mask=mask)
+            val_v = tl.load(kv_cache_ptr + kvcache_offset_v + offset, mask=mask)
+
+            tl.store(dst_ptr + dst_offset_k + offset, val_k, mask=mask)
+            tl.store(dst_ptr + dst_offset_v + offset, val_v, mask=mask)
+
+
+def scatter_kv_caches(
+    kv_caches_ptrs: torch.Tensor,
+    total_token_in_kvcache: int,
+    src_tensor: torch.Tensor,
+    token_indices: list[int],
+    is_mla: bool = False,
+) -> None:
+    """Scatter KV cache data from source tensor to KV cache storage.
+
+    Args:
+        kv_caches_ptrs: Tensor of KV cache pointers (one per layer)
+        total_token_in_kvcache: Total number of tokens in KV cache
+        src_tensor: Source tensor containing data to scatter
+            - MHA format: [num_layers, 2, num_tokens_in_block, hidden_size]
+            - MLA format: [num_layers, num_tokens_in_block, hidden_size]
+        token_indices: List of token positions to update
+        is_mla: Whether using MLA model format
+    """
+    num_layers = len(kv_caches_ptrs)
+    num_tokens_in_block = len(token_indices)
+
+    if is_mla:
+        # MLA: src_tensor is [num_layers, num_tokens_in_block, hidden_size]
+        assert len(src_tensor.shape) == 3, (
+            f"MLA src_tensor should be 3D, got {src_tensor.shape}"
+        )
+        hidden_size = src_tensor.shape[2]
+    else:
+        # MHA: src_tensor is [num_layers, 2, num_tokens_in_block, hidden_size]
+        assert len(src_tensor.shape) == 4, (
+            f"MHA src_tensor should be 4D, got {src_tensor.shape}"
+        )
+        hidden_size = src_tensor.shape[3]
+
+    device = src_tensor.device
+    token_indices_tensor = torch.tensor(
+        token_indices, dtype=torch.int32, device="cpu"
+    ).to(device, non_blocking=True)
+
+    grid = (num_layers, num_tokens_in_block)
+    BLOCK_SIZE = 128
+
+    kv_cache_scatter_kernel[grid](
+        kv_caches_ptrs,
+        src_tensor,
+        token_indices_tensor,
+        num_tokens_in_block,
+        hidden_size,
+        total_token_in_kvcache,
+        num_layers,
+        is_mla,
+        BLOCK_SIZE=BLOCK_SIZE,
+    )
+
+
+def gather_kv_caches(
+    kv_caches_ptrs: torch.Tensor,
+    total_token_in_kvcache: int,
+    dst_tensor: torch.Tensor,
+    token_indices: list[int],
+    is_mla: bool = False,
+) -> None:
+    """Gather KV cache data from KV cache storage to destination tensor.
+
+    Args:
+        kv_caches_ptrs: Tensor of KV cache pointers (one per layer)
+        total_token_in_kvcache: Total number of tokens in KV cache
+        dst_tensor: Destination tensor to store gathered data
+            - MHA format: [num_layers, 2, num_tokens_in_block, hidden_size]
+            - MLA format: [num_layers, num_tokens_in_block, hidden_size]
+        token_indices: List of token positions to gather
+        is_mla: Whether using MLA model format
+    """
+    num_layers = kv_caches_ptrs.shape[0]
+    num_tokens_in_block = len(token_indices)
+
+    if is_mla:
+        # MLA: dst_tensor is [num_layers, num_tokens_in_block, hidden_size]
+        assert len(dst_tensor.shape) == 3, (
+            f"MLA dst_tensor should be 3D, got {dst_tensor.shape}"
+        )
+        assert dst_tensor.shape[0] == num_layers, (
+            f"Layer count mismatch: {dst_tensor.shape[0]} vs {num_layers}"
+        )
+        assert dst_tensor.shape[1] == num_tokens_in_block, (
+            f"Token count mismatch: {dst_tensor.shape[1]} vs {num_tokens_in_block}"
+        )
+        hidden_size = dst_tensor.shape[2]
+    else:
+        # MHA: dst_tensor is [num_layers, 2, num_tokens_in_block, hidden_size]
+        assert len(dst_tensor.shape) == 4, (
+            f"MHA dst_tensor should be 4D, got {dst_tensor.shape}"
+        )
+        assert dst_tensor.shape[0] == num_layers, (
+            f"Layer count mismatch: {dst_tensor.shape[0]} vs {num_layers}"
+        )
+        assert dst_tensor.shape[1] == 2, (
+            f"MHA should have 2 (K,V) components, got {dst_tensor.shape[1]}"
+        )
+        assert dst_tensor.shape[2] == num_tokens_in_block, (
+            f"Token count mismatch: {dst_tensor.shape[2]} vs {num_tokens_in_block}"
+        )
+        hidden_size = dst_tensor.shape[3]
+
+    device = dst_tensor.device
+    token_indices_tensor = torch.tensor(
+        token_indices, dtype=torch.int32, device="cpu"
+    ).to(device, non_blocking=True)
+
+    grid = (num_layers, num_tokens_in_block)
+    BLOCK_SIZE = 128
+
+    kv_cache_gather_kernel[grid](
+        kv_caches_ptrs,
+        dst_tensor,
+        token_indices_tensor,
+        num_tokens_in_block,
+        hidden_size,
+        total_token_in_kvcache,
+        num_layers,
+        is_mla,
+        BLOCK_SIZE=BLOCK_SIZE,
+    )
+
+
+class CopyBufferAllocator:
+    """Memory pool for tensor buffers to avoid frequent allocation/deallocation."""
+
+    def __init__(
+        self, device: torch.device, dtype: torch.dtype, shape: list, max_count: int
+    ):
+        self._shape = shape
+        self._max_count = max_count
+        self._device = device
+        self._free_buffers = [
+            torch.empty(shape, dtype=dtype, device=device) for _ in range(max_count)
+        ]
+        self._inuse_count = 0
+
+    def alloc_buffer(self, count: int) -> list[torch.Tensor] | None:
+        """Allocate buffers from the pool."""
+        if count == 0:
+            return []
+
+        if self._inuse_count + count <= self._max_count:
+            self._inuse_count += count
+            result = self._free_buffers[-count:]
+            del self._free_buffers[-count:]
+            return result
+        return None
+
+    def free_buffer(self, buffers: list[torch.Tensor]) -> None:
+        """Return buffers to the pool."""
+        if not buffers:
+            return
+
+        if self._inuse_count >= len(buffers):
+            self._inuse_count -= len(buffers)
+            self._free_buffers.extend(buffers)
+        else:
+            raise RuntimeError("Attempted to free more buffers than allocated")
+
+
+logger = init_logger(__name__)

vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/hf3fs_mock_client.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import logging
+import os
+
+import torch
+
+logger = logging.getLogger(__name__)
+HF3FS_AVAILABLE = True
+
+
+class Hf3fsClient:
+    """Mock HF3FS client using file backend for debugging and testing."""
+
+    def __init__(self, path: str, size: int, bytes_per_page: int, entries: int):
+        self._size = size
+        self._bytes_per_page = bytes_per_page
+        self._entries = entries
+        self._file_path = path
+
+        self._ensure_file_exists()
+        logger.debug("Initialized mock HF3FS client: %s (%d bytes)", path, size)
+
+    def _ensure_file_exists(self) -> None:
+        """Create file if it doesn't exist."""
+        if not os.path.exists(self._file_path):
+            with open(self._file_path, "w+b") as f:
+                f.truncate(self._size)
+
+    def batch_read(self, offsets: list[int], tensors: list[torch.Tensor]) -> list[int]:
+        """Read data from file at specified offsets into tensors."""
+        results = []
+
+        try:
+            with open(self._file_path, "rb") as f:
+                for offset, tensor in zip(offsets, tensors):
+                    num_bytes = tensor.numel() * tensor.element_size()
+
+                    if offset < 0 or offset + num_bytes > self._size:
+                        results.append(-1)
+                        continue
+
+                    f.seek(offset)
+                    buffer_data = f.read(num_bytes)
+
+                    if len(buffer_data) == num_bytes == self._bytes_per_page:
+                        tensor_data = self._convert_buffer_to_tensor(
+                            buffer_data, tensor.dtype
+                        )
+                        tensor.copy_(
+                            tensor_data.reshape(tensor.shape).to(tensor.device)
+                        )
+                        results.append(self._bytes_per_page)
+                    else:
+                        logger.error(
+                            "Read size mismatch: got %d, expected %d",
+                            len(buffer_data),
+                            num_bytes,
+                        )
+                        results.append(-1)
+        except Exception as e:
+            logger.error("Batch read error: %s", e)
+            results.extend([-1] * (len(offsets) - len(results)))
+
+        return results
+
+    def _convert_buffer_to_tensor(
+        self, buffer_data: bytes, dtype: torch.dtype
+    ) -> torch.Tensor:
+        """Convert buffer data to tensor with proper dtype handling."""
+        if dtype == torch.bfloat16:
+            tensor_data = torch.frombuffer(buffer_data, dtype=torch.uint16)
+            return tensor_data.view(dtype=torch.bfloat16)
+        else:
+            return torch.frombuffer(buffer_data, dtype=dtype)
+
+    def batch_write(
+        self, offsets: list[int], tensors: list[torch.Tensor], event: torch.cuda.Event
+    ) -> list[int]:
+        """Write data from tensors to file at specified offsets."""
+        results = []
+
+        try:
+            torch.cuda.current_stream().wait_event(event)
+
+            # Convert tensors to bytes
+            data_bytes_list = [self._tensor_to_bytes(tensor) for tensor in tensors]
+
+            # Write to file
+            with open(self._file_path, "r+b") as f:
+                for offset, data_bytes in zip(offsets, data_bytes_list):
+                    if offset < 0 or offset + len(data_bytes) > self._size:
+                        results.append(-1)
+                        continue
+
+                    f.seek(offset)
+                    bytes_written = f.write(data_bytes)
+
+                    if bytes_written == len(data_bytes) == self._bytes_per_page:
+                        results.append(self._bytes_per_page)
+                    else:
+                        logger.error(
+                            "Write size mismatch: wrote %d, expected %d",
+                            bytes_written,
+                            self._bytes_per_page,
+                        )
+                        results.append(-1)
+
+        except Exception as e:
+            logger.error("Batch write error: %s", e)
+            results.extend([-1] * (len(offsets) - len(results)))
+
+        return results
+
+    def _tensor_to_bytes(self, tensor: torch.Tensor) -> bytes:
+        """Convert tensor to bytes with proper dtype handling."""
+        cpu_tensor = tensor.cpu()
+        if cpu_tensor.dtype == torch.bfloat16:
+            return cpu_tensor.view(dtype=torch.uint16).numpy().tobytes()
+        else:
+            return cpu_tensor.numpy().tobytes()
+
+    def get_size(self) -> int:
+        """Get the total size of the storage file."""
+        return self._size
+
+    def close(self) -> None:
+        """Close the client (no-op for file backend)."""
+        pass
+
+    def flush(self) -> None:
+        """Flush any pending writes (no-op for file backend)."""
+        pass

vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/hf3fs_utils.cpp

@@ -0,0 +1,57 @@
+#include <cuda_runtime.h>
+#include <torch/extension.h>
+#include <cstring>
+#include <vector>
+
+void read_shm(const torch::Tensor& shm, const torch::Tensor& pin,
+              std::vector<torch::Tensor> dst, uint64_t stream_ptr) {
+  py::gil_scoped_release release;
+
+  cudaStream_t stream = reinterpret_cast<cudaStream_t>(stream_ptr);
+
+  // Copy from shared memory to pinned memory
+  char* shm_ptr = static_cast<char*>(shm.data_ptr());
+  char* src_ptr = static_cast<char*>(pin.data_ptr());
+  std::memcpy(src_ptr, shm_ptr, shm.numel() * shm.element_size());
+
+  // Copy from pinned memory to GPU tensors
+  size_t current = 0;
+  for (size_t i = 0; i < dst.size(); ++i) {
+    auto& t = dst[i];
+    size_t t_bytes = t.numel() * t.element_size();
+    char* dst_ptr = static_cast<char*>(t.data_ptr());
+    cudaMemcpyAsync(dst_ptr, src_ptr + current, t_bytes, cudaMemcpyHostToDevice,
+                    stream);
+    current += t_bytes;
+  }
+  cudaStreamSynchronize(stream);
+}
+
+void write_shm(const std::vector<torch::Tensor> src, torch::Tensor& shm,
+               const torch::Tensor& pin, uint64_t stream_ptr) {
+  py::gil_scoped_release release;
+
+  cudaStream_t stream = reinterpret_cast<cudaStream_t>(stream_ptr);
+
+  // Copy from GPU tensors to pinned memory
+  char* dst_ptr = static_cast<char*>(pin.data_ptr());
+  size_t current = 0;
+  for (size_t i = 0; i < src.size(); ++i) {
+    auto& t = src[i];
+    size_t t_bytes = t.numel() * t.element_size();
+    char* src_ptr = static_cast<char*>(t.data_ptr());
+    cudaMemcpyAsync(dst_ptr + current, src_ptr, t_bytes, cudaMemcpyDeviceToHost,
+                    stream);
+    current += t_bytes;
+  }
+  cudaStreamSynchronize(stream);
+
+  // Copy from pinned memory to shared memory
+  char* shm_ptr = static_cast<char*>(shm.data_ptr());
+  std::memcpy(shm_ptr, dst_ptr, shm.numel() * shm.element_size());
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("read_shm", &read_shm, "Read tensors from shared memory");
+  m.def("write_shm", &write_shm, "Write tensors to shared memory");
+}

vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py

@@ -0,0 +1,354 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.distributed.kv_events import (
+    BlockStored,
+    KVCacheEvent,
+    KVConnectorKVEvents,
+    KVEventAggregator,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorBase_V1,
+    KVConnectorMetadata,
+    KVConnectorRole,
+)
+from vllm.logger import init_logger
+from vllm.v1.attention.backend import AttentionMetadata
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.outputs import KVConnectorOutput
+
+if TYPE_CHECKING:
+    from vllm.forward_context import ForwardContext
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+    from vllm.v1.request import Request
+
+logger = init_logger(__name__)
+
+
+class LMCacheKVEvents(KVConnectorKVEvents):
+    """
+    Concrete implementation of KVConnectorKVEvents using KVEventAggregator.
+    """
+
+    def __init__(self, num_workers: int) -> None:
+        self._aggregator = KVEventAggregator(num_workers)
+
+    def add_events(self, events: list[KVCacheEvent]) -> None:
+        self._aggregator.add_events(events)
+
+    def aggregate(self) -> "LMCacheKVEvents":
+        """
+        Aggregate KV events and retain only common events.
+        """
+        common_events = self._aggregator.get_common_events()
+        self._aggregator.clear_events()
+        self._aggregator.add_events(common_events)
+        self._aggregator.reset_workers()
+        return self
+
+    def increment_workers(self, count: int = 1) -> None:
+        self._aggregator.increment_workers(count)
+
+    def get_all_events(self) -> list[KVCacheEvent]:
+        return self._aggregator.get_all_events()
+
+    def get_number_of_workers(self) -> int:
+        return self._aggregator.get_number_of_workers()
+
+    def clear_events(self) -> None:
+        self._aggregator.clear_events()
+        self._aggregator.reset_workers()
+
+    def __repr__(self) -> str:
+        return f"<LMCacheKVEvents events={self.get_all_events()}>"
+
+
+class LMCacheConnectorV1(KVConnectorBase_V1):
+    @classmethod
+    def requires_piecewise_for_cudagraph(cls, extra_config: dict[str, Any]) -> bool:
+        """
+        LMCache requires PIECEWISE CUDA graph mode when layerwise
+        operations are enabled. The wait_for_layer_load and save_kv_layer
+        methods perform actual async synchronization that cannot be
+        captured in CUDA graphs.
+        """
+        return extra_config.get("use_layerwise", False)
+
+    def __init__(
+        self,
+        vllm_config: "VllmConfig",
+        role: KVConnectorRole,
+        kv_cache_config: "KVCacheConfig",
+    ):
+        super().__init__(
+            vllm_config=vllm_config, role=role, kv_cache_config=kv_cache_config
+        )
+        assert vllm_config.kv_transfer_config is not None
+        use_native = vllm_config.kv_transfer_config.get_from_extra_config(
+            "use_native", False
+        )
+        if use_native:
+            logger.info("Initializing native LMCache connector")
+            # lazy import
+            from vllm.distributed.kv_transfer.kv_connector.v1 import lmcache_integration
+
+            _adapter = lmcache_integration.vllm_v1_adapter
+
+            cls = _adapter.LMCacheConnectorV1Impl
+        else:
+            logger.info("Initializing latest dev LMCache connector")
+            # lazy import
+            from lmcache.integration.vllm.vllm_v1_adapter import (
+                LMCacheConnectorV1Impl as LMCacheConnectorLatestImpl,
+            )
+
+            cls = LMCacheConnectorLatestImpl
+
+        self._lmcache_engine = cls(vllm_config, role, self)
+
+        self._kv_cache_events: LMCacheKVEvents | None = None
+
+    # ==============================
+    # Worker-side methods
+    # ==============================
+    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
+        """
+        Initialize with the KV caches. Useful for pre-registering the
+        KV Caches in the KVConnector (e.g. for NIXL).
+
+        Args:
+            kv_caches: dictionary of layer names, kv cache
+        """
+        if hasattr(self._lmcache_engine, "register_kv_caches"):
+            self._lmcache_engine.register_kv_caches(kv_caches)
+        else:
+            logger.warning(
+                "LMCache engine does not support register_kv_caches, "
+                "please check and use the latest version"
+            )
+
+    def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None:
+        """
+        Start loading the KV cache from the connector to vLLM's paged
+        KV buffer. This is called from the forward context before the
+        forward pass to enable async loading during model execution.
+
+        Args:
+            forward_context (ForwardContext): the forward context.
+            **kwargs: additional arguments for the load operation
+
+        Note:
+            The number of elements in kv_caches and layer_names should be
+            the same.
+
+        """
+        self._lmcache_engine.start_load_kv(forward_context, **kwargs)
+
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        """
+        Block until the KV for a specific layer is loaded into vLLM's
+        paged buffer. This is called from within attention layer to ensure
+        async copying from start_load_kv is complete.
+
+        This interface will be useful for layer-by-layer pipelining.
+
+        Args:
+            layer_name: the name of that layer
+        """
+        self._lmcache_engine.wait_for_layer_load(layer_name)
+
+    def save_kv_layer(
+        self,
+        layer_name: str,
+        kv_layer: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Start saving the a layer of KV cache from vLLM's paged buffer
+        to the connector. This is called from within attention layer to
+        enable async copying during execution.
+
+        Args:
+            layer_name (str): the name of the layer.
+            kv_layer (torch.Tensor): the paged KV buffer of the current
+                layer in vLLM.
+            attn_metadata (AttentionMetadata): the attention metadata.
+            **kwargs: additional arguments for the save operation.
+        """
+        self._lmcache_engine.save_kv_layer(
+            layer_name, kv_layer, attn_metadata, **kwargs
+        )
+
+    def wait_for_save(self):
+        """
+        Block until all the save operations is done. This is called
+        as the forward context exits to ensure that the async saving
+        from save_kv_layer is complete before finishing the forward.
+
+        This prevents overwrites of paged KV buffer before saving done.
+        """
+        self._lmcache_engine.wait_for_save()
+
+    def get_finished(
+        self, finished_req_ids: set[str]
+    ) -> tuple[set[str] | None, set[str] | None]:
+        """
+        Notifies worker-side connector ids of requests that have
+        finished generating tokens.
+
+        Returns:
+            ids of requests that have finished asynchronous transfer
+            (requests that previously returned True from request_finished()),
+            tuple of (sending/saving ids, recving/loading ids).
+            The finished saves/sends req ids must belong to a set provided in a
+            call to this method (this call or a prior one).
+        """
+        return self._lmcache_engine.get_finished(finished_req_ids)
+
+    def get_block_ids_with_load_errors(self) -> set[int]:
+        """
+        Get the set of block IDs that failed to load.
+
+        Returns:
+            Set of block IDs that encountered load errors.
+            Empty set if no load errors occurred.
+        """
+        method = getattr(self._lmcache_engine, "get_block_ids_with_load_errors", None)
+        if callable(method):
+            return method()
+
+        # Fallback for older versions that don't support this method
+        return set()
+
+    def get_kv_connector_kv_cache_events(self) -> LMCacheKVEvents | None:
+        """
+        Get the KV connector kv cache events collected during the last interval.
+        """
+
+        events = self._lmcache_engine.get_kv_events()  # type: ignore [attr-defined]
+        if not events:
+            return None
+
+        blocks: list[BlockStored] = [
+            BlockStored(
+                block_hashes=e.block_hashes,
+                parent_block_hash=e.parent_block_hash,
+                token_ids=e.token_ids,
+                lora_id=e.lora_id,
+                block_size=e.block_size,
+                medium=e.medium,
+                lora_name=getattr(e, "lora_name", None),
+            )
+            for e in events
+        ]
+
+        lmcache_kv_events = LMCacheKVEvents(num_workers=1)
+        lmcache_kv_events.add_events(blocks)
+        return lmcache_kv_events
+
+    # ==============================
+    # Scheduler-side methods
+    # ==============================
+    def get_num_new_matched_tokens(
+        self,
+        request: "Request",
+        num_computed_tokens: int,
+    ) -> tuple[int | None, bool]:
+        """
+        Get number of new tokens that can be loaded from the
+        external KV cache beyond the num_computed_tokens.
+
+        Args:
+            request (Request): the request object.
+            num_computed_tokens (int): the number of locally
+                computed tokens for this request
+
+        Returns:
+            the number of tokens that can be loaded from the
+            external KV cache beyond what is already computed.
+        """
+        return self._lmcache_engine.get_num_new_matched_tokens(
+            request, num_computed_tokens
+        ), False
+
+    def update_state_after_alloc(
+        self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
+    ):
+        """
+        Update KVConnector state after block allocation.
+        """
+        self._lmcache_engine.update_state_after_alloc(request, num_external_tokens)
+
+    def build_connector_meta(
+        self, scheduler_output: SchedulerOutput
+    ) -> KVConnectorMetadata:
+        """
+        Build the connector metadata for this step.
+
+        This function should NOT modify fields in the scheduler_output.
+        Also, calling this function will reset the state of the connector.
+
+        Args:
+            scheduler_output (SchedulerOutput): the scheduler output object.
+        """
+        return self._lmcache_engine.build_connector_meta(scheduler_output)
+
+    def update_connector_output(self, connector_output: KVConnectorOutput):
+        """
+        Update KVConnector state from worker-side connectors output.
+
+        Args:
+            connector_output (KVConnectorOutput): the worker-side
+                connectors output.
+        """
+        # Get the KV events
+        kv_cache_events = connector_output.kv_cache_events
+        if not kv_cache_events or not isinstance(kv_cache_events, LMCacheKVEvents):
+            return
+
+        if self._kv_cache_events is None:
+            self._kv_cache_events = kv_cache_events
+        else:
+            self._kv_cache_events.add_events(kv_cache_events.get_all_events())
+            self._kv_cache_events.increment_workers(
+                kv_cache_events.get_number_of_workers()
+            )
+        return
+
+    def request_finished(
+        self,
+        request: "Request",
+        block_ids: list[int],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        """
+        Called when a request has finished, before its blocks are freed.
+
+        Returns:
+            True if the request is being saved/sent asynchronously and blocks
+            should not be freed until the request_id is returned from
+            get_finished().
+            Optional KVTransferParams to be included in the request outputs
+            returned by the engine.
+        """
+        return self._lmcache_engine.request_finished(request, block_ids)
+
+    def take_events(self) -> Iterable["KVCacheEvent"]:
+        """
+        Take the KV cache events from the connector.
+
+        Yields:
+            New KV cache events since the last call.
+        """
+        if self._kv_cache_events is not None:
+            self._kv_cache_events.aggregate()
+            kv_cache_events = self._kv_cache_events.get_all_events()
+            yield from kv_cache_events
+            self._kv_cache_events.clear_events()
+            self._kv_cache_events = None

vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/__init__.py

@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+from . import multi_process_adapter, vllm_v1_adapter
+from .multi_process_adapter import (
+    LMCacheMPSchedulerAdapter,
+    LMCacheMPWorkerAdapter,
+    LoadStoreOp,
+    ParallelStrategy,
+)
+
+__all__ = [
+    "vllm_v1_adapter",
+    "multi_process_adapter",
+    "LMCacheMPSchedulerAdapter",
+    "LMCacheMPWorkerAdapter",
+    "LoadStoreOp",
+    "ParallelStrategy",
+]

vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/multi_process_adapter.py

@@ -0,0 +1,736 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import os
+from collections.abc import Iterable
+from dataclasses import dataclass
+from itertools import islice
+from typing import Any
+
+import torch
+import zmq
+from lmcache.utils import _lmcache_nvtx_annotate, init_logger
+from lmcache.v1.multiprocess.custom_types import (
+    CudaIPCWrapper,
+    IPCCacheEngineKey,
+    KVCache,
+)
+from lmcache.v1.multiprocess.mq import MessageQueueClient, MessagingFuture
+from lmcache.v1.multiprocess.protocol import RequestType, get_response_class
+
+logger = init_logger(__name__)
+
+
+def wrap_kv_caches(kv_caches: dict[str, torch.Tensor]) -> KVCache:
+    logger.info("KV caches keys are %s", list(kv_caches.keys()))
+    return [CudaIPCWrapper(tensor) for tensor in kv_caches.values()]
+
+
+def striding_block_hashes(
+    block_hashes: list[bytes], blocks_in_chunk: int
+) -> Iterable[bytes]:
+    """Extract chunk-level hashes from block hashes by striding.
+
+    In hash-based vLLM, each vLLM block has its own hash.  LMCache chunks
+    span ``blocks_in_chunk`` consecutive blocks.  The representative hash
+    for a chunk is the hash of the **last** block in that chunk (because
+    each block hash already encodes its prefix).  So we start at index
+    ``blocks_in_chunk - 1`` and stride by ``blocks_in_chunk``.
+    """
+    return islice(block_hashes, blocks_in_chunk - 1, None, blocks_in_chunk)
+
+
+def send_lmcache_request(
+    mq_client: MessageQueueClient,
+    request_type: RequestType,
+    payloads: list[Any],
+) -> MessagingFuture[Any]:
+    """
+    Helper function to send the request to the LMCache multiprocess server
+
+    Args:
+        mq_client: The LMCache multiprocess mode message queue client
+        request_type: The request type
+        payloads: The request payloads
+
+    Returns:
+        A messaging future for the request
+    """
+
+    future = mq_client.submit_request(
+        request_type, payloads, get_response_class(request_type)
+    )
+    return future
+
+
+def get_lmcache_chunk_size(
+    mq_client: MessageQueueClient,
+) -> int:
+    """
+    Helper function to get the LMCache chunk size from the server
+
+    Args:
+        mq_client: The LMCache multiprocess mode message queue client
+
+    Returns:
+        An integer representing the LMCache chunk size
+    """
+    future = send_lmcache_request(mq_client, RequestType.GET_CHUNK_SIZE, [])
+    chunk_size = future.result()
+    return chunk_size
+
+
+@dataclass
+class ParallelStrategy:
+    use_mla: bool
+    """Whether to use the MLA."""
+
+    kv_world_size: int
+    """
+    The kv world size, kv_world_size may not be equal to the actual_world_size, 
+    in the case of mla, it will 'exclude' the effect of TP, the value is 
+    calculated by `extract_world_size_and_kv_rank` in `lmcache_mp_connector.py`.
+    """
+
+    kv_worker_id: int
+    """
+    The kv worker id of the sub-process, kv_worker_id may not be equal to the 
+    actual_worker_id, in the case of mla, it will 'exclude' the effect of TP, 
+    the value is calculated by `extract_world_size_and_kv_rank` in 
+    `lmcache_mp_connector.py`.
+    """
+
+    actual_world_size: int
+    """The actual world size."""
+
+    actual_worker_id: int
+    """The actual worker id of the sub-process."""
+
+    tp_size: int
+    """The tensor parallel size."""
+
+    pp_size: int
+    """The pipeline parallel size."""
+
+
+@dataclass
+class LoadStoreOp:
+    block_ids: list[int]
+    """Block ids for the load/store operation"""
+
+    token_ids: list[int] | None = None
+    """Token IDs for the load/store operation (token mode)"""
+
+    block_hashes: list[bytes] | None = None
+    """Block hashes for the load/store operation (hash mode)"""
+
+    start: int = 0
+    """Start token index (token mode only)"""
+
+    end: int = 0
+    """End token index (token mode only)"""
+
+    def __len__(self) -> int:
+        return len(self.block_ids)
+
+
+StoreResult = bool
+RetrieveResult = list[bool]
+LookupResult = int
+
+
+class LMCacheMPSchedulerAdapter:
+    def __init__(
+        self,
+        server_url: str,
+        context: zmq.Context,
+        model_name: str,
+        vllm_block_size: int,
+        parallel_strategy: ParallelStrategy,
+    ):
+        """
+        Args:
+            server_url: The server URL for the LMCache message queue
+            context: The ZMQ context
+
+            model_name: The model name used for LMCache keys
+            vllm_block_size: The block size used in vLLM
+            parallel_strategy:
+                The parallel strategy, which includes `use_mla`,
+                `world_size`, `worker_id` and so on
+        """
+        self.mq_client = MessageQueueClient(server_url, context)
+
+        # Request futures
+        self.lookup_futures: dict[str, MessagingFuture[LookupResult]] = {}
+
+        self.model_name = model_name
+        self.parallel_strategy = parallel_strategy
+
+        # Read chunk size from lmcache
+        self.chunk_size = get_lmcache_chunk_size(self.mq_client)
+        assert self.chunk_size % vllm_block_size == 0, (
+            "LMCache chunk size should be a multiple of vLLM block size"
+        )
+        self.blocks_in_chunk = self.chunk_size // vllm_block_size
+
+    @property
+    def world_size(self) -> int:
+        """The world size."""
+        return self.parallel_strategy.kv_world_size
+
+    @property
+    def worker_id(self) -> int:
+        """The worker id."""
+        return self.parallel_strategy.kv_worker_id
+
+    @property
+    def tp_size(self) -> int:
+        """The tensor parallel size."""
+        return self.parallel_strategy.tp_size
+
+    @_lmcache_nvtx_annotate
+    def maybe_submit_lookup_request(
+        self,
+        request_id: str,
+        block_hashes: list[bytes] | None = None,
+        token_ids: list[int] | None = None,
+    ) -> None:
+        """
+        Submit a new lookup request to LMCache if there is no ongoing request.
+
+        Supports both token-based and hash-based vLLM:
+        - token_ids: token IDs (token-based vLLM) -> single token-mode key
+        - block_hashes: block hashes (hash-based vLLM) -> strided hash-mode keys
+
+        Exactly one of block_hashes or token_ids must be provided.
+
+        Args:
+            request_id: The ID of the lookup request. The same ID indicates it's
+                from the same request
+            block_hashes: Block hashes to lookup from LMCache (hash mode)
+            token_ids: Token IDs to lookup from LMCache (token mode)
+
+        Returns:
+            None
+
+        Notes:
+            This function will have a side-effect: submitting a look up request to
+            LMCache, which will essentially 'lock' the KV cache chunks in the LMCache
+            for later retrieve operations.
+            In the meantime, this function will record the lookup request, and the
+            status of the look up request can be checked by `check_lookup_result`.
+        """
+        if request_id in self.lookup_futures:
+            # Skip if there is already a lookup request
+            return
+
+        assert (block_hashes is None) != (token_ids is None), (
+            "Exactly one of block_hashes or token_ids must be provided"
+        )
+
+        if block_hashes is not None:
+            # Hash mode: stride block hashes -> N hash-mode keys
+            chunk_hashes = list(
+                striding_block_hashes(block_hashes, self.blocks_in_chunk)
+            )
+            keys = [
+                self._create_hash_key(ch, request_id=request_id) for ch in chunk_hashes
+            ]
+        else:
+            # Token mode: truncate to chunk-aligned length
+            assert token_ids is not None
+            aligned_end = (len(token_ids) // self.chunk_size) * self.chunk_size
+            if aligned_end == 0:
+                return
+            keys = [
+                self._create_key(
+                    token_ids,
+                    start=0,
+                    end=aligned_end,
+                    request_id=request_id,
+                ).no_worker_id_version()
+            ]
+
+        future = send_lmcache_request(
+            self.mq_client,
+            RequestType.LOOKUP,
+            [keys],
+        )
+        self.lookup_futures[request_id] = future
+
+    @_lmcache_nvtx_annotate
+    def check_lookup_result(self, request_id: str) -> int | None:
+        """
+        Check the result of a previously submitted lookup request.
+
+        Args:
+            request_id: The ID of the lookup request submitted in
+                `maybe_submit_lookup_request`
+
+        Returns:
+            An integer representing the total number of tokens matched
+            in LMCache (prefix matching), or
+            None if the lookup request is not finished yet.
+        """
+        assert request_id in self.lookup_futures, (
+            f"Lookup request for request_id={request_id} has not been submitted"
+        )
+
+        future = self.lookup_futures[request_id]
+        if not future.query():
+            return None
+
+        result = future.result()
+        num_chunks = result
+        return num_chunks * self.chunk_size
+
+    def num_blocks_per_chunk(self) -> int:
+        """
+        Returns:
+            The number of vllm blocks in a LMCache data chunk
+        """
+        return self.blocks_in_chunk
+
+    def cleanup_lookup_result(self, request_id: str) -> None:
+        """
+        Clean up lookup future for a finished request to prevent memory leak.
+        Args:
+            request_id: The ID of the finished request.
+        """
+        self.lookup_futures.pop(request_id, None)
+
+    def end_session(self, request_id: str) -> None:
+        """
+        Notify LMCache server to remove the session for a finished request.
+        Args:
+            request_id: The ID of the finished request.
+        """
+        send_lmcache_request(
+            self.mq_client,
+            RequestType.END_SESSION,
+            [request_id],
+        )
+
+    # Helper functions
+    def _create_key(
+        self,
+        token_ids: list[int],
+        start: int = 0,
+        end: int = 0,
+        request_id: str | None = None,
+    ) -> IPCCacheEngineKey:
+        """Convert token IDs to an IPC cache engine key"""
+        return IPCCacheEngineKey(
+            model_name=self.model_name,
+            world_size=self.world_size,
+            worker_id=self.worker_id,
+            token_ids=tuple(token_ids),
+            start=start,
+            end=end,
+            request_id=request_id,
+            tp_size=self.tp_size,
+        )
+
+    def _create_hash_key(
+        self, chunk_hash: bytes, request_id: str | None = None
+    ) -> IPCCacheEngineKey:
+        """Create a hash-mode IPC cache engine key"""
+        return IPCCacheEngineKey(
+            model_name=self.model_name,
+            world_size=self.world_size,
+            worker_id=None,
+            chunk_hash=chunk_hash,
+            request_id=request_id,
+            tp_size=self.tp_size,
+        )
+
+
+class LMCacheMPWorkerAdapter:
+    def __init__(
+        self,
+        server_url: str,
+        context: zmq.Context,
+        model_name: str,
+        vllm_block_size: int,
+        parallel_strategy: ParallelStrategy,
+    ):
+        self.mq_client = MessageQueueClient(server_url, context)
+
+        # Instance id for GPU worker
+        self.instance_id = os.getpid()
+
+        # Registered kv caches from vLLM
+        self.kv_caches: dict[str, torch.Tensor] = {}
+
+        # Request futures
+        # request_id -> (future, other merged requests)
+        self.store_futures: dict[
+            str, tuple[MessagingFuture[StoreResult], list[str]]
+        ] = {}
+        self.retrieve_futures: dict[
+            str, tuple[MessagingFuture[RetrieveResult], list[str]]
+        ] = {}
+
+        # The store requests that have finished execution in LMCache
+        self.finished_stores: set[str] = set()
+        # The finished request ids that are passed via vLLM and also
+        # have corresponding store requests submitted to LMCache before
+        self.previously_finished: set[str] = set()
+
+        self.model_name = model_name
+        self.parallel_strategy = parallel_strategy
+
+        # Read chunk size from lmcache
+        chunk_size = get_lmcache_chunk_size(self.mq_client)
+        assert chunk_size % vllm_block_size == 0, (
+            "LMCache chunk size should be a multiple of vLLM block size"
+        )
+        self.blocks_in_chunk = chunk_size // vllm_block_size
+
+    @property
+    def world_size(self) -> int:
+        """The world size."""
+        return self.parallel_strategy.kv_world_size
+
+    @property
+    def worker_id(self) -> int:
+        """The worker id."""
+        return self.parallel_strategy.kv_worker_id
+
+    @property
+    def use_mla(self) -> bool:
+        """Whether to use MLA."""
+        return self.parallel_strategy.use_mla
+
+    @property
+    def is_first_rank_of_pp_group(self) -> bool:
+        """Is the first rank of the pipeline parallel group."""
+        return (
+            self.parallel_strategy.actual_worker_id % self.parallel_strategy.tp_size
+            == 0
+        )
+
+    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
+        """
+        Register the kv caches with LMCache server
+
+        Args:
+            kv_caches: A dict of kv caches to register. The keys are the
+                layer names and the values are the corresponding tensors.
+        """
+        # Register kv cache and send the request
+        self.kv_caches = kv_caches
+        logger.info("Registering kv caches")
+        future = send_lmcache_request(
+            self.mq_client,
+            RequestType.REGISTER_KV_CACHE,
+            [self.instance_id, wrap_kv_caches(kv_caches)],
+        )
+        future.result()
+
+    @_lmcache_nvtx_annotate
+    def submit_store_request(
+        self, request_id: str, op: LoadStoreOp, event: torch.cuda.Event
+    ):
+        """
+        Submit a KV cache store request to LMCache
+
+        Args:
+            request_id: The ID of the request
+            op: The LoadStoreOp describing the store operation.
+            event: The CUDA event that is recorded after the current
+                model inference step
+        """
+        if op.block_hashes is not None:
+            # Hash mode
+            chunk_hashes = list(
+                striding_block_hashes(op.block_hashes, self.blocks_in_chunk)
+            )
+            keys = [
+                self._create_hash_key(ch, request_id=request_id) for ch in chunk_hashes
+            ]
+        else:
+            # Token mode
+            assert op.token_ids is not None
+            keys = [
+                self._create_key(op.token_ids, op.start, op.end, request_id=request_id)
+            ]
+        future = send_lmcache_request(
+            self.mq_client,
+            RequestType.STORE,
+            [keys, self.instance_id, op.block_ids, event.ipc_handle()],
+        ).to_cuda_future()
+        self.store_futures[request_id] = (future, [])
+
+    @_lmcache_nvtx_annotate
+    def submit_retrieve_request(
+        self, request_id: str, op: LoadStoreOp, event: torch.cuda.Event
+    ):
+        """
+        Submit a KV cache retrieve request to LMCache
+
+        Args:
+            request_id: The ID of the request
+            op: The LoadStoreOp describing the retrieve operation.
+            event: The CUDA event that is recorded after the current
+                model inference step
+        """
+        if op.block_hashes is not None:
+            # Hash mode
+            chunk_hashes = list(
+                striding_block_hashes(op.block_hashes, self.blocks_in_chunk)
+            )
+            keys = [
+                self._create_hash_key(ch, request_id=request_id) for ch in chunk_hashes
+            ]
+        else:
+            # Token mode
+            assert op.token_ids is not None
+            keys = [
+                self._create_key(op.token_ids, op.start, op.end, request_id=request_id)
+            ]
+        future = send_lmcache_request(
+            self.mq_client,
+            RequestType.RETRIEVE,
+            [keys, self.instance_id, op.block_ids, event.ipc_handle()],
+        ).to_cuda_future()
+        self.retrieve_futures[request_id] = (future, [])
+
+    @_lmcache_nvtx_annotate
+    def batched_submit_store_requests(
+        self,
+        request_ids: list[str],
+        ops: list[LoadStoreOp],
+        event: torch.cuda.Event,
+    ):
+        """
+        Submit a batched store request to LMCache
+
+        Args:
+            request_ids: The IDs of the requests
+            ops: The LoadStoreOps describing the store operations. Should have
+                the same length as request_ids
+            event: The CUDA event that is recorded after the current
+                model inference step
+        """
+        all_keys: list[IPCCacheEngineKey] = []
+        block_ids: list[int] = []
+        for request_id, op in zip(request_ids, ops, strict=False):
+            if op.block_hashes is not None:
+                chunk_hashes = list(
+                    striding_block_hashes(op.block_hashes, self.blocks_in_chunk)
+                )
+                keys = [
+                    self._create_hash_key(ch, request_id=request_id)
+                    for ch in chunk_hashes
+                ]
+                all_keys.extend(keys)
+            else:
+                assert op.token_ids is not None
+                all_keys.append(
+                    self._create_key(
+                        op.token_ids, op.start, op.end, request_id=request_id
+                    )
+                )
+            block_ids.extend(op.block_ids)
+        future = send_lmcache_request(
+            self.mq_client,
+            RequestType.STORE,
+            [
+                all_keys,
+                self.instance_id,
+                block_ids,
+                event.ipc_handle(),
+            ],
+        ).to_cuda_future()
+        self.store_futures[request_ids[0]] = (future, list(request_ids[1:]))
+
+    @_lmcache_nvtx_annotate
+    def batched_submit_retrieve_requests(
+        self,
+        request_ids: list[str],
+        ops: list[LoadStoreOp],
+        event: torch.cuda.Event,
+    ):
+        """
+        Submit a batched retrieve request to LMCache
+
+        Args:
+            request_ids: The IDs of the requests
+            ops: The LoadStoreOps describing the retrieve operations. Should have
+                the same length as request_ids
+            event: The CUDA event that is recorded after the current
+                model inference step
+        """
+        all_keys: list[IPCCacheEngineKey] = []
+        block_ids: list[int] = []
+        for request_id, op in zip(request_ids, ops, strict=False):
+            if op.block_hashes is not None:
+                chunk_hashes = list(
+                    striding_block_hashes(op.block_hashes, self.blocks_in_chunk)
+                )
+                keys = [
+                    self._create_hash_key(ch, request_id=request_id)
+                    for ch in chunk_hashes
+                ]
+                all_keys.extend(keys)
+            else:
+                assert op.token_ids is not None
+                all_keys.append(
+                    self._create_key(
+                        op.token_ids, op.start, op.end, request_id=request_id
+                    )
+                )
+            block_ids.extend(op.block_ids)
+        future = send_lmcache_request(
+            self.mq_client,
+            RequestType.RETRIEVE,
+            [
+                all_keys,
+                self.instance_id,
+                block_ids,
+                event.ipc_handle(),
+            ],
+        ).to_cuda_future()
+        self.retrieve_futures[request_ids[0]] = (future, list(request_ids[1:]))
+
+    @_lmcache_nvtx_annotate
+    def get_finished(
+        self, finished_req_ids_from_engine: set[str]
+    ) -> tuple[set[str] | None, set[str] | None]:
+        """
+        Check and get the finished store and retrieve requests.
+
+        Args:
+            finished_req_ids_from_engine: the set of request ids that are
+                reported as finished from the vLLM engine side.
+
+        Returns:
+            A tuple of two sets:
+            - The first set contains the finished store request ids. The returned
+                store request ids MUST be seen before in the
+                `finished_req_ids_from_engine`.
+            - The second set contains the finished retrieve request ids.
+
+        Notes:
+            When enabling async scheduling in vLLM, the same request ID may appear
+            multiple times in `finished_req_ids_from_engine`. The adapter should
+            take care of deduplicating the request IDs and only return the request
+            IDs that have not been returned before.
+        """
+        finished_stores = set()
+        finished_retrieves = set()
+        for request_id, (s_future, other_reqs) in self.store_futures.items():
+            if not s_future.query():
+                continue
+
+            s_result = s_future.result()
+            finished_stores.add(request_id)
+            finished_stores.update(other_reqs)
+
+            if not s_result:
+                # TODO: add error handling here
+                logger.error(
+                    "Something went wrong when processing the "
+                    "store request for request_id=%s",
+                    request_id,
+                )
+
+        for request_id, (r_future, other_reqs) in self.retrieve_futures.items():
+            if not r_future.query():
+                continue
+
+            r_result = r_future.result()
+            finished_retrieves.add(request_id)
+            finished_retrieves.update(other_reqs)
+
+            if not all(r_result):
+                # TODO: add error handing here
+                logger.error(
+                    "Something went wrong when processing the "
+                    "retrieve request for request_id=%s, result=%s",
+                    request_id,
+                    r_result,
+                )
+
+        # Remove the finished requests from the tracking dicts
+        for request_id in finished_stores:
+            self.store_futures.pop(request_id, None)
+        for request_id in finished_retrieves:
+            self.retrieve_futures.pop(request_id, None)
+
+        # Update the internal states
+        self.finished_stores.update(finished_stores)
+
+        ret_stores = set()
+        for req_id in finished_req_ids_from_engine:
+            if req_id in self.finished_stores or req_id in self.store_futures:
+                self.previously_finished.add(req_id)
+            else:
+                ret_stores.add(req_id)
+
+        # Calculate the final finished stores
+        ret_stores.update(self._update_and_get_finished_store())
+
+        return ret_stores, finished_retrieves
+
+    def num_blocks_per_chunk(self) -> int:
+        """
+        Returns:
+            The number of vllm blocks in a LMCache data chunk
+        """
+        return self.blocks_in_chunk
+
+    def shutdown(self):
+        """
+        Shutdown the LMCache MP worker adapter
+        """
+        logger.info("Unregistering kv caches")
+        send_lmcache_request(
+            self.mq_client, RequestType.UNREGISTER_KV_CACHE, [self.instance_id]
+        ).result()
+
+        self.mq_client.close()
+
+    # Helper functions
+    def _update_and_get_finished_store(
+        self,
+    ) -> set[str]:
+        """Converge the internal states about finished stores
+        and returns the 'safe finished store request ids' back
+        """
+        safe_finished_s = self.finished_stores.intersection(self.previously_finished)
+        self.finished_stores.difference_update(self.previously_finished)
+        self.previously_finished.difference_update(safe_finished_s)
+
+        return safe_finished_s
+
+    def _create_key(
+        self,
+        token_ids: list[int],
+        start: int = 0,
+        end: int = 0,
+        request_id: str | None = None,
+    ) -> IPCCacheEngineKey:
+        """Convert token IDs to an IPC cache engine key"""
+        return IPCCacheEngineKey(
+            model_name=self.model_name,
+            world_size=self.world_size,
+            worker_id=self.worker_id,
+            token_ids=tuple(token_ids),
+            start=start,
+            end=end,
+            request_id=request_id,
+        )
+
+    def _create_hash_key(
+        self, chunk_hash: bytes, request_id: str | None = None
+    ) -> IPCCacheEngineKey:
+        """Create a hash-mode IPC cache engine key"""
+        return IPCCacheEngineKey(
+            model_name=self.model_name,
+            world_size=self.world_size,
+            worker_id=self.worker_id,
+            chunk_hash=chunk_hash,
+            request_id=request_id,
+        )

vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/utils.py

@@ -0,0 +1,211 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Standard
+import os
+import threading
+from typing import TYPE_CHECKING, Union
+
+import torch
+from lmcache.logging import init_logger
+from lmcache.v1.config import LMCacheEngineConfig as V1Config
+
+if TYPE_CHECKING:
+    from vllm.config import ModelConfig
+    from vllm.multimodal.inputs import PlaceholderRange
+    from vllm.v1.core.sched.output import NewRequestData
+    from vllm.v1.request import Request
+
+logger = init_logger(__name__)
+ENGINE_NAME = "vllm-instance"
+
+# Thread-safe singleton storage
+_config_instance: V1Config | None = None
+_config_lock = threading.Lock()
+
+
+def is_false(value: str) -> bool:
+    """Check if the given string value is equivalent to 'false'."""
+    return value.lower() in ("false", "0", "no", "n", "off")
+
+
+def lmcache_get_or_create_config() -> V1Config:
+    """Get the LMCache configuration from the environment variable
+    `LMCACHE_CONFIG_FILE`. If the environment variable is not set, this
+    function will return the default configuration.
+
+    This function is thread-safe and implements singleton pattern,
+    ensuring the configuration is loaded only once.
+    """
+    global _config_instance
+
+    # Double-checked locking for thread-safe singleton
+    if _config_instance is None:
+        with _config_lock:
+            if _config_instance is None:  # Check again within lock
+                LMCacheEngineConfig = V1Config  # type: ignore[assignment]
+
+                if "LMCACHE_CONFIG_FILE" not in os.environ:
+                    logger.warning(
+                        "No LMCache configuration file is set. Trying to read"
+                        " configurations from the environment variables."
+                    )
+                    logger.warning(
+                        "You can set the configuration file through "
+                        "the environment variable: LMCACHE_CONFIG_FILE"
+                    )
+                    _config_instance = LMCacheEngineConfig.from_env()
+                else:
+                    config_file = os.environ["LMCACHE_CONFIG_FILE"]
+                    logger.info("Loading LMCache config file %s", config_file)
+                    _config_instance = LMCacheEngineConfig.from_file(config_file)
+                    # Update config from environment variables
+                    _config_instance.update_config_from_env()
+    return _config_instance
+
+
+def hex_hash_to_int16(s: str) -> int:
+    """
+    Convert a hex hash string to a 16-bit integer.
+    """
+    return int(s, 16) & 0xFFFF
+
+
+def apply_mm_hashes_to_token_ids(
+    token_ids: torch.Tensor,
+    mm_hashes: list[str],
+    mm_positions: list["PlaceholderRange"],
+) -> torch.Tensor:
+    """
+    Overwrite token_ids in-place for multimodal placeholders using
+    efficient slice assignments.
+    """
+    n = token_ids.size(0)
+    for hash_str, placeholder in zip(mm_hashes, mm_positions):
+        start, length = placeholder.offset, placeholder.length
+        if start >= n:
+            continue
+        end = min(start + length, n)
+        token_ids[start:end] = hex_hash_to_int16(hash_str)
+    return token_ids
+
+
+def mla_enabled(model_config: "ModelConfig") -> bool:
+    return (
+        hasattr(model_config, "use_mla")
+        and isinstance(model_config.use_mla, bool)
+        and model_config.use_mla
+    )
+
+
+def create_lmcache_metadata(
+    vllm_config=None, model_config=None, parallel_config=None, cache_config=None
+):
+    """
+    Create LMCacheEngineMetadata from vLLM configuration.
+
+    This function extracts common metadata creation logic that was duplicated
+    across multiple files.
+
+    Args:
+        vllm_config (VllmConfig): vLLM configuration object containing model,
+                                  parallel, and cache configs (alternative to
+                                  individual config parameters)
+        model_config (ModelConfig): Model configuration (alternative to
+                                    vllm_config)
+        parallel_config (ParallelConfig): Parallel configuration (alternative
+                                          to vllm_config)
+        cache_config (CacheConfig): Cache configuration (alternative to
+                                    vllm_config)
+    """
+    # Third Party
+    # First Party
+    from lmcache.config import LMCacheEngineMetadata
+
+    from vllm.utils.torch_utils import get_kv_cache_torch_dtype
+
+    config = lmcache_get_or_create_config()
+    # Support both vllm_config object and individual config parameters
+    if vllm_config is not None:
+        model_cfg = vllm_config.model_config
+        parallel_cfg = vllm_config.parallel_config
+        cache_cfg = vllm_config.cache_config
+    else:
+        if model_config is None or parallel_config is None or cache_config is None:
+            raise ValueError(
+                "Either vllm_config must be provided, or all of "
+                "model_config, parallel_config, and cache_config must be provided."
+            )
+        model_cfg = model_config
+        parallel_cfg = parallel_config
+        cache_cfg = cache_config
+
+    # Get KV cache dtype
+    kv_dtype = get_kv_cache_torch_dtype(cache_cfg.cache_dtype, model_cfg.dtype)
+
+    # Check if MLA is enabled
+    use_mla = mla_enabled(model_cfg)
+
+    # Construct KV shape (for memory pool)
+    num_layer = model_cfg.get_num_layers(parallel_cfg)
+    chunk_size = config.chunk_size
+    num_kv_head = model_cfg.get_num_kv_heads(parallel_cfg)
+    head_size = model_cfg.get_head_size()
+    kv_shape = (num_layer, 1 if use_mla else 2, chunk_size, num_kv_head, head_size)
+
+    # Create metadata
+    metadata = LMCacheEngineMetadata(
+        model_cfg.model,
+        parallel_cfg.world_size,
+        parallel_cfg.rank,
+        "vllm",
+        kv_dtype,
+        kv_shape,
+        use_mla,
+    )
+
+    return metadata, config
+
+
+def extract_mm_features(
+    request: Union["Request", "NewRequestData"], modify: bool = False
+) -> tuple[list[str], list["PlaceholderRange"]]:
+    """
+    Normalize multimodal information from a Request into parallel lists.
+
+    This helper reads either:
+      1) `request.mm_features` (objects each exposing `.identifier` and
+      `.mm_position`), or
+      2) legacy fields `request.mm_hashes` and `request.mm_positions`.
+
+    It returns two equally sized lists: the multimodal hash identifiers and
+    their corresponding positions. If the request contains no multimodal info,
+    it returns `([], [])`.
+
+    Args:
+        request (Request): The source object.
+        modify (bool):
+            Controls copy semantics for the legacy-path return values.
+            - If True and legacy fields are used, shallow-copies are returned so
+              the caller can mutate the lists without affecting `request`.
+            - If False, the original legacy sequences are returned as-is
+              (zero-copy); treat them as read-only.
+
+    Returns:
+        tuple[list[str], list[PlaceholderRange]]: (`mm_hashes`, `mm_positions`).
+        May be `([], [])` when no multimodal data is present.
+    """
+    if getattr(request, "mm_features", None):
+        mm_hashes, mm_positions = zip(
+            *((f.identifier, f.mm_position) for f in request.mm_features)
+        )
+        return (list(mm_hashes), list(mm_positions))
+    elif getattr(request, "mm_hashes", None):
+        if modify:
+            return (
+                request.mm_hashes.copy(),  # type: ignore
+                request.mm_positions.copy(),  # type: ignore
+            )
+        else:
+            return (request.mm_hashes, request.mm_positions)  # type: ignore
+    else:
+        return ([], [])

vllm/distributed/kv_transfer/kv_connector/v1/metrics.py

@@ -0,0 +1,175 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import dataclass, field
+from typing import Any, TypeAlias, TypeVar
+
+from prometheus_client import Counter, Gauge, Histogram
+
+from vllm.config import KVTransferConfig, VllmConfig
+from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
+from vllm.logger import init_logger
+
+PromMetric: TypeAlias = Gauge | Counter | Histogram
+PromMetricT = TypeVar("PromMetricT", bound=PromMetric)
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class KVConnectorStats:
+    """
+    Base class for KV Connector Stats, a container for transfer performance
+    metrics or otherwise important telemetry from the connector.
+    All sub-classes need to be serializable as stats are sent from worker to
+    logger process.
+    """
+
+    data: dict[str, Any] = field(default_factory=dict)
+
+    def reset(self):
+        """Reset the stats, clear the state."""
+        raise NotImplementedError
+
+    def aggregate(self, other: "KVConnectorStats") -> "KVConnectorStats":
+        """
+        Aggregate stats with another `KVConnectorStats` object.
+        """
+        raise NotImplementedError
+
+    def reduce(self) -> dict[str, int | float]:
+        """
+        Reduce the observations collected during a time interval to one or
+        more representative values (eg avg/median/sum of the series).
+        This is meant to be called by the logger to produce a summary of the
+        stats for the last time interval.
+        """
+        raise NotImplementedError
+
+    def is_empty(self) -> bool:
+        """Return True if the stats are empty."""
+        raise NotImplementedError
+
+
+class KVConnectorLogging:
+    def __init__(self, kv_transfer_config: KVTransferConfig | None):
+        # Instantiate the connector's stats class.
+        if kv_transfer_config and kv_transfer_config.kv_connector:
+            self.connector_cls = KVConnectorFactory.get_connector_class(
+                kv_transfer_config
+            )
+        self.reset()
+
+    def reset(self):
+        self.transfer_stats_accumulator: KVConnectorStats | None = None
+
+    def observe(self, transfer_stats_data: dict[str, Any]):
+        # Should not be called when a KVConnector is not configured.
+        assert self.connector_cls is not None
+        # Called periodically when connector syncs with the scheduler.
+        # Note that this is not the same as the logging interval.
+        # We expect transfer_stats_data to be aggregated across all workers and
+        # consist of observations from a single connector or a MultiConnector.
+        transfer_stats = self.connector_cls.build_kv_connector_stats(
+            transfer_stats_data
+        )
+        if transfer_stats is None:
+            logger.warning_once(
+                "The connector %s is collecting stats but "
+                "does not implement the "
+                "`build_kv_connector_stats` method. "
+                "Stats will not be logged.",
+                self.connector_cls,
+            )
+            return
+
+        if self.transfer_stats_accumulator is None:
+            self.transfer_stats_accumulator = transfer_stats
+        else:
+            # Accumulate last interval stats.
+            self.transfer_stats_accumulator = self.transfer_stats_accumulator.aggregate(
+                transfer_stats
+            )
+
+    def log(self, log_fn=logger.info):
+        """Log transfer metrics periodically, similar to throughput logging"""
+        if (
+            self.transfer_stats_accumulator
+            and not self.transfer_stats_accumulator.is_empty()
+        ):
+            # Produce a single cumulative stats object for the last time
+            # interval from the recorded observations.
+            xfer_metrics = self.transfer_stats_accumulator.reduce()
+            xfer_metrics_str = ", ".join(f"{k}={v}" for k, v in xfer_metrics.items())
+            log_fn("KV Transfer metrics: %s", xfer_metrics_str)
+
+            # Reset metrics for next interval
+            self.reset()
+
+
+class KVConnectorPromMetrics:
+    """
+    A base class for per-connector Prometheus metric registration
+    and recording.
+    """
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        metric_types: dict[type[PromMetric], type[PromMetricT]],
+        labelnames: list[str],
+        per_engine_labelvalues: dict[int, list[object]],
+    ):
+        self._kv_transfer_config = vllm_config.kv_transfer_config
+        self._gauge_cls = metric_types[Gauge]
+        self._counter_cls = metric_types[Counter]
+        self._histogram_cls = metric_types[Histogram]
+        self._labelnames = labelnames
+        self.per_engine_labelvalues = per_engine_labelvalues
+
+    def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
+        """
+        Record the supplied transfer statistics to Prometheus metrics. These
+        statistics are engine-specific, and should be recorded to a metric
+        with the appropriate 'engine' label. These metric instances can be
+        created using the create_metric_per_engine() helper method.
+        """
+        raise NotImplementedError
+
+
+class KVConnectorProm:
+    """
+    Support for registering per-connector Prometheus metrics, and
+    recording transfer statistics to those metrics. Uses
+    KVConnectorBase.build_prom_metrics().
+    """
+
+    _gauge_cls = Gauge
+    _counter_cls = Counter
+    _histogram_cls = Histogram
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        labelnames: list[str],
+        per_engine_labelvalues: dict[int, list[object]],
+    ):
+        self.prom_metrics: KVConnectorPromMetrics | None = None
+        kv_transfer_config = vllm_config.kv_transfer_config
+        if kv_transfer_config and kv_transfer_config.kv_connector:
+            connector_cls = KVConnectorFactory.get_connector_class(kv_transfer_config)
+            metric_types = {
+                Gauge: self._gauge_cls,
+                Counter: self._counter_cls,
+                Histogram: self._histogram_cls,
+            }
+            self.prom_metrics = connector_cls.build_prom_metrics(
+                vllm_config,
+                metric_types,
+                labelnames,
+                per_engine_labelvalues,
+            )
+
+    def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
+        if self.prom_metrics is None:
+            return
+        self.prom_metrics.observe(transfer_stats_data, engine_idx)

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py

@@ -0,0 +1,136 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import threading
+import time
+from dataclasses import dataclass
+
+import uvicorn
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+
+from vllm.config import ParallelConfig
+from vllm.distributed.kv_transfer.kv_connector.utils import EngineId
+from vllm.logger import init_logger
+
+WorkerAddr = str
+
+logger = init_logger(__name__)
+
+
+def get_mooncake_dp_engine_index(parallel_config: ParallelConfig) -> int:
+    """Return the per-engine DP index used for Mooncake side channels."""
+    if parallel_config.local_engines_only:
+        assert parallel_config.data_parallel_rank_local is not None
+        return parallel_config.data_parallel_rank_local
+
+    return parallel_config.data_parallel_index
+
+
+class RegisterWorkerPayload(BaseModel):
+    engine_id: EngineId
+    dp_rank: int
+    tp_rank: int
+    pp_rank: int
+    addr: WorkerAddr
+
+
+@dataclass
+class EngineEntry:
+    engine_id: EngineId
+    # {tp_rank: {pp_rank: worker_addr}}
+    worker_addr: dict[int, dict[int, WorkerAddr]]
+
+
+class MooncakeBootstrapServer:
+    """
+    A centralized server running on the global rank 0 prefiller worker.
+    Prefiller workers register their connection info (IP, port, ranks) here.
+    """
+
+    def __init__(self, host: str, port: int):
+        self.workers: dict[int, EngineEntry] = {}
+
+        self.host = host
+        self.port = port
+        self.app = FastAPI()
+        self._register_routes()
+        self.server_thread: threading.Thread | None = None
+        self.server: uvicorn.Server | None = None
+
+    def __del__(self):
+        self.shutdown()
+
+    def _register_routes(self):
+        # All methods are async. No need to use lock to protect data.
+        self.app.post("/register")(self.register_worker)
+        self.app.get("/query", response_model=dict[int, EngineEntry])(self.query)
+
+    def start(self):
+        if self.server_thread:
+            return
+
+        config = uvicorn.Config(app=self.app, host=self.host, port=self.port)
+        self.server = uvicorn.Server(config=config)
+        self.server_thread = threading.Thread(
+            target=self.server.run, name="mooncake_bootstrap_server", daemon=True
+        )
+        self.server_thread.start()
+        while not self.server.started:
+            time.sleep(0.1)  # Wait for the server to start
+        logger.info("Mooncake Bootstrap Server started at %s:%d", self.host, self.port)
+
+    def shutdown(self):
+        if self.server_thread is None or self.server is None or not self.server.started:
+            return
+
+        self.server.should_exit = True
+        self.server_thread.join()
+        logger.info("Mooncake Bootstrap Server stopped.")
+
+    async def register_worker(self, payload: RegisterWorkerPayload):
+        """Handles registration of a prefiller worker."""
+        if payload.dp_rank not in self.workers:
+            self.workers[payload.dp_rank] = EngineEntry(
+                engine_id=payload.engine_id,
+                worker_addr={},
+            )
+
+        dp_entry = self.workers[payload.dp_rank]
+        if dp_entry.engine_id != payload.engine_id:
+            raise HTTPException(
+                status_code=400,
+                detail=(
+                    f"Engine ID mismatch for dp_rank={payload.dp_rank}: "
+                    f"expected {dp_entry.engine_id}, got {payload.engine_id}"
+                ),
+            )
+        if payload.tp_rank not in dp_entry.worker_addr:
+            dp_entry.worker_addr[payload.tp_rank] = {}
+
+        tp_entry = dp_entry.worker_addr[payload.tp_rank]
+        if payload.pp_rank in tp_entry:
+            raise HTTPException(
+                status_code=400,
+                detail=(
+                    f"Worker with dp_rank={payload.dp_rank}, "
+                    f"tp_rank={payload.tp_rank}, pp_rank={payload.pp_rank} "
+                    f"is already registered at "
+                    f"{tp_entry[payload.pp_rank]}, "
+                    f"but still want to register at {payload.addr}"
+                ),
+            )
+
+        tp_entry[payload.pp_rank] = payload.addr
+        logger.debug(
+            "Registered worker: engine_id=%s, dp_rank=%d, tp_rank=%d, pp_rank=%d at %s",
+            payload.engine_id,
+            payload.dp_rank,
+            payload.tp_rank,
+            payload.pp_rank,
+            payload.addr,
+        )
+
+        return {"status": "ok"}
+
+    async def query(self) -> dict[int, EngineEntry]:
+        return self.workers

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/rdma_utils.py

@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Mooncake requester config helpers."""
+
+from collections.abc import Mapping
+from typing import Any
+
+import vllm.envs as envs
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def normalize_string_override(value: Any) -> str | None:
+    if not isinstance(value, str):
+        return None
+    normalized = value.strip()
+    return normalized or None
+
+
+def get_requester_local_hostname(local_ip: str) -> str:
+    override = normalize_string_override(envs.MOONCAKE_REQUESTER_LOCAL_HOSTNAME)
+    if override is not None:
+        return override
+    return local_ip
+
+
+def get_configured_preferred_segment(
+    extra_config: Mapping[str, Any],
+) -> str | None:
+    preferred_segment = normalize_string_override(extra_config.get("preferred_segment"))
+    if preferred_segment is not None:
+        return preferred_segment
+    if extra_config.get("preferred_segment") is not None:
+        raise ValueError(
+            "Mooncake preferred_segment override must be a non-empty string"
+        )
+
+    env_value = normalize_string_override(envs.MOONCAKE_PREFERRED_SEGMENT)
+    if env_value is not None:
+        logger.info(
+            "Mooncake preferred_segment from MOONCAKE_PREFERRED_SEGMENT: %s",
+            env_value,
+        )
+        return env_value
+    return None

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/stats.py

@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Stats container for the Mooncake connector."""
+
+import threading
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
+    KVConnectorStats,
+)
+
+# TODO(mooncake-stats): add MooncakePromMetrics (mirror NixlPromMetrics)
+# and wire it via MooncakeConnector.build_prom_metrics in a follow-up PR.
+
+
+@dataclass
+class MooncakeKVConnectorStats(KVConnectorStats):
+    """Container for Mooncake KV transfer performance metrics.
+
+    `_lock` serializes record_* against clone_and_reset so each row's
+    appends are atomic and column lengths stay aligned. Writers run on
+    the sender pool / receiver loop / sender loop; reader runs on the
+    main worker thread.
+    """
+
+    def __post_init__(self):
+        self._lock = threading.Lock()
+        if not self.data:
+            self.reset()
+
+    # threading.Lock is not picklable; strip it from the wire form and
+    # rebuild a fresh per-process lock on the receiver side.
+    def __getstate__(self) -> dict[str, Any]:
+        state = self.__dict__.copy()
+        state.pop("_lock", None)
+        return state
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        self.__dict__.update(state)
+        self._lock = threading.Lock()
+
+    def reset(self):
+        self.data: dict[str, list[float | int]] = {
+            "transfer_duration": [],
+            "bytes_transferred": [],
+            "num_descriptors": [],
+            "num_failed_transfers": [],
+            "num_failed_recvs": [],
+            "num_kv_expired_reqs": [],
+        }
+
+    def record_transfer(self, duration_s: float, total_bytes: int, num_descs: int):
+        with self._lock:
+            self.data["transfer_duration"].append(duration_s)
+            self.data["bytes_transferred"].append(total_bytes)
+            self.data["num_descriptors"].append(num_descs)
+
+    # Failure counters store a list of 1s so a future Prom counter can iterate
+    # with .inc(list_item), mirroring NIXL's NixlPromMetrics.observe.
+    def record_failed_transfer(self):
+        with self._lock:
+            self.data["num_failed_transfers"].append(1)
+
+    def record_failed_recv(self):
+        with self._lock:
+            self.data["num_failed_recvs"].append(1)
+
+    def record_kv_expired_req(self):
+        with self._lock:
+            self.data["num_kv_expired_reqs"].append(1)
+
+    def clone_and_reset(self) -> "MooncakeKVConnectorStats":
+        # Copy lists under the lock for length alignment; return a fresh
+        # instance so the snapshot has its own _lock.
+        with self._lock:
+            snapshot_data: dict[str, list[float | int]] = {
+                k: list(v) for k, v in self.data.items()
+            }
+            self.reset()
+        return MooncakeKVConnectorStats(data=snapshot_data)
+
+    def is_empty(self) -> bool:
+        return (
+            self.num_successful_transfers == 0
+            and len(self.data["num_failed_transfers"]) == 0
+            and len(self.data["num_failed_recvs"]) == 0
+            and len(self.data["num_kv_expired_reqs"]) == 0
+        )
+
+    def aggregate(self, other: KVConnectorStats) -> KVConnectorStats:
+        if not other.is_empty():
+            for k, v in other.data.items():
+                accumulator = self.data[k]
+                assert isinstance(accumulator, list)
+                accumulator.extend(v)
+        return self
+
+    def reduce(self) -> dict[str, int | float]:
+        num_failed_transfers = len(self.data["num_failed_transfers"])
+        num_failed_recvs = len(self.data["num_failed_recvs"])
+        num_kv_expired_reqs = len(self.data["num_kv_expired_reqs"])
+
+        if self.num_successful_transfers == 0:
+            return {
+                "Num successful transfers": 0,
+                "Avg xfer time (ms)": 0,
+                "P90 xfer time (ms)": 0,
+                "Avg MB per transfer": 0,
+                "Throughput (MB/s)": 0,
+                "Avg number of descriptors": 0,
+                "Num failed transfers": num_failed_transfers,
+                "Num failed recvs": num_failed_recvs,
+                "Num KV expired reqs": num_kv_expired_reqs,
+            }
+
+        xfer_time = np.asarray(self.data["transfer_duration"])
+        mb = np.asarray(self.data["bytes_transferred"]) / 2**20
+        descs = np.asarray(self.data["num_descriptors"], dtype=np.uint32)
+        n = len(descs)
+        assert n == self.num_successful_transfers
+
+        total_mb = mb.sum()
+        avg_mb = total_mb / n
+        total_time_seconds = xfer_time.sum()
+        throughput_mb_s = (
+            total_mb / total_time_seconds if total_time_seconds > 0 else 0.0
+        )
+
+        return {
+            "Num successful transfers": n,
+            "Avg xfer time (ms)": round(xfer_time.mean() * 1e3, 3),
+            "P90 xfer time (ms)": round(np.percentile(xfer_time, 90).item() * 1e3, 3),
+            "Avg MB per transfer": round(avg_mb, 3),
+            "Throughput (MB/s)": round(throughput_mb_s, 3),
+            "Avg number of descriptors": round(descs.mean(), 1),
+            "Num failed transfers": num_failed_transfers,
+            "Num failed recvs": num_failed_recvs,
+            "Num KV expired reqs": num_kv_expired_reqs,
+        }
+
+    @property
+    def num_successful_transfers(self) -> int:
+        return len(self.data["transfer_duration"])

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py

@@ -0,0 +1,346 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Adapted from vllm-project/vllm-ascend
+# (vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/).
+"""MooncakeStoreConnector - KV cache connector using MooncakeDistributedStore.
+
+Unlike MooncakeConnector which does direct P2P transfer, this connector
+uses MooncakeDistributedStore as a shared KV cache pool. Both producer
+and consumer instances read/write KV to/from the store independently,
+enabling prefix caching via hash-based deduplication.
+"""
+
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.distributed.kv_events import (
+    KVCacheEvent,
+    KVConnectorKVEvents,
+    KVEventAggregator,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorBase_V1,
+    KVConnectorMetadata,
+    KVConnectorRole,
+    SupportsHMA,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
+    KVConnectorPromMetrics,
+    KVConnectorStats,
+    PromMetric,
+    PromMetricT,
+)
+from vllm.forward_context import ForwardContext
+from vllm.logger import init_logger
+from vllm.v1.attention.backend import AttentionMetadata
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.outputs import KVConnectorOutput
+from vllm.v1.request import Request
+
+from .data import MooncakeStoreConnectorMetadata
+from .metrics import MooncakeStoreConnectorStats, MooncakeStorePromMetrics
+from .scheduler import MooncakeStoreScheduler
+from .worker import MooncakeStoreWorker
+
+logger = init_logger(__name__)
+
+
+class MooncakeStoreKVEvents(KVConnectorKVEvents):
+    """KV event aggregation for MooncakeStoreConnector."""
+
+    def __init__(self, num_workers: int) -> None:
+        self._aggregator = KVEventAggregator(num_workers)
+
+    def add_events(self, events: list[KVCacheEvent]) -> None:
+        self._aggregator.add_events(events)
+
+    def aggregate(self) -> "MooncakeStoreKVEvents":
+        common_events = self._aggregator.get_common_events()
+        self._aggregator.clear_events()
+        self._aggregator.add_events(common_events)
+        self._aggregator.reset_workers()
+        return self
+
+    def increment_workers(self, count: int = 1) -> None:
+        self._aggregator.increment_workers(count)
+
+    def get_all_events(self) -> list[KVCacheEvent]:
+        return self._aggregator.get_all_events()
+
+    def get_number_of_workers(self) -> int:
+        return self._aggregator.get_number_of_workers()
+
+    def clear_events(self) -> None:
+        self._aggregator.clear_events()
+        self._aggregator.reset_workers()
+
+    def __repr__(self) -> str:
+        return f"<MooncakeStoreKVEvents events={self.get_all_events()}>"
+
+
+class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):
+    """KV connector using MooncakeDistributedStore as shared KV pool."""
+
+    @property
+    def prefer_cross_layer_blocks(self) -> bool:
+        extra_config = self._kv_transfer_config.kv_connector_extra_config
+        return (
+            str(extra_config.get("enable_cross_layers_blocks", "False")).lower()
+            == "true"
+        )
+
+    @staticmethod
+    def _validate_kv_cache_config(
+        vllm_config: VllmConfig, kv_cache_config: KVCacheConfig
+    ) -> None:
+        from vllm.v1.kv_cache_interface import CrossAttentionSpec, MambaSpec
+
+        unsupported: list[str] = []
+        cache_block_size = vllm_config.cache_config.block_size
+        for g_idx, g in enumerate(kv_cache_config.kv_cache_groups):
+            spec = g.kv_cache_spec
+            if isinstance(spec, CrossAttentionSpec):
+                unsupported.append(f"group {g_idx}: CrossAttentionSpec")
+            # Enforce Mamba align mode
+            if isinstance(spec, MambaSpec) and spec.block_size != cache_block_size:
+                unsupported.append(
+                    f"group {g_idx}: MambaSpec with block_size="
+                    f"{spec.block_size} != cache_config.block_size="
+                    f"{cache_block_size} (mamba_cache_mode != 'align')"
+                )
+        pcp = vllm_config.parallel_config.prefill_context_parallel_size
+        dcp = vllm_config.parallel_config.decode_context_parallel_size
+        if len(kv_cache_config.kv_cache_groups) > 1 and pcp * dcp > 1:
+            unsupported.append(
+                f"PCP/DCP > 1 (pcp={pcp}, dcp={dcp}) with hybrid attention"
+            )
+        if unsupported:
+            raise ValueError(
+                "MooncakeStoreConnector does not support: " + "; ".join(unsupported)
+            )
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        role: KVConnectorRole,
+        kv_cache_config: KVCacheConfig | None = None,
+    ):
+        super().__init__(
+            vllm_config=vllm_config,
+            role=role,
+            kv_cache_config=kv_cache_config,  # type: ignore[arg-type]
+        )
+        assert vllm_config.kv_transfer_config is not None
+        assert kv_cache_config is not None, "kv_cache_config is required"
+        self._validate_kv_cache_config(vllm_config, kv_cache_config)
+        self._kv_cache_config = kv_cache_config
+        self.kv_role = vllm_config.kv_transfer_config.kv_role
+        self._kv_cache_events: MooncakeStoreKVEvents | None = None
+
+        self.connector_scheduler: MooncakeStoreScheduler | None = None
+        self.connector_worker: MooncakeStoreWorker | None = None
+
+        if role == KVConnectorRole.SCHEDULER:
+            self.connector_scheduler = MooncakeStoreScheduler(
+                vllm_config, kv_cache_config
+            )
+        else:
+            self.connector_worker = MooncakeStoreWorker(vllm_config, kv_cache_config)
+
+    def shutdown(self):
+        """Release connector resources on teardown.
+
+        Closes the worker's MooncakeDistributedStore handle so its
+        TransferEngine and RDMA registrations are released. Invoked from the
+        engine's explicit shutdown path and as a backstop from ``__del__``;
+        a no-op on the scheduler role, which holds no store handle.
+        """
+        worker = getattr(self, "connector_worker", None)
+        if worker is not None:
+            worker.close()
+
+    def __del__(self):
+        self.shutdown()
+
+    # ============================================================
+    # Scheduler-side methods
+    # ============================================================
+
+    def get_num_new_matched_tokens(
+        self,
+        request: Request,
+        num_computed_tokens: int,
+    ) -> tuple[int | None, bool]:
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.get_num_new_matched_tokens(
+            request, num_computed_tokens
+        )
+
+    def update_state_after_alloc(
+        self,
+        request: Request,
+        blocks: KVCacheBlocks,
+        num_external_tokens: int,
+    ):
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.update_state_after_alloc(
+            request, blocks, num_external_tokens
+        )
+
+    def build_connector_meta(
+        self,
+        scheduler_output: SchedulerOutput,
+    ) -> KVConnectorMetadata:
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.build_connector_meta(scheduler_output)
+
+    def request_finished(
+        self,
+        request: Request,
+        block_ids: list[int],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        return self.request_finished_all_groups(request, (block_ids,))
+
+    def request_finished_all_groups(
+        self,
+        request: Request,
+        block_ids: tuple[list[int], ...],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.request_finished(request, block_ids)
+
+    def reset_cache(self) -> bool | None:
+        """Reset the external Mooncake store on prefix-cache reset.
+
+        Drains the worker send queue, then runs ``remove_all`` on the
+        Mooncake master. Caller must first pause generation (e.g.
+        ``pause_generation``) so no new puts are enqueued during drain.
+
+        Returns True on ack, False on failure, None for the worker role.
+        """
+        if self.role == KVConnectorRole.SCHEDULER:
+            assert self.connector_scheduler is not None
+            # Clear local references to keys we're about to wipe.
+            self.connector_scheduler.load_specs.clear()
+            self._kv_cache_events = None
+            return self.connector_scheduler.reset_store()
+        return None
+
+    def update_connector_output(self, connector_output: KVConnectorOutput):
+        kv_cache_events = connector_output.kv_cache_events
+        if not kv_cache_events or not isinstance(
+            kv_cache_events, MooncakeStoreKVEvents
+        ):
+            return
+
+        if self._kv_cache_events is None:
+            self._kv_cache_events = kv_cache_events
+        else:
+            self._kv_cache_events.add_events(kv_cache_events.get_all_events())
+            self._kv_cache_events.increment_workers(
+                kv_cache_events.get_number_of_workers()
+            )
+
+    def take_events(self) -> Iterable[KVCacheEvent]:
+        if self._kv_cache_events is not None:
+            self._kv_cache_events.aggregate()
+            yield from self._kv_cache_events.get_all_events()
+            self._kv_cache_events.clear_events()
+            self._kv_cache_events = None
+
+    # ============================================================
+    # Worker-side methods
+    # ============================================================
+
+    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
+        assert self.connector_worker is not None
+        self.connector_worker.register_kv_caches(kv_caches)
+
+    def register_cross_layers_kv_cache(
+        self, kv_cache: torch.Tensor, attn_backend: type
+    ):
+        assert self.connector_worker is not None
+        assert (
+            self._kv_cache_config is not None
+            and len(self._kv_cache_config.kv_cache_groups) == 1
+        ), "Cross-layer KV cache does not supported with hybrid models"
+        self.connector_worker.register_cross_layers_kv_caches(kv_cache)
+
+    def start_load_kv(self, forward_context: ForwardContext, **kwargs: Any) -> None:
+        # No-op: loads are issued in get_finished() for compute overlap.
+        pass
+
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        # No layerwise support - no-op
+        return
+
+    def save_kv_layer(
+        self,
+        layer_name: str,
+        kv_layer: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        **kwargs: Any,
+    ) -> None:
+        # No layerwise support - no-op
+        return
+
+    def wait_for_save(self):
+        # No-op: stores are issued in get_finished() for compute overlap.
+        pass
+
+    def get_finished(
+        self, finished_req_ids: set[str]
+    ) -> tuple[set[str] | None, set[str] | None]:
+        assert self.connector_worker is not None
+        metadata = self._get_connector_metadata()
+        assert isinstance(metadata, MooncakeStoreConnectorMetadata)
+        return self.connector_worker.get_finished(finished_req_ids, metadata)
+
+    def get_block_ids_with_load_errors(self) -> set[int]:
+        assert self.connector_worker is not None
+        return self.connector_worker.get_block_ids_with_load_errors()
+
+    def get_kv_connector_kv_cache_events(
+        self,
+    ) -> MooncakeStoreKVEvents | None:
+        assert self.connector_worker is not None
+        events = self.connector_worker.get_kv_events()
+        if not events:
+            return None
+
+        kv_events = MooncakeStoreKVEvents(num_workers=1)
+        kv_events.add_events(events)
+        return kv_events
+
+    def get_kv_connector_stats(self) -> KVConnectorStats | None:
+        if self.connector_worker is None:
+            return None
+        return self.connector_worker.get_kv_connector_stats()
+
+    @classmethod
+    def build_kv_connector_stats(
+        cls, data: dict[str, Any] | None = None
+    ) -> KVConnectorStats | None:
+        return (
+            MooncakeStoreConnectorStats(data=data)
+            if data is not None
+            else MooncakeStoreConnectorStats()
+        )
+
+    @classmethod
+    def build_prom_metrics(
+        cls,
+        vllm_config: VllmConfig,
+        metric_types: dict[type[PromMetric], type[PromMetricT]],
+        labelnames: list[str],
+        per_engine_labelvalues: dict[int, list[object]],
+    ) -> KVConnectorPromMetrics:
+        return MooncakeStorePromMetrics(
+            vllm_config, metric_types, labelnames, per_engine_labelvalues
+        )

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py

@@ -0,0 +1,364 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""External-store cache-hit coordinator for MooncakeStoreConnector."""
+
+from collections.abc import Sequence
+from typing import cast
+
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
+    chunk_hashes_for_block_size,
+)
+from vllm.utils.math_utils import cdiv
+from vllm.v1.core.block_pool import BlockPool
+from vllm.v1.core.kv_cache_utils import (
+    BlockHash,
+    KVCacheBlock,
+)
+from vllm.v1.core.single_type_kv_cache_manager import (
+    SingleTypeKVCacheManager,
+)
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheGroupSpec,
+    KVCacheSpec,
+    UniformTypeKVCacheSpecs,
+)
+from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
+
+
+class ExternalCachedBlockPool:
+    """Duck-typed BlockPool backed by a ``(group_id, hash)`` exists set."""
+
+    def __init__(
+        self,
+        hash_block_size: int,
+        exists: set[tuple[int, bytes]] | None = None,
+    ) -> None:
+        # ``exists=None`` is used on the recv side where hit_length is already
+        # determined and we just want each spec's manager to apply its own mask.
+        self._exists = exists
+        self.hash_block_size = hash_block_size
+        self.null_block = KVCacheBlock(block_id=0)
+        # Dummy ID 1 for present block for duck-typing.
+        self._present_block = KVCacheBlock(block_id=1)
+
+    def get_cached_block(
+        self,
+        block_hash: BlockHash,
+        group_ids: list[int],
+    ) -> list[KVCacheBlock] | None:
+        # Mirrors BlockPool.get_cached_block: hit only when every group_id
+        # (groups sharing a spec) has the hash cached.
+        if self._exists is None:
+            return [self._present_block] * len(group_ids)
+        h = bytes(block_hash)
+        if all((g, h) in self._exists for g in group_ids):
+            return [self._present_block] * len(group_ids)
+        return None
+
+
+class MooncakeStoreCoordinator:
+    """Mirror of ``HybridKVCacheCoordinator.find_longest_cache_hit`` over an
+    ``ExternalCachedBlockPool``."""
+
+    def __init__(
+        self,
+        kv_cache_groups: list[KVCacheGroupSpec],
+        scheduler_block_size: int,
+        hash_block_size: int,
+        use_eagle: bool = False,
+        retention_interval: int | None = None,
+    ) -> None:
+        assert all(
+            g.kv_cache_spec.block_size % hash_block_size == 0 for g in kv_cache_groups
+        ), "block_size must be divisible by hash_block_size"
+        assert scheduler_block_size % hash_block_size == 0, (
+            f"scheduler_block_size ({scheduler_block_size}) must be a multiple of "
+            f"hash_block_size ({hash_block_size})"
+        )
+        assert all(
+            scheduler_block_size % g.kv_cache_spec.block_size == 0
+            for g in kv_cache_groups
+        ), "scheduler_block_size must be a multiple of each group's block_size"
+        self.kv_cache_groups = kv_cache_groups
+        self.hash_block_size = hash_block_size
+        self.lcm_block_size = scheduler_block_size
+        self.use_eagle = use_eagle
+        # Mirror vLLM core's KVCacheCoordinator.retention_interval.
+        self.retention_interval = retention_interval
+        self.eagle_group_ids = {
+            i for i, g in enumerate(kv_cache_groups) if g.is_eagle_group
+        }
+        if use_eagle and not self.eagle_group_ids:
+            self.eagle_group_ids = set(range(len(kv_cache_groups)))
+        self._verify_and_split_kv_cache_groups()
+
+    def _verify_and_split_kv_cache_groups(self) -> None:
+        """Mirrors KVCacheCoordinator.verify_and_split_kv_cache_groups but
+        dispatches via spec_manager_map (we don't allocate managers).
+        """
+        attention_groups: list[
+            tuple[KVCacheSpec, list[int], type[SingleTypeKVCacheManager]]
+        ] = []
+        for i, g in enumerate(self.kv_cache_groups):
+            spec = _unwrap_spec(g.kv_cache_spec)
+            manager_cls = KVCacheSpecRegistry.get_manager_class(spec)
+            assert manager_cls is not None, (
+                f"No manager registered for KVCacheSpec {spec}"
+            )
+            for existing_spec, group_ids, existing_cls in attention_groups:
+                if existing_spec == spec:
+                    assert manager_cls is existing_cls
+                    group_ids.append(i)
+                    break
+            else:
+                attention_groups.append((spec, [i], manager_cls))
+        # Full attention first (matches upstream convergence ordering).
+        self.attention_groups = sorted(
+            attention_groups,
+            key=lambda x: not isinstance(x[0], FullAttentionSpec),
+        )
+        self.eagle_attn_group_indices: set[int] = {
+            i
+            for i, (_, group_ids, _) in enumerate(self.attention_groups)
+            if any(self.kv_cache_groups[gid].is_eagle_group for gid in group_ids)
+        }
+        if self.use_eagle and not self.eagle_attn_group_indices:
+            self.eagle_attn_group_indices = set(range(len(self.attention_groups)))
+
+    def find_longest_cache_hit(
+        self,
+        block_hashes: Sequence[BlockHash],
+        max_length: int,
+        cached_block_pool: ExternalCachedBlockPool,
+        *,
+        apply_eagle: bool = True,
+    ) -> tuple[tuple[list[bool], ...], int]:
+        """Returns ``(load_mask_per_group, hit_length)``. ``mask[g][i]`` is True iff
+        group ``g`` populates chunk ``i`` locally (e.g. SWA and Mamba tail-only);
+        recv-side callers skip False slots.
+
+        ``apply_eagle`` controls whether the per-spec ``use_eagle`` last-block
+        pop is applied. Lookup callers want it (the drafter requires recomputing
+        the last block); per-chunk mask callers must not, because ``token_len``
+        already reflects the eagle-pruned hit length and a second pop would
+        leave the trailing block unloaded.
+        """
+        blocks_per_group, hit_length = self._find_hit_blocks(
+            block_hashes, max_length, cached_block_pool, apply_eagle=apply_eagle
+        )
+        masks = tuple(
+            [blk is not cached_block_pool.null_block for blk in blocks]
+            for blocks in blocks_per_group
+        )
+        return masks, hit_length
+
+    def load_mask(
+        self,
+        block_hashes: Sequence[BlockHash],
+        token_len: int,
+    ) -> tuple[list[bool], ...]:
+        """Per-group load masks: ``mask[g][i]`` is True iff group ``g``'s
+        spec would populate chunk ``i`` locally at length ``token_len``
+        (e.g. SWA / Mamba tail-only).
+        """
+        # ``apply_eagle=False`` because ``token_len`` is already the
+        # eagle-pruned hit length returned by ``client.lookup``. Re-applying
+        # the pop here would shorten the mask by one extra block; the recv
+        # thread would then silently skip the trailing chunk yielded by
+        # ``db.process_tokens`` and leave that block uninitialized in the
+        # local KV pool.
+        masks, _ = self.find_longest_cache_hit(
+            block_hashes,
+            token_len,
+            ExternalCachedBlockPool(self.hash_block_size),
+            apply_eagle=False,
+        )
+        return masks
+
+    def store_mask(
+        self,
+        aligned_token_len: int,
+        start_token: int = 0,
+        num_prompt_tokens: int | None = None,
+    ) -> tuple[list[bool] | None, ...]:
+        """Per-group store masks for the suffix starting at ``start_token``.
+
+        ``mask[g][i]`` is True iff the i-th chunk of group ``g`` *after*
+        ``start_token`` should be written to the store so a future cache hit
+        can consume it. ``None`` is the all-True sentinel for the suffix.
+
+        Reuses the engine's ``SingleTypeKVCacheManager.reachable_block_mask``
+        so the store retains exactly the blocks the local prefix cache would.
+        """
+        return self._reachable_masks(
+            aligned_token_len,
+            start_token,
+            retention_interval=self.retention_interval,
+            num_prompt_tokens=num_prompt_tokens,
+        )
+
+    def lookup_mask(
+        self,
+        aligned_token_len: int,
+    ) -> tuple[list[bool] | None, ...]:
+        """Per-group lookup masks.
+
+        ``mask[g][i]`` is True iff chunk ``i`` of group ``g`` should be
+        looked up as an aligned hit boundary. ``None`` is the all-True
+        sentinel.
+        """
+        return self._reachable_masks(
+            aligned_token_len,
+            0,
+            retention_interval=None,
+            num_prompt_tokens=None,
+        )
+
+    def _reachable_masks(
+        self,
+        aligned_token_len: int,
+        start_token: int,
+        *,
+        retention_interval: int | None,
+        num_prompt_tokens: int | None,
+    ) -> tuple[list[bool] | None, ...]:
+        assert aligned_token_len % self.lcm_block_size == 0, (
+            f"aligned_token_len ({aligned_token_len}) must be a multiple of "
+            f"lcm_block_size ({self.lcm_block_size})"
+        )
+        masks: list[list[bool] | None] = []
+        for g_idx, g in enumerate(self.kv_cache_groups):
+            spec = _unwrap_spec(g.kv_cache_spec)
+            end_chunk = aligned_token_len // spec.block_size
+            start_chunk = min(end_chunk, max(0, cdiv(start_token, spec.block_size)))
+            manager_cls = KVCacheSpecRegistry.get_manager_class(spec)
+            assert manager_cls is not None
+            use_eagle = g_idx in self.eagle_group_ids
+            mask = manager_cls.reachable_block_mask(
+                start_block=start_chunk,
+                end_block=end_chunk,
+                alignment_tokens=self.lcm_block_size,
+                kv_cache_spec=spec,
+                use_eagle=use_eagle,
+                retention_interval=retention_interval,
+                num_prompt_tokens=num_prompt_tokens,
+            )
+            if mask is not None:
+                assert len(mask) == end_chunk - start_chunk
+            masks.append(mask)
+        return tuple(masks)
+
+    def block_hashes_for_spec(
+        self, block_hashes: Sequence[BlockHash], spec: KVCacheSpec
+    ) -> Sequence[BlockHash]:
+        return chunk_hashes_for_block_size(
+            block_hashes, self.hash_block_size, spec.block_size
+        )
+
+    def _find_hit_blocks(
+        self,
+        block_hashes: Sequence[BlockHash],
+        max_length: int,
+        cached_block_pool: ExternalCachedBlockPool,
+        *,
+        apply_eagle: bool = True,
+    ) -> tuple[tuple[list[KVCacheBlock], ...], int]:
+        """Mirrors HybridKVCacheCoordinator.find_longest_cache_hit but
+        dispatches via spec_manager_map (we don't allocate managers).
+
+        When ``apply_eagle`` is False, ignore ``eagle_attn_group_indices`` —
+        used by ``load_mask`` to avoid popping a second block on top of the
+        one already removed by the lookup.
+        """
+        eagle_indices = self.eagle_attn_group_indices if apply_eagle else set()
+        if len(self.attention_groups) == 1:
+            spec, group_ids, manager_cls = self.attention_groups[0]
+            hashes = self.block_hashes_for_spec(block_hashes, spec)
+            hit_blocks, hit_length = manager_cls.find_longest_cache_hit(
+                block_hashes=hashes,  # type: ignore[arg-type]
+                max_length=max_length,
+                kv_cache_group_ids=group_ids,
+                block_pool=cast(BlockPool, cached_block_pool),
+                kv_cache_spec=spec,
+                drop_eagle_block=(0 in eagle_indices),
+                alignment_tokens=spec.block_size,
+            )
+            num_groups = len(self.kv_cache_groups)
+            blocks_by_group: list[list[KVCacheBlock]] = [[] for _ in range(num_groups)]
+            for gid, blks in zip(group_ids, hit_blocks, strict=True):
+                blocks_by_group[gid] = blks
+            return tuple(blocks_by_group), hit_length
+
+        num_groups = len(self.kv_cache_groups)
+        hit_length = max_length
+        hit_blocks_by_group: list[list[KVCacheBlock] | None] = [None] * num_groups
+        hit_length_by_group: list[int] = [0] * num_groups
+
+        is_simple_hybrid = len(self.attention_groups) == 2 and isinstance(
+            self.attention_groups[0][0], FullAttentionSpec
+        )
+        eagle_verified: set[int] = set()
+
+        while True:
+            curr_hit_length = hit_length
+
+            for idx, (spec, group_ids, manager_cls) in enumerate(self.attention_groups):
+                first_group_id = group_ids[0]
+                cached = hit_blocks_by_group[first_group_id]
+                if isinstance(spec, FullAttentionSpec) and cached is not None:
+                    curr_hit_length = min(
+                        curr_hit_length, hit_length_by_group[first_group_id]
+                    )
+                    continue
+
+                drop_eagle_block = idx in eagle_indices and idx not in eagle_verified
+                _max_length = curr_hit_length
+                if drop_eagle_block:
+                    _max_length = min(curr_hit_length + spec.block_size, max_length)
+                hashes = self.block_hashes_for_spec(block_hashes, spec)
+                hit_blocks, _new_hit_length = manager_cls.find_longest_cache_hit(
+                    block_hashes=hashes,  # type: ignore[arg-type]
+                    max_length=_max_length,
+                    kv_cache_group_ids=group_ids,
+                    block_pool=cast(BlockPool, cached_block_pool),
+                    kv_cache_spec=spec,
+                    drop_eagle_block=drop_eagle_block,
+                    alignment_tokens=self.lcm_block_size,
+                )
+                if drop_eagle_block:
+                    eagle_verified.add(idx)
+                elif _new_hit_length < curr_hit_length:
+                    eagle_verified.clear()
+                curr_hit_length = _new_hit_length
+                for gid, blocks in zip(group_ids, hit_blocks, strict=True):
+                    hit_blocks_by_group[gid] = blocks
+                    hit_length_by_group[gid] = _new_hit_length
+
+            if curr_hit_length >= hit_length:
+                break
+            hit_length = curr_hit_length
+            if is_simple_hybrid:
+                break
+
+        # Truncate full-attention hit_blocks to final converged length;
+        # other specs already trim themselves inside their hit logic.
+        spec0, group_ids0, _ = self.attention_groups[0]
+        if isinstance(spec0, FullAttentionSpec):
+            num_blocks = hit_length // spec0.block_size
+            for gid in group_ids0:
+                full_blks = hit_blocks_by_group[gid]
+                assert full_blks is not None
+                del full_blks[num_blocks:]
+                hit_length_by_group[gid] = hit_length
+
+        return (
+            tuple(blks if blks is not None else [] for blks in hit_blocks_by_group),
+            hit_length,
+        )
+
+
+def _unwrap_spec(spec: KVCacheSpec) -> KVCacheSpec:
+    if isinstance(spec, UniformTypeKVCacheSpecs):
+        return next(iter(spec.kv_cache_specs.values()))
+    return spec

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py

@@ -0,0 +1,406 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Adapted from vllm-project/vllm-ascend
+# (vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/).
+"""Data classes for MooncakeStoreConnector."""
+
+from collections.abc import Iterable, Sequence
+from dataclasses import dataclass
+from typing import cast
+
+import torch
+
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorMetadata,
+)
+from vllm.logger import init_logger
+from vllm.utils.math_utils import cdiv
+from vllm.v1.core.kv_cache_utils import (
+    BlockHash,
+    BlockHashListWithBlockSize,
+)
+
+logger = init_logger(__name__)
+
+
+class BlobBlockHashes(Sequence[BlockHash]):
+    """Lazy view over a flat buffer of fixed-size block hashes to avoid the overhead
+    of materializing all hashes upfront.
+    """
+
+    def __init__(self, blob: memoryview, hash_len: int):
+        self._blob = blob
+        self._hash_len = hash_len
+        self._n = len(blob) // hash_len if hash_len else 0
+
+    def __len__(self) -> int:
+        return self._n
+
+    def __getitem__(self, idx):
+        if isinstance(idx, slice):
+            return [self[i] for i in range(*idx.indices(self._n))]
+        if idx < 0:
+            idx += self._n
+        if not 0 <= idx < self._n:
+            raise IndexError(idx)
+        off = idx * self._hash_len
+        return BlockHash(self._blob[off : off + self._hash_len])
+
+
+class _CompactChunkHashList(BlockHashListWithBlockSize):
+    """View that keys each ``block_size`` chunk by the last constituent
+    ``hash_block_size`` hash instead of concatenating all of them.
+
+    The engine chains block hashes (each hash folds in the previous one), so the
+    final sub-block hash of a chunk already uniquely identifies the whole chunk
+    and its prefix. Using it keeps a Mooncake key at a single hash digest
+    regardless of the ``block_size`` / ``hash_block_size`` ratio, instead of
+    growing the key linearly with it (e.g. 64x for ``block_size=256``,
+    ``hash_block_size=4``).
+    """
+
+    def __init__(
+        self,
+        block_hashes: Sequence[BlockHash],
+        hash_block_size: int,
+        target_block_size: int,
+    ):
+        # Accept any indexable sequence (e.g. the lazy ``BlobBlockHashes``), not
+        # just ``list``; the base only indexes/sizes it.
+        assert target_block_size % hash_block_size == 0
+        self.block_hashes = block_hashes  # type: ignore[assignment]
+        self.scale_factor = target_block_size // hash_block_size
+
+    def _get_value_at(self, idx: int) -> BlockHash:
+        return self.block_hashes[idx * self.scale_factor + self.scale_factor - 1]
+
+
+def chunk_hashes_for_block_size(
+    block_hashes: Sequence[BlockHash],
+    hash_block_size: int,
+    block_size: int,
+) -> Sequence[BlockHash]:
+    """Map ``hash_block_size``-granular block hashes to one compact hash per
+    ``block_size`` chunk (the chunk's last sub-hash). Returns ``block_hashes``
+    unchanged when the two sizes are equal.
+    """
+    if block_size == hash_block_size:
+        return block_hashes
+    # Structurally a Sequence[BlockHash] (indexable + sized); the base class
+    # just isn't declared as one.
+    return cast(
+        "Sequence[BlockHash]",
+        _CompactChunkHashList(block_hashes, hash_block_size, block_size),
+    )
+
+
+@dataclass
+class KeyMetadata:
+    """Metadata for constructing pool keys."""
+
+    model_name: str
+    tp_rank: int
+    pcp_rank: int
+    dcp_rank: int
+    pp_rank: int
+    group_id: int = 0
+    # Optional namespace prepended to every key. Lets separate deployments
+    # share one Mooncake master without colliding on identical block hashes.
+    # Empty (the default) keeps keys byte-identical to the unprefixed format.
+    cache_prefix: str = ""
+
+
+@dataclass(order=True)
+class PoolKey:
+    """Key for addressing KV cache blocks in the distributed store."""
+
+    key_metadata: KeyMetadata
+    chunk_hash: str
+
+    def __hash__(self):
+        return hash(
+            (
+                self.key_metadata.cache_prefix,
+                self.key_metadata.model_name,
+                self.key_metadata.tp_rank,
+                self.key_metadata.pcp_rank,
+                self.key_metadata.dcp_rank,
+                self.key_metadata.pp_rank,
+                self.key_metadata.group_id,
+                self.chunk_hash,
+            )
+        )
+
+    @staticmethod
+    def build_prefix(
+        key_metadata: KeyMetadata,
+        *,
+        tp_rank: int | None = None,
+        pcp_rank: int | None = None,
+        dcp_rank: int | None = None,
+        pp_rank: int | None = None,
+    ) -> str:
+        """Return the stable prefix for a Mooncake pool key."""
+        prefix = f"{key_metadata.cache_prefix}@" if key_metadata.cache_prefix else ""
+        return (
+            f"{prefix}"
+            f"{key_metadata.model_name}"
+            f"@tp_rank:{key_metadata.tp_rank if tp_rank is None else tp_rank}"
+            f"@pcp{key_metadata.pcp_rank if pcp_rank is None else pcp_rank}"
+            f"@dcp{key_metadata.dcp_rank if dcp_rank is None else dcp_rank}"
+            f"@pp_rank:{key_metadata.pp_rank if pp_rank is None else pp_rank}"
+            f"@group:{key_metadata.group_id}"
+        )
+
+    @staticmethod
+    def build_key_string(key_prefix: str, chunk_hash: str) -> str:
+        return f"{key_prefix}@{chunk_hash}"
+
+    def to_string(self) -> str:
+        return self.build_key_string(
+            self.build_prefix(self.key_metadata), self.chunk_hash
+        )
+
+
+class ChunkedTokenDatabase:
+    """Maps token positions to store keys and GPU memory addresses."""
+
+    def __init__(
+        self,
+        metadata: KeyMetadata,
+        block_size: int,
+        hash_block_size: int | None = None,
+    ):
+        self.metadata = metadata
+        self.block_size = block_size
+        self.hash_block_size = hash_block_size or block_size
+        if self.block_size % self.hash_block_size != 0:
+            raise ValueError(
+                f"block_size ({self.block_size}) must be a multiple of "
+                f"hash_block_size ({self.hash_block_size})"
+            )
+        self.kv_caches_base_addr: list[int] = []
+        self.block_len: list[int] = []
+        self._key_prefix = PoolKey.build_prefix(metadata)
+
+    def key_for(self, chunk_hash: BlockHash) -> str:
+        return PoolKey.build_key_string(self._key_prefix, chunk_hash.hex())
+
+    def set_kv_caches_base_addr(self, kv_caches_base_addr: list[int]):
+        self.kv_caches_base_addr = kv_caches_base_addr
+
+    def set_block_len(self, block_len: list[int]):
+        self.block_len = block_len
+
+    def prepare_value(
+        self, start: int, end: int, block_ids: list[int]
+    ) -> tuple[list[int], list[int], int]:
+        """Compute memory addresses and sizes for a token range.
+
+        Returns:
+            (addr_list, size_list, block_id)
+        """
+        addr_list = []
+        size_list = []
+        block_id = block_ids[start // self.block_size]
+        length = len(self.block_len)
+        for index, base_addr in enumerate(self.kv_caches_base_addr):
+            addr = base_addr + block_id * self.block_len[index % length]
+            assert (end - start) % self.block_size == 0
+            size = self.block_len[index % length] * cdiv(end - start, self.block_size)
+            addr_list.append(addr)
+            size_list.append(size)
+        return addr_list, size_list, block_id
+
+    def process_tokens(
+        self,
+        token_len: int,
+        block_hashes: list[BlockHash],
+        mask_num: int = 0,
+        *,
+        chunk_mask: list[bool] | None = None,
+        put_step: int = 1,
+        put_step_rank: int = 0,
+    ) -> Iterable[tuple[int, int, BlockHash]]:
+        """Process tokens and yield (start_idx, end_idx, block_hash) tuples.
+
+        When there are fewer KV heads than TP ranks, chunks are distributed
+        across TP ranks to avoid duplicate load/store. The assignment keys off
+        the absolute ``chunk_id`` so a given chunk always lands on the same
+        rank regardless of where the processed suffix begins.
+
+        Args:
+            token_len: Total number of tokens.
+            block_hashes: Block hashes computed at ``hash_block_size`` granularity.
+                When ``block_size > hash_block_size`` each group's ``block_size`` chunk
+                is keyed by its last sub-hash via ``chunk_hashes_for_block_size``.
+            mask_num: Number of tokens to skip from the beginning.
+            chunk_mask: Optional mask relative to the first chunk after
+                ``mask_num``. False entries are skipped before hash access.
+            put_step: Stride for distributing chunks across ranks.
+            put_step_rank: ``chunk_id % put_step`` value this rank stores.
+        """
+        assert put_step > 0
+        if not block_hashes:
+            return
+        chunk_hashes: Sequence[BlockHash] = chunk_hashes_for_block_size(
+            block_hashes, self.hash_block_size, self.block_size
+        )
+        start_chunk = max(0, cdiv(mask_num, self.block_size))
+        max_chunks = min(len(chunk_hashes), cdiv(token_len, self.block_size))
+        if chunk_mask is not None:
+            max_chunks = min(max_chunks, start_chunk + len(chunk_mask))
+        for chunk_id in range(start_chunk, max_chunks):
+            if chunk_mask is not None and not chunk_mask[chunk_id - start_chunk]:
+                continue
+            if chunk_id % put_step != put_step_rank:
+                continue
+            h = chunk_hashes[chunk_id]
+            start_idx = chunk_id * self.block_size
+            end_idx = min(start_idx + self.block_size, token_len)
+            yield start_idx, end_idx, h
+
+
+@dataclass
+class LoadSpec:
+    """Specification for loading KV cache from external store."""
+
+    vllm_cached_tokens: int
+    kvpool_cached_tokens: int
+    can_load: bool
+    token_len: int = 0
+
+
+@dataclass
+class RequestTracker:
+    """Tracks per-request state across scheduler ticks."""
+
+    req_id: str
+    token_len: int
+    allocated_block_ids: tuple[list[int], ...]
+    num_saved_tokens: int = 0
+    token_ids: list[int] | None = None
+    # Snapshot of the prefill range length at tracker creation time.
+    # For a fresh request this is len(prompt). For a resumed-from-preemption
+    # request it includes previously-generated tokens, which are re-prefilled.
+    prefill_end_tokens: int = 0
+
+    def reset(self) -> None:
+        self.token_len = 0
+        self.allocated_block_ids = ()
+        self.num_saved_tokens = 0
+        self.token_ids = None
+        self.prefill_end_tokens = 0
+
+    def update(
+        self,
+        new_block_ids: tuple[list[int], ...] | list[int],
+    ) -> None:
+        # Backward-compat: accept a single list (broadcast to single group).
+        if isinstance(new_block_ids, list):
+            new_block_ids = (new_block_ids,)
+        if len(new_block_ids) != len(self.allocated_block_ids):
+            raise ValueError(
+                f"Group count mismatch: tracker has "
+                f"{len(self.allocated_block_ids)} groups, update has "
+                f"{len(new_block_ids)}"
+            )
+        for existing, new in zip(self.allocated_block_ids, new_block_ids, strict=True):
+            if new:
+                existing.extend(new)
+
+
+@dataclass
+class ReqMeta:
+    """Per-request metadata for store put/get operations."""
+
+    req_id: str
+    token_len_chunk: int
+    block_ids: tuple[list[int], ...]
+    block_hashes: list[BlockHash]
+
+    can_save: bool | None = None
+    load_spec: LoadSpec | None = None
+    is_last_chunk: bool | None = None
+    current_event: torch.cuda.Event | None = None
+
+    token_ids: list[int] | None = None
+    num_prompt_tokens: int | None = None
+
+    @staticmethod
+    def from_request_tracker(
+        tracker: RequestTracker,
+        block_size: int,
+        load_spec: LoadSpec | None = None,
+        skip_save: bool | None = False,
+        block_hashes: list[BlockHash] | None = None,
+        is_last_chunk: bool | None = None,
+    ) -> "ReqMeta | None":
+        """Create ReqMeta from a RequestTracker."""
+        if block_hashes is None:
+            block_hashes = []
+        input_token_len = tracker.token_len
+
+        chunk_boundary = cdiv(tracker.num_saved_tokens + 1, block_size) * block_size
+        num_tokens_to_save = input_token_len // block_size * block_size
+
+        skip_save = skip_save or num_tokens_to_save < chunk_boundary
+        # A ReqMeta must never carry both a save AND a load.
+        # The save would also be wasted work — the bytes are being looked up
+        # in the store right now. Later cached_reqs steps save new tokens
+        # normally.
+        if load_spec is not None and load_spec.can_load:
+            skip_save = True
+        if skip_save and load_spec is None:
+            return None
+
+        if not skip_save:
+            tracker.num_saved_tokens = num_tokens_to_save
+
+        token_ids = None
+        if tracker.token_ids:
+            token_ids = tracker.token_ids
+
+        if load_spec is not None and load_spec.can_load:
+            logger.debug(
+                "Scheduled to load %d tokens for request %s",
+                load_spec.kvpool_cached_tokens,
+                tracker.req_id,
+            )
+        else:
+            load_spec = None
+
+        logger.debug(
+            "request:%s, meta save spec:%s, meta load spec:%s",
+            tracker.req_id,
+            not skip_save,
+            load_spec,
+        )
+        return ReqMeta(
+            req_id=tracker.req_id,
+            token_len_chunk=num_tokens_to_save,
+            block_ids=tracker.allocated_block_ids,
+            can_save=not skip_save,
+            load_spec=load_spec,
+            block_hashes=block_hashes,
+            is_last_chunk=is_last_chunk,
+            token_ids=token_ids,
+            num_prompt_tokens=tracker.prefill_end_tokens,
+        )
+
+
+class MooncakeStoreConnectorMetadata(KVConnectorMetadata):
+    """Metadata passed from scheduler to worker."""
+
+    def __init__(
+        self,
+        unfinished_request_ids: set[str],
+        preempted_req_ids: set[str],
+    ):
+        self.requests: list[ReqMeta] = []
+        self.unfinished_request_ids = unfinished_request_ids
+        self.preempted_req_ids = preempted_req_ids
+
+    def add_request(self, req_meta: ReqMeta) -> None:
+        self.requests.append(req_meta)

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/metrics.py

@@ -0,0 +1,189 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Per-operation telemetry for MooncakeStoreConnector.
+
+Records one row per Mooncake RPC (``save_exists``, ``save_put``, ``load_get``,
+``lookup_exists``) with duration, key/byte counts, status, and failed-key
+count. Exposed to the logger via ``KVConnectorLogging`` and to Prometheus
+via ``MooncakeStorePromMetrics``.
+"""
+
+from dataclasses import dataclass
+from statistics import fmean
+from typing import Any
+
+from vllm.config import VllmConfig
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
+    KVConnectorPromMetrics,
+    KVConnectorStats,
+    PromMetric,
+    PromMetricT,
+)
+
+
+def _nearest_rank_percentile(values: list[float], percentile: float) -> float:
+    if not values:
+        return 0.0
+    sorted_values = sorted(values)
+    rank = max(
+        0, min(len(sorted_values) - 1, int(percentile * len(sorted_values) - 1e-12))
+    )
+    return sorted_values[rank]
+
+
+@dataclass
+class MooncakeStoreConnectorStats(KVConnectorStats):
+    """Serializable Mooncake store communication telemetry."""
+
+    def __post_init__(self):
+        if not self.data:
+            self.reset()
+
+    def reset(self):
+        self.data: dict[str, list[dict[str, int | float | str]]] = {}
+
+    def is_empty(self) -> bool:
+        return not self.data
+
+    def aggregate(self, other: KVConnectorStats) -> KVConnectorStats:
+        if other.is_empty():
+            return self
+        for operation, records in other.data.items():
+            self.data.setdefault(operation, []).extend(records)
+        return self
+
+    def reduce(self) -> dict[str, int | float]:
+        reduced: dict[str, int | float] = {}
+        for operation, records in sorted(self.data.items()):
+            if not records:
+                continue
+            durations = [float(record["duration_seconds"]) for record in records]
+            reduced[f"{operation}_count"] = len(records)
+            reduced[f"{operation}_avg_ms"] = round(fmean(durations) * 1e3, 3)
+            reduced[f"{operation}_p90_ms"] = round(
+                _nearest_rank_percentile(durations, 0.9) * 1e3, 3
+            )
+            reduced[f"{operation}_total_keys"] = sum(
+                int(record["num_keys"]) for record in records
+            )
+            reduced[f"{operation}_total_bytes"] = sum(
+                int(record["num_bytes"]) for record in records
+            )
+            reduced[f"{operation}_failed_keys"] = sum(
+                int(record["num_failed_keys"]) for record in records
+            )
+            reduced[f"{operation}_error_count"] = sum(
+                1 for record in records if record["status"] == "error"
+            )
+        return reduced
+
+    def record_operation(
+        self,
+        operation: str,
+        duration_seconds: float,
+        num_keys: int,
+        *,
+        num_bytes: int = 0,
+        status: str = "ok",
+        num_failed_keys: int = 0,
+    ) -> None:
+        self.data.setdefault(operation, []).append(
+            {
+                "duration_seconds": duration_seconds,
+                "num_keys": num_keys,
+                "num_bytes": num_bytes,
+                "status": status,
+                "num_failed_keys": num_failed_keys,
+            }
+        )
+
+
+class MooncakeStorePromMetrics(KVConnectorPromMetrics):
+    """Prometheus metrics for Mooncake store communication."""
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        metric_types: dict[type[PromMetric], type[PromMetricT]],
+        labelnames: list[str],
+        per_engine_labelvalues: dict[int, list[object]],
+    ):
+        super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues)
+        metric_labelnames = labelnames + ["operation", "status"]
+        self._metric_cache: dict[tuple[int, str, str], dict[str, PromMetric]] = {}
+
+        self._histogram_operation_time = self._histogram_cls(
+            name="vllm:mooncake_store_operation_time_seconds",
+            documentation="Histogram of Mooncake store communication time.",
+            buckets=[
+                1e-3,
+                5e-3,
+                1e-2,
+                5e-2,
+                1e-1,
+                2e-1,
+                3e-1,
+                4e-1,
+                5e-1,
+                7.5e-1,
+                1.0,
+                1.5,
+                2.0,
+                3.0,
+                4.0,
+            ],
+            labelnames=metric_labelnames,
+        )
+        self._counter_operation_calls = self._counter_cls(
+            name="vllm:mooncake_store_operation_total",
+            documentation="Number of Mooncake store communication operations.",
+            labelnames=metric_labelnames,
+        )
+        self._counter_operation_keys = self._counter_cls(
+            name="vllm:mooncake_store_operation_keys_total",
+            documentation="Number of Mooncake store keys touched by operations.",
+            labelnames=metric_labelnames,
+        )
+        self._counter_operation_bytes = self._counter_cls(
+            name="vllm:mooncake_store_operation_bytes_total",
+            documentation="Number of bytes transferred by Mooncake store operations.",
+            labelnames=metric_labelnames,
+        )
+        self._counter_failed_keys = self._counter_cls(
+            name="vllm:mooncake_store_operation_failed_keys_total",
+            documentation="Number of Mooncake store keys that failed in operations.",
+            labelnames=metric_labelnames,
+        )
+
+    def _get_metrics(
+        self,
+        engine_idx: int,
+        operation: str,
+        status: str,
+    ) -> dict[str, PromMetric]:
+        cache_key = (engine_idx, operation, status)
+        if cache_key not in self._metric_cache:
+            label_values = self.per_engine_labelvalues[engine_idx] + [operation, status]
+            self._metric_cache[cache_key] = {
+                "time": self._histogram_operation_time.labels(*label_values),
+                "calls": self._counter_operation_calls.labels(*label_values),
+                "keys": self._counter_operation_keys.labels(*label_values),
+                "bytes": self._counter_operation_bytes.labels(*label_values),
+                "failed_keys": self._counter_failed_keys.labels(*label_values),
+            }
+        return self._metric_cache[cache_key]
+
+    def observe(self, transfer_stats_data: dict[str, Any] | None, engine_idx: int = 0):
+        if not transfer_stats_data:
+            return
+        for operation, records in transfer_stats_data.items():
+            assert isinstance(records, list)
+            for record in records:
+                assert isinstance(record, dict)
+                status = str(record["status"])
+                metrics = self._get_metrics(engine_idx, operation, status)
+                metrics["time"].observe(float(record["duration_seconds"]))
+                metrics["calls"].inc()
+                metrics["keys"].inc(int(record["num_keys"]))
+                metrics["bytes"].inc(int(record["num_bytes"]))
+                metrics["failed_keys"].inc(int(record["num_failed_keys"]))

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/protocol.py

@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Wire-format constants for the LookupKey ZMQ admin channel.
+
+This is the single source of truth shared by ``LookupKeyClient`` and
+``LookupKeyServer`` on the scheduler<->worker rank-0 admin channel.
+
+Wire format (REQ/REP over IPC):
+
+    Request: [msg_type: bytes] [payload_frames...]
+
+      msg_type == LOOKUP_MSG:
+          frame 1: token_len (u32 big-endian, 4 bytes)
+          frame 2: hash_len (u16 big-endian, 2 bytes) — byte length of each
+                   fixed-size block hash (0 when there are no hashes)
+          frame 3: raw block hashes concatenated back-to-back (each hash_len
+                   bytes); the server splits on hash_len
+        Response: [hit_count: u32 big-endian, 4 bytes]
+
+      msg_type == RESET_MSG:
+          (no payload frames)
+        Response: [RESP_OK] or [RESP_ERR]
+
+The first frame of every request is a named bytes tag (not a numeric
+sentinel that aliases the data field) so the protocol stays
+self-describing and extensible: adding new admin commands requires
+only a new tag and a new dispatch branch.
+
+Mirrors the named-tag convention used by the NIXL connector (see
+``vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py``).
+"""
+
+# Request message-type tags. Frame 0 of every request.
+LOOKUP_MSG: bytes = b"lookup"
+RESET_MSG: bytes = b"reset"
+
+# Single-byte response status codes for admin commands.
+RESP_OK: bytes = b"\x01"
+RESP_ERR: bytes = b"\x00"

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py

@@ -0,0 +1,402 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Adapted from vllm-project/vllm-ascend
+# (vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/).
+"""Scheduler-side logic for MooncakeStoreConnector."""
+
+from typing import Any
+
+from vllm.config import VllmConfig
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorMetadata,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (  # noqa: E501
+    LoadSpec,
+    MooncakeStoreConnectorMetadata,
+    ReqMeta,
+    RequestTracker,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.worker import (  # noqa: E501
+    LookupKeyClient,
+)
+from vllm.logger import init_logger
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes
+from vllm.v1.core.sched.output import NewRequestData, SchedulerOutput
+from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.request import Request
+
+logger = init_logger(__name__)
+
+
+def _new_req_prefill_tokens(request: NewRequestData) -> list[int]:
+    """Tokens this prefill will compute KV for.
+
+    Under the v2 model runner, resumed-from-preemption requests appear in
+    ``scheduled_new_reqs`` with ``prefill_token_ids`` set to the request's full
+    token list (prompt + previously-generated). For all other cases this falls
+    back to the original prompt.
+    """
+    if request.prefill_token_ids is not None:
+        return request.prefill_token_ids
+    assert request.prompt_token_ids is not None
+    return request.prompt_token_ids
+
+
+class MooncakeStoreScheduler:
+    """Scheduler-side component for MooncakeStoreConnector."""
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        kv_cache_config: KVCacheConfig,
+    ):
+        assert vllm_config.kv_transfer_config is not None
+        self.kv_role = vllm_config.kv_transfer_config.kv_role
+        kvc_extra_config = vllm_config.kv_transfer_config.kv_connector_extra_config
+        self.load_async = kvc_extra_config.get("load_async", True)
+        self.lookup_async = kvc_extra_config.get("lookup_async", False)
+        self.client = LookupKeyClient(vllm_config)
+
+        # Align with the engine's own scheduler_block_size and hash_block_size.
+        self._block_size, self._hash_block_size = resolve_kv_cache_block_sizes(
+            kv_cache_config, vllm_config
+        )
+
+        # Per-request state
+        self.load_specs: dict[str, LoadSpec] = {}  # to be loaded
+        self._request_trackers: dict[str, RequestTracker] = {}  # scheduled new requests
+        self._unfinished_requests: dict[str, tuple[Request, tuple[list[int], ...]]] = {}
+        self._unfinished_request_ids: set[str] = set()
+
+    def get_num_new_matched_tokens(
+        self,
+        request: Request,
+        num_computed_tokens: int,
+    ) -> tuple[int | None, bool]:
+        """Check for external KV cache hit.
+
+        Returns ``(None, False)`` when an async lookup is still in flight,
+        signaling the scheduler to retry this request on a later step.
+        """
+        # Look up against the full prefill range, not just the prompt.
+        token_len = request.num_tokens // self._block_size * self._block_size
+        if token_len < self._block_size:
+            return 0, False
+
+        num_external_hit_tokens = self.client.lookup(
+            request.request_id,
+            token_len,
+            request.block_hashes,
+            non_block=self.lookup_async,
+        )
+        if num_external_hit_tokens is None:
+            # Lookup not ready yet; scheduler will retry on a later step.
+            return None, False
+
+        if num_external_hit_tokens == request.num_tokens:
+            # Leave a sub-block tail uncomputed for sampling, on a block
+            # boundary so the recv-side load mask covers every yielded chunk.
+            num_external_hit_tokens = max(
+                0,
+                (request.num_tokens - 1) // self._block_size * self._block_size,
+            )
+
+        if num_external_hit_tokens < num_computed_tokens:
+            need_to_allocate = 0
+        else:
+            need_to_allocate = num_external_hit_tokens - num_computed_tokens
+
+        logger.debug(
+            "Reqid: %s, Total tokens %d, kvpool hit tokens: %d, need to load: %d",
+            request.request_id,
+            request.num_tokens,
+            num_external_hit_tokens,
+            need_to_allocate,
+        )
+
+        if need_to_allocate <= 0:
+            return 0, False
+
+        self.load_specs[request.request_id] = LoadSpec(
+            vllm_cached_tokens=num_computed_tokens,
+            kvpool_cached_tokens=num_external_hit_tokens,
+            can_load=False,
+        )
+
+        return need_to_allocate, self.load_async
+
+    def update_state_after_alloc(
+        self,
+        request: Request,
+        blocks: KVCacheBlocks,
+        num_external_tokens: int,
+    ):
+        """Update state after block allocation."""
+        local_block_ids: tuple[list[int], ...] = ()
+        if num_external_tokens > 0:
+            local_block_ids = blocks.get_block_ids()
+
+        self._unfinished_requests[request.request_id] = (request, local_block_ids)
+        self._unfinished_request_ids.add(request.request_id)
+
+        if request.request_id not in self.load_specs:
+            return
+
+        if num_external_tokens == 0:
+            self.load_specs[request.request_id].can_load = False
+            return
+
+        assert (
+            num_external_tokens > 0
+            and num_external_tokens
+            == self.load_specs[request.request_id].kvpool_cached_tokens
+            - self.load_specs[request.request_id].vllm_cached_tokens
+        ), (
+            f"Mismatch in number of tokens: {num_external_tokens} vs "
+            f"{self.load_specs[request.request_id].kvpool_cached_tokens} - "
+            f"{self.load_specs[request.request_id].vllm_cached_tokens}"
+            f" for request {request.request_id}"
+        )
+
+        self.load_specs[request.request_id].can_load = True
+
+    def build_connector_meta(
+        self, scheduler_output: SchedulerOutput
+    ) -> KVConnectorMetadata:
+        """Build connector metadata for this scheduler step."""
+        force_skip_save = self.kv_role == "kv_consumer"
+
+        for finished_req_id in scheduler_output.finished_req_ids:
+            self.client.discard(finished_req_id)
+            self.load_specs.pop(finished_req_id, None)
+            self._request_trackers.pop(finished_req_id, None)
+            self._unfinished_requests.pop(finished_req_id, None)
+            self._unfinished_request_ids.discard(finished_req_id)
+
+        preempted_ids = scheduler_output.preempted_req_ids or set()
+        for req_id in preempted_ids:
+            self.load_specs.pop(req_id, None)
+            if request_tracker := self._request_trackers.get(req_id):
+                request_tracker.reset()
+            self._unfinished_requests.pop(req_id, None)
+
+        meta = MooncakeStoreConnectorMetadata(
+            self._unfinished_request_ids,
+            preempted_ids,
+        )
+
+        # Handle new requests
+        for request in scheduler_output.scheduled_new_reqs:
+            load_spec = self.load_specs.pop(request.req_id, None)
+            num_tokens_to_compute = (
+                request.num_computed_tokens
+                + scheduler_output.num_scheduled_tokens[request.req_id]
+            )
+            assert request.req_id in self._unfinished_requests
+            request_tuple = self._unfinished_requests.get(request.req_id)
+            request_real = request_tuple[0]  # type: ignore[index]
+
+            if isinstance(request.block_ids, tuple):
+                # Multi-group: preserve per-group structure.
+                unfolded_block_ids = tuple(b.copy() for b in request.block_ids)
+            else:
+                # Single-group legacy: list[int] -> 1-tuple.
+                unfolded_block_ids = (request.block_ids.copy(),)
+
+            prefill_tokens = _new_req_prefill_tokens(request)
+            request_tracker = RequestTracker(
+                req_id=request.req_id,
+                token_len=num_tokens_to_compute,
+                allocated_block_ids=unfolded_block_ids,
+                num_saved_tokens=0,
+                token_ids=prefill_tokens[:num_tokens_to_compute],
+                prefill_end_tokens=len(prefill_tokens),
+            )
+            self._request_trackers[request.req_id] = request_tracker
+
+            last_chunk_tokens_num = (
+                len(prefill_tokens) // self._block_size * self._block_size
+            )
+
+            req_meta = ReqMeta.from_request_tracker(
+                request_tracker,
+                self._block_size,
+                load_spec=load_spec,
+                skip_save=force_skip_save,
+                block_hashes=request_real.block_hashes,
+                is_last_chunk=(request_tracker.token_len >= last_chunk_tokens_num),
+            )
+            if req_meta is not None:
+                meta.add_request(req_meta)
+
+        # Handle cached (running, or MRV1 resumed-from-preemption) requests
+        cached_reqs = scheduler_output.scheduled_cached_reqs
+        if not force_skip_save:
+            for i, req_id in enumerate(cached_reqs.req_ids):
+                new_block_ids = cached_reqs.new_block_ids[i]
+                if not new_block_ids:
+                    continue
+
+                req_meta = None
+                if req_id in cached_reqs.resumed_req_ids:
+                    # Resumed after preemption
+                    if isinstance(new_block_ids, tuple):
+                        new_block_ids = tuple(b.copy() for b in new_block_ids)
+                    else:
+                        new_block_ids = (new_block_ids.copy(),)
+                    load_spec = self.load_specs.pop(req_id, None)
+                    request_tuple = self._unfinished_requests.get(req_id)
+                    request_real = request_tuple[0]  # type: ignore[index]
+                    num_tokens_to_compute = (
+                        request_real.num_computed_tokens
+                        + scheduler_output.num_scheduled_tokens[req_id]
+                    )
+                    # On resume, the request re-prefills prompt + previously
+                    # generated tokens (all_token_ids).
+                    prefill_tokens = list(request_real.all_token_ids)
+                    request_tracker = RequestTracker(
+                        req_id=req_id,
+                        token_len=num_tokens_to_compute,
+                        allocated_block_ids=new_block_ids,
+                        num_saved_tokens=0,
+                        token_ids=prefill_tokens[:num_tokens_to_compute].copy(),
+                        prefill_end_tokens=len(prefill_tokens),
+                    )
+                    self._request_trackers[req_id] = request_tracker
+
+                    last_chunk_tokens_num = (
+                        len(prefill_tokens) // self._block_size * self._block_size
+                    )
+                    req_meta = ReqMeta.from_request_tracker(
+                        request_tracker,
+                        self._block_size,
+                        load_spec=load_spec,
+                        skip_save=force_skip_save,
+                        block_hashes=request_real.block_hashes,
+                        is_last_chunk=(
+                            request_tracker.token_len >= last_chunk_tokens_num
+                        ),
+                    )
+                else:
+                    # Decode/chunked request
+                    request_tracker = self._request_trackers[req_id]
+                    num_new_tokens = scheduler_output.num_scheduled_tokens[req_id]
+                    req_tuple = self._unfinished_requests.get(req_id)
+                    if req_tuple:
+                        unfinished_req = req_tuple[0]
+                        num_current_tokens = request_tracker.token_len
+                        new_token_ids = unfinished_req.all_token_ids[
+                            num_current_tokens : num_current_tokens + num_new_tokens
+                        ]
+                        request_tracker.token_len += len(new_token_ids)
+                    else:
+                        raise ValueError(
+                            f"Request {req_id} is not in _unfinished_requests"
+                        )
+                    num_computed_token = cached_reqs.num_computed_tokens[i]
+                    # Use the tracker's snapshot of the prefill range so resumed
+                    # requests keep saving past the original prompt boundary.
+                    prefill_end = request_tracker.prefill_end_tokens
+                    if num_computed_token >= prefill_end:
+                        continue
+                    request_tracker.update(new_block_ids)
+
+                    last_chunk_tokens_num = (
+                        prefill_end // self._block_size * self._block_size
+                    )
+                    req_meta = ReqMeta.from_request_tracker(
+                        request_tracker,
+                        self._block_size,
+                        load_spec=None,
+                        skip_save=force_skip_save,
+                        block_hashes=unfinished_req.block_hashes,
+                        is_last_chunk=(
+                            request_tracker.token_len >= last_chunk_tokens_num
+                        ),
+                    )
+
+                if req_meta is not None:
+                    meta.add_request(req_meta)
+
+        # Handle requests with pending load specs not yet scheduled
+        request_ids = [req.req_id for req in scheduler_output.scheduled_new_reqs]
+        for request_id, (
+            unfinished_req,
+            block_ids,
+        ) in self._unfinished_requests.items():
+            if request_id not in request_ids and request_id not in cached_reqs.req_ids:
+                load_spec = self.load_specs.pop(request_id, None)
+                if not load_spec:
+                    continue
+                num_tokens_to_compute = load_spec.kvpool_cached_tokens
+                request_tracker = RequestTracker(
+                    req_id=request_id,
+                    token_len=num_tokens_to_compute,
+                    allocated_block_ids=block_ids,
+                    num_saved_tokens=0,
+                )
+                self._request_trackers[request_id] = request_tracker
+                req_meta = ReqMeta.from_request_tracker(
+                    request_tracker,
+                    self._block_size,
+                    load_spec=load_spec,
+                    skip_save=None,
+                    block_hashes=unfinished_req.block_hashes,
+                )
+                if req_meta is not None:
+                    meta.add_request(req_meta)
+
+        return meta
+
+    def request_finished(
+        self,
+        request: Request,
+        block_ids: tuple[list[int], ...],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        """Determine whether to delay freeing blocks for async save."""
+        if self.kv_role == "kv_consumer":
+            return False, None
+        tracker = self._request_trackers.get(request.request_id)
+        # Missing tracker can happen when the request is aborted before the
+        # connector observes the normal finished lifecycle or is preempted
+        # before finishing.
+        if tracker is None or tracker.num_saved_tokens <= 0:
+            return False, None
+        total_blocks = sum(len(g) for g in block_ids)
+        delay_free_blocks = total_blocks > 0
+        if delay_free_blocks:
+            logger.debug(
+                "Delaying free of %d blocks for request %s",
+                total_blocks,
+                request.request_id,
+            )
+        return delay_free_blocks, None
+
+    def reset_store(self) -> bool:
+        """Trigger a global ``remove_all(force=True)`` on the Mooncake master.
+
+        Routes through the existing LookupKey ZMQ admin channel to worker
+        rank 0, which owns the ``MooncakeDistributedStore`` handle.
+
+        Ordering assumption: caller (typically
+        ``Scheduler.reset_connector_cache``, invoked via
+        ``reset_prefix_cache(reset_connector=True)``) MUST ensure no
+        in-flight Mooncake lookups or transfers. For RL workflows this is
+        satisfied at the step boundary after weight updates and rollout
+        drain. Violating this can allow stale KV to be served on the next
+        request, defeating the hard-reset guarantee.
+
+        Returns True on ACK from worker, False on NACK or RPC error.
+        """
+        try:
+            ok = self.client.reset()
+            if ok:
+                logger.info("Mooncake store reset via remove_all succeeded.")
+            else:
+                logger.warning("Mooncake store reset returned NACK from worker.")
+            return ok
+        except Exception as e:
+            logger.error("Mooncake reset_store RPC failed: %s", e)
+            return False

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py

@@ -0,0 +1,500 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import contextlib
+import os
+import threading
+import time
+from collections.abc import Iterator
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, NamedTuple
+
+import msgspec
+import regex as re
+import torch
+import zmq
+
+from vllm.config import KVTransferConfig, VllmConfig
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorMetadata,
+)
+from vllm.distributed.parallel_state import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+)
+from vllm.logger import init_logger
+from vllm.utils.network_utils import (
+    get_ip,
+    get_open_port,
+    make_zmq_socket,
+)
+
+if TYPE_CHECKING:
+    pass
+
+from dataclasses import field
+from enum import Enum
+
+logger = init_logger(__name__)
+
+
+Transfer = tuple[int, float]
+EngineId = str
+ReqId = str
+TransferId = str
+
+
+class MoRIIOTransferAck(NamedTuple):
+    transfer_id: TransferId
+    consumer_tp_size: int = 1
+
+
+@dataclass
+class WriteTask:
+    request_id: ReqId
+    transfer_id: TransferId
+    dst_engine_id: str
+    local_block_ids: list[int]
+    remote_block_ids_hint: list[int] | None
+    layer_name: str
+    event: torch.cuda.Event
+    remote_notify_port: int
+    remote_ip: str
+    enqueue_time: float = field(default_factory=time.perf_counter)
+    retried: int = 0
+
+
+@dataclass
+class LayerTransferPlan:
+    """Plan for transferring a single layer."""
+
+    request_id: ReqId
+    transfer_id: TransferId
+    layer_name: str
+    sess_idx: int
+    transfer_local_offsets: list[int]
+    transfer_remote_offsets: list[int]
+    transfer_sizes: list[int]
+    use_batch: bool = True
+
+
+@dataclass
+class RemoteAllocInfo:
+    """Information about remote block allocation."""
+
+    block_ids: list[int]
+    writes_done: int = 0
+    writes_expected: int | None = None
+    decode_dp_rank: int = 0
+    completion_request_id: str | None = None
+    completion_remote_notify_port: int | None = None
+    completion_remote_ip: str | None = None
+    completion_notified: bool = False
+    transfer_statuses: list[Any] = field(default_factory=list)
+    transfer_offsets: dict[
+        tuple[tuple[int, ...], tuple[int, ...], torch.dtype],
+        tuple[list[int], list[int], list[int]],
+    ] = field(default_factory=dict)
+
+
+class ROLE(Enum):
+    PRODUCER = "producer"
+    CONSUMER = "consumer"
+    NOTINIT = "notinit"
+
+
+class MoRIIOAgentMetadata(
+    msgspec.Struct,
+    omit_defaults=True,  # type: ignore[call-arg]
+    # required for @cached_property.d
+    dict=True,
+):
+    engine_id: str
+    agent_metadata: bytes
+    kv_caches_base_addr: list[int]
+    num_blocks: int
+    block_len: int
+    attn_backend_name: str
+
+
+class RoleManager:
+    """Manages role state across the connector."""
+
+    _instance: "RoleManager | None" = None
+    _lock = threading.Lock()
+
+    def __init__(self) -> None:
+        self._role: ROLE = ROLE.NOTINIT
+
+    @classmethod
+    def get_instance(cls) -> "RoleManager":
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = cls()
+        return cls._instance
+
+    def set_role(self, role: ROLE) -> None:
+        """Set the current role."""
+        with self._lock:
+            self._role = role
+
+    def get_role(self) -> ROLE:
+        """Get the current role."""
+        return self._role
+
+
+def set_role(role: ROLE):
+    """Set the global role."""
+    RoleManager.get_instance().set_role(role)
+
+
+def get_role() -> ROLE:
+    """Get the global role."""
+    return RoleManager.get_instance().get_role()
+
+
+class MoRIIOMode(Enum):
+    READ = "read"
+    WRITE = "write"
+
+
+class MoRIIOError(Exception):
+    """Base exception for MoRIIO operations."""
+
+    pass
+
+
+class HandshakeError(MoRIIOError):
+    """Exception raised when handshake fails."""
+
+    pass
+
+
+class TransferError(MoRIIOError):
+    """Exception raised when transfer fails."""
+
+    pass
+
+
+def get_moriio_mode(kv_transfer_config: KVTransferConfig) -> MoRIIOMode:
+    read_mode = str(
+        kv_transfer_config.kv_connector_extra_config.get("read_mode", "false")
+    ).lower().strip() in ("true", "1")
+    logger.debug("MoRIIO Connector read_mode: %s", read_mode)
+    if read_mode:
+        return MoRIIOMode.READ
+    else:
+        return MoRIIOMode.WRITE
+
+
+def get_port_offset(dp_rank: int, tp_rank: int, tp_size: int = 1) -> int:
+    return (dp_rank) * tp_size + tp_rank
+
+
+def resolve_host_ip(extra_config: dict) -> str:
+    """The IP this MoRIIO process advertises for KV transfer.
+
+    Honors an explicit ``host_ip`` in ``kv_connector_extra_config`` before
+    falling back to ``get_ip()``. An external router/orchestrator can set it to
+    the node's routable address; this is required under frameworks (e.g. Ray)
+    where ``get_ip()`` resolves to an unroutable public IP and ``VLLM_HOST_IP``
+    cannot be propagated to the worker processes that bind the transfer engine.
+    """
+    return extra_config.get("host_ip") or get_ip()
+
+
+_DEPRECATED_ENV_VARS: dict[str, str] = {
+    "VLLM_MORIIO_CONNECTOR_READ_MODE": "read_mode",
+    "VLLM_MORIIO_QP_PER_TRANSFER": "qp_per_transfer",
+    "VLLM_MORIIO_POST_BATCH_SIZE": "post_batch_size",
+    "VLLM_MORIIO_NUM_WORKERS": "num_workers",
+}
+
+
+def _warn_deprecated_env_vars() -> None:
+    for env_var, new_key in _DEPRECATED_ENV_VARS.items():
+        if env_var in os.environ:
+            logger.warning_once(
+                "The environment variable %s is deprecated and ignored. "
+                "Set %r inside kv_transfer_config.kv_connector_extra_config "
+                "instead.",
+                env_var,
+                new_key,
+            )
+
+
+@dataclass
+class MoRIIOConfig:
+    local_ip: str
+    local_kv_port: int
+    proxy_ip: str
+    local_ping_port: int
+    proxy_ping_port: int
+    http_port: int
+    handshake_port: int
+    notify_port: int
+    tp_rank: int
+    dp_rank: int
+    dp_size: int
+    tp_size: int
+    transfer_timeout: float
+    defer_timeout: float
+    read_mode: bool = False
+    qp_per_transfer: int = 1
+    post_batch_size: int = -1
+    num_workers: int = 1
+    backend: str = "rdma"
+
+    @classmethod
+    def from_vllm_config(cls, vllm_config: VllmConfig) -> "MoRIIOConfig":
+        # Port Configuration:
+        # local_ping_port   -> Outgoing heartbeat to proxy
+        # proxy_ping_port   -> Remote proxy's heartbeat ingress port
+        # http_port         -> Instance's HTTP service endpoint
+        # local_kv_port     -> service port for mori engine
+        # notify_port       -> For synchronizing stages between prefill and decode
+        # handshake_port    -> For initial handshake between mori engine
+
+        # Optional tuning knobs
+        # read_mode        -> If true, run the connector in READ mode (consumer
+        #                     pulls KV from producer) instead of the default
+        #                     WRITE mode.
+        # transfer_timeout -> Timeout for waiting_for_transfer_complete before
+        #                     raising TransferError (sec).
+        # defer_timeout    -> Timeout before a deferred send with no finished_sending
+        #                     notification is reaped and its blocks force-freed (sec).
+
+        # Knobs for RDMA transfers, ignored if on xgmi backend
+        # qp_per_transfer  -> Number of RDMA Queue Pairs per KV transfer.
+        # post_batch_size  -> Batch size for posting transfer work requests
+        #                     (-1 lets the MoRI backend choose).
+        # num_workers      -> Number of background worker threads the MoRI
+        #                     engine uses for transfer processing.
+
+        # TODO : merge notify_port and handshake_port to simplify port management
+        #        supports non-contiguous ports
+        assert vllm_config.kv_transfer_config is not None, (
+            "kv_transfer_config must be set for MoRIIOConnector"
+        )
+        _warn_deprecated_env_vars()
+        kv_transfer_config = vllm_config.kv_transfer_config
+        extra_config = kv_transfer_config.kv_connector_extra_config
+        tp_rank = get_tensor_model_parallel_rank()
+        dp_rank = vllm_config.parallel_config.data_parallel_rank
+        base_notify_port = int(extra_config["notify_port"])
+        dp_size = vllm_config.parallel_config.data_parallel_size
+        tp_size = get_tensor_model_parallel_world_size()
+        port_offset = get_port_offset(dp_rank, tp_rank)
+        backend = str(extra_config.get("backend", "rdma")).lower()
+        if backend not in ("rdma", "xgmi"):
+            raise ValueError(
+                f"Invalid MoRIIO backend {backend!r} in kv_connector_extra_config; "
+                "must be one of 'rdma' or 'xgmi'."
+            )
+
+        transfer_timeout = float(
+            extra_config.get(
+                "transfer_timeout", MoRIIOConstants.DEFAULT_TRANSFER_TIMEOUT
+            )
+        )
+        defer_timeout = float(
+            extra_config.get("defer_timeout", MoRIIOConstants.DEFAULT_DEFER_TIMEOUT)
+        )
+
+        return cls(
+            local_ip=resolve_host_ip(extra_config),
+            local_kv_port=get_open_port(),
+            proxy_ip=extra_config["proxy_ip"],
+            local_ping_port=get_open_port(),
+            proxy_ping_port=int(extra_config["proxy_ping_port"]),
+            http_port=int(extra_config["http_port"]),
+            handshake_port=int(extra_config["handshake_port"]),
+            notify_port=base_notify_port + port_offset,
+            tp_rank=tp_rank,
+            dp_rank=dp_rank,
+            dp_size=dp_size,
+            tp_size=tp_size,
+            read_mode=get_moriio_mode(kv_transfer_config) == MoRIIOMode.READ,
+            qp_per_transfer=int(extra_config.get("qp_per_transfer", 1)),
+            post_batch_size=int(extra_config.get("post_batch_size", -1)),
+            num_workers=int(extra_config.get("num_workers", 1)),
+            backend=backend,
+            transfer_timeout=transfer_timeout,
+            defer_timeout=defer_timeout,
+        )
+
+
+class MoRIIOConstants:
+    """Constants for MoRIIO connector."""
+
+    # ZMQ message types
+    GET_META_MSG = b"get_meta_msg"
+    POP_DONE_RECV = b"pop_done_recv"
+    OVER = b"OVER"
+    COMPLETION_PREFIX = "cmpl"
+    TRANSFER_PREFIX = "tx"
+
+    PING_INTERVAL = 3
+    MAX_PING_RETRIES = 100
+    DEFAULT_HANDSHAKE_PORT = "6301"
+    DEFAULT_NOTIFY_PORT = "61005"
+
+    VLLM_MORI_READ_ABORT_REQUEST_TIMEOUT = 3600
+
+    # Timeout (seconds) for waiting_for_transfer_complete before raising TransferError.
+    # Overridable via kv_connector_extra_config["transfer_timeout"].
+    DEFAULT_TRANSFER_TIMEOUT = 30.0
+    # Timeout (seconds) before a deferred send with no finished_sending
+    # notification is reaped and its blocks force-freed.
+    # Overridable via kv_connector_extra_config["defer_timeout"].
+    DEFAULT_DEFER_TIMEOUT = 60.0
+
+
+# The router embeds both zmq_addresses in the request_id:
+#   "___prefill_addr_{zmq}___decode_addr_{zmq}_{32-hex-uuid}"
+# MoRIIO zmq_address format: "host:IP,handshake:PORT,notify:PORT"
+#
+# This lets each connector side parse the peer's connection info without
+# requiring the router to pass it explicitly in kv_transfer_params.
+_PREFILL_ZMQ_RE = re.compile(r"___prefill_addr_(.+?)___decode_addr_")
+# vLLM wraps the router's X-Request-Id as "cmpl-<id>-<seq>-<hex>" so there may
+# be a trailing "-<seq>-<hex>" suffix after the 32-char UUID.  Allow it.
+_DECODE_ZMQ_RE = re.compile(r"___decode_addr_(.+)_[0-9a-f]{32}(?:-.*)?$")
+
+
+def parse_moriio_zmq_address(
+    zmq_address: str,
+) -> tuple[str, int, int]:
+    """Parse the MoRI-IO zmq address into its components.
+
+    Parses ``"host:IP,handshake:PORT,notify:PORT"`` into
+        (host, handshake_port, notify_port).
+
+    Each key-value pair is split on the *first* colon so that IPv6 addresses
+    (e.g. ``host:::1``) are handled correctly.  Raises ``ValueError`` if any
+    of ``host``, ``handshake``, or ``notify`` keys are absent or if the port
+    values are non-numeric.
+    """
+    parts: dict[str, str] = {}
+    for segment in zmq_address.split(","):
+        key, _, val = segment.partition(":")
+        parts[key.strip()] = val.strip()
+    try:
+        host = parts["host"]
+        handshake_port = int(parts["handshake"])
+        notify_port = int(parts["notify"])
+    except (KeyError, ValueError) as e:
+        raise ValueError(
+            f"Malformed zmq_address {zmq_address!r}: expected "
+            f"'host:IP,handshake:PORT,notify:PORT' format"
+        ) from e
+    return host, handshake_port, notify_port
+
+
+def get_peer_zmq_from_request_id(request_id: str, is_producer: bool) -> str:
+    """Extract the *peer's* zmq_address from the vLLM router request_id.
+
+    The producer (prefill) needs the decode's address; the consumer (decode)
+    needs the prefill's address.
+    """
+    if is_producer:
+        m = _DECODE_ZMQ_RE.search(request_id)
+    else:
+        m = _PREFILL_ZMQ_RE.search(request_id)
+    if m is None:
+        raise ValueError(
+            f"Cannot parse peer zmq_address from request_id: {request_id!r}"
+        )
+    return m.group(1)
+
+
+@dataclass
+class ReqMeta:
+    """Metadata for a single request."""
+
+    transfer_id: TransferId
+    local_block_ids: list[int]
+    remote_block_ids: list[int]
+    remote_host: str
+    remote_port: int
+    remote_handshake_port: int
+    remote_notify_port: int
+    remote_engine_id: str
+    tp_size: int
+    remote_dp_size: int
+
+
+class MoRIIOConnectorMetadata(KVConnectorMetadata):
+    def __init__(self):
+        self.reqs_to_recv: dict[ReqId, ReqMeta] = {}
+        self.reqs_to_save: dict[ReqId, ReqMeta] = {}
+        self.reqs_to_send: dict[ReqId, float] = {}
+        self.transfer_id_to_request_id: dict[TransferId, ReqId] = {}
+
+    def __repr__(self):
+        return (
+            f"MoRIIOConnectorMetadata: reqs_to_recv={self.reqs_to_recv}, "
+            f"reqs_to_save={self.reqs_to_save}, "
+            f"reqs_to_send={self.reqs_to_send}, "
+            f"transfer_id_to_request_id={self.transfer_id_to_request_id}"
+        )
+
+    def add_new_req(
+        self,
+        request_id: ReqId,
+        local_block_ids: list[int],
+        kv_transfer_params: dict[str, Any],
+        write_mode=False,
+    ):
+        transfer_id = kv_transfer_params["transfer_id"]
+
+        remote_host = kv_transfer_params.get("remote_host")
+        remote_handshake_port = kv_transfer_params.get("remote_handshake_port")
+        remote_notify_port = kv_transfer_params.get("remote_notify_port")
+        if (
+            remote_host is None
+            or remote_handshake_port is None
+            or remote_notify_port is None
+        ):
+            # Parse host/ports from the request_id. The router embeds both
+            # zmq_addresses in PD request IDs, but WRITE decode requests may carry
+            # a plain request ID and get the remote address via kv_transfer_params.
+            peer_zmq = get_peer_zmq_from_request_id(request_id, is_producer=write_mode)
+            remote_host, remote_handshake_port, remote_notify_port = (
+                parse_moriio_zmq_address(peer_zmq)
+            )
+
+        _req = ReqMeta(
+            transfer_id=transfer_id,
+            local_block_ids=local_block_ids,
+            remote_block_ids=kv_transfer_params["remote_block_ids"],
+            remote_engine_id=kv_transfer_params["remote_engine_id"],
+            remote_host=remote_host,
+            remote_port=int(remote_handshake_port),
+            remote_handshake_port=int(remote_handshake_port),
+            remote_notify_port=int(remote_notify_port),
+            tp_size=kv_transfer_params.get("tp_size", 1),
+            remote_dp_size=kv_transfer_params.get("remote_dp_size", 1),
+        )
+        if write_mode:
+            self.reqs_to_save[request_id] = _req
+        else:
+            self.reqs_to_recv[request_id] = _req
+
+
+@contextlib.contextmanager
+def zmq_ctx(socket_type: Any, addr: str) -> Iterator[zmq.Socket]:
+    """Context manager for a ZMQ socket"""
+
+    if socket_type not in (zmq.ROUTER, zmq.REQ, zmq.DEALER):
+        raise ValueError(f"Unexpected socket type: {socket_type}")
+
+    ctx: zmq.Context | None = None
+    try:
+        ctx = zmq.Context()  # type: ignore[attr-defined]
+        yield make_zmq_socket(
+            ctx=ctx, path=addr, socket_type=socket_type, bind=socket_type == zmq.ROUTER
+        )
+    finally:
+        if ctx is not None:
+            ctx.destroy(linger=0)

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py

@@ -0,0 +1,885 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import threading
+import time
+from collections import OrderedDict, defaultdict
+from queue import Empty, Queue
+from typing import TYPE_CHECKING, Any
+from weakref import ref as weakref_ref
+
+import msgpack
+import torch
+import zmq
+
+from vllm.logger import init_logger
+from vllm.utils.network_utils import (
+    make_zmq_path,
+    make_zmq_socket,
+)
+
+if TYPE_CHECKING:
+    from mori.io import BackendType
+
+from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
+    ROLE,
+    HandshakeError,
+    LayerTransferPlan,
+    MoRIIOAgentMetadata,
+    MoRIIOConstants,
+    MoRIIOError,
+    MoRIIOTransferAck,
+    RemoteAllocInfo,
+    TransferError,
+    TransferId,
+    WriteTask,
+    get_port_offset,
+    get_role,
+    zmq_ctx,
+)
+
+if TYPE_CHECKING:
+    from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import (
+        MoRIIOConnectorWorker,
+    )
+
+logger = init_logger(__name__)
+try:
+    from mori.io import (
+        BackendType,
+        EngineDesc,
+        IOEngine,
+        MemoryDesc,
+        PollCqMode,
+        RdmaBackendConfig,
+        XgmiBackendConfig,
+    )
+
+    logger.info("MoRIIO is available")
+except ImportError:
+    logger.error("MoRIIO is not available")
+
+
+"""Write task execution logic for MoRIIO connector."""
+
+
+_MAX_TERMINAL_TRANSFER_IDS = 4096
+
+
+WriteGeometryKey = tuple[tuple[int, ...], tuple[int, ...], torch.dtype]
+
+
+def _get_write_geometry_key(kv_cache: torch.Tensor) -> WriteGeometryKey:
+    return (tuple(kv_cache.shape), tuple(kv_cache.stride()), kv_cache.dtype)
+
+
+class MoRIIOWriter:
+    """Handles write operations for KV cache transfers.
+
+    WRITE mode state machine:
+    D sends destination block allocation, P schedules one write per layer
+    after the layer CUDA event, P seals the scheduled write count after
+    forward, then P notifies D and releases P blocks after all scheduled
+    writes complete.
+    """
+
+    def __init__(self, worker: "MoRIIOConnectorWorker"):
+        """Initialize the writer.
+
+        Args:
+            worker: Reference to the parent worker
+        """
+        self._worker_ref: weakref_ref[MoRIIOConnectorWorker] = weakref_ref(worker)
+        self._write_task_q: Queue[WriteTask] = Queue()
+        self._write_worker_started = False
+        self._write_worker_lock = threading.Lock()
+        self._write_state_lock = threading.Lock()
+        self._deferred_tasks: list[WriteTask] = []
+        self._scheduled_writes: dict[TransferId, int] = defaultdict(int)
+        self._scheduled_layers: dict[TransferId, set[str]] = defaultdict(set)
+        self._sealed_writes: dict[TransferId, int] = {}
+        self._defer_timeout = worker.moriio_config.defer_timeout
+
+    @property
+    def worker(self) -> "MoRIIOConnectorWorker":
+        """Get the worker instance.
+
+        Returns:
+            The parent worker instance
+
+        Raises:
+            RuntimeError: If worker has been garbage collected
+        """
+        worker = self._worker_ref()
+        if worker is None:
+            raise RuntimeError("Parent worker has been garbage collected")
+        return worker
+
+    def ensure_worker_started(self) -> None:
+        """Ensure the background write worker is running."""
+        if self._write_worker_started:
+            return
+        self._write_worker_started = True
+        with self._write_worker_lock:
+            thread = threading.Thread(
+                target=self._write_worker_loop, daemon=True, name="moriio-write-worker"
+            )
+            thread.start()
+            logger.info("Started MoRIIO write worker thread")
+
+    def schedule_write(self, task: WriteTask) -> bool:
+        """Schedule a write task.
+
+        Args:
+            task: The write task to schedule
+        """
+        self.ensure_worker_started()
+        if self._is_transfer_terminal(task.transfer_id):
+            return False
+
+        with self._write_state_lock:
+            if self._is_transfer_terminal(task.transfer_id):
+                return False
+            if task.layer_name in self._scheduled_layers[task.transfer_id]:
+                return False
+            self._scheduled_layers[task.transfer_id].add(task.layer_name)
+            self._scheduled_writes[task.transfer_id] += 1
+        self._write_task_q.put(task)
+        return True
+
+    def is_scheduled(self, transfer_id: TransferId, layer_name: str) -> bool:
+        with self._write_state_lock:
+            return layer_name in self._scheduled_layers.get(transfer_id, set())
+
+    def seal_pending_transfers(self) -> None:
+        """Seal expected WRITE counts after the model forward has run.
+
+        `save_kv_layer` is only invoked for attention layers whose backend uses
+        the standard KV connector hook. Hybrid models can register more KV
+        cache tensors than the number of hooks that fire in a forward, so WRITE
+        completion must be based on the tasks actually queued for the transfer.
+        """
+        pending: list[tuple[TransferId, RemoteAllocInfo]] = []
+        with self._write_state_lock:
+            for transfer_id, write_count in self._scheduled_writes.items():
+                if transfer_id in self._sealed_writes:
+                    continue
+                self._sealed_writes[transfer_id] = write_count
+                request_info = (
+                    self.worker.moriio_wrapper.done_remote_allocate_req_dict.get(
+                        transfer_id
+                    )
+                )
+                if request_info is not None:
+                    request_info.writes_expected = write_count
+                    pending.append((transfer_id, request_info))
+
+        for transfer_id, request_info in pending:
+            self._finalize_if_complete(transfer_id, request_info)
+
+    def _write_worker_loop(self) -> None:
+        """Main loop for the write worker thread."""
+
+        while True:
+            # Process deferred tasks first
+            self._process_deferred_tasks()
+
+            # Get new task
+            try:
+                task = self._write_task_q.get(timeout=0.01)
+            except Empty:
+                continue
+
+            if self._is_transfer_terminal(task.transfer_id):
+                continue
+
+            # Check if remote blocks are ready
+            if not self._is_remote_ready(task):
+                # task.retry_count += 1
+                self._deferred_tasks.append(task)
+                # logger.debug(
+                #     "Deferred task for request %s (retry %d)",
+                #     task.request_id, task.retry_count
+                # )
+                continue
+
+            # Execute the task
+            try:
+                self._execute_write_task(task)
+            except Exception:
+                logger.exception(
+                    "Write task failed for request %s, marking done",
+                    task.request_id,
+                )
+                self._mark_request_done(task.transfer_id)
+
+    def _process_deferred_tasks(self) -> None:
+        """Process tasks that were previously deferred."""
+        if not self._deferred_tasks:
+            return
+
+        defer_timeout = self._defer_timeout
+        now = time.perf_counter()
+        still_deferred: list[WriteTask] = []
+
+        for task in self._deferred_tasks:
+            if self._is_transfer_terminal(task.transfer_id):
+                continue
+            if now - task.enqueue_time > defer_timeout:
+                logger.error(
+                    "Deferred write task for request %s expired after %.1fs "
+                    "(remote blocks never arrived), marking done",
+                    task.request_id,
+                    now - task.enqueue_time,
+                )
+                self._mark_request_done(task.transfer_id)
+                continue
+            if self._is_remote_ready(task):
+                try:
+                    self._execute_write_task(task)
+                except Exception:
+                    logger.exception(
+                        "Deferred write task failed for request %s, marking done",
+                        task.request_id,
+                    )
+                    self._mark_request_done(task.transfer_id)
+            else:
+                still_deferred.append(task)
+
+        self._deferred_tasks = still_deferred
+
+    def _clear_transfer_state(self, transfer_id: TransferId) -> None:
+        with self._write_state_lock:
+            self._scheduled_writes.pop(transfer_id, None)
+            self._scheduled_layers.pop(transfer_id, None)
+            self._sealed_writes.pop(transfer_id, None)
+
+    def _is_transfer_terminal(self, transfer_id: TransferId) -> bool:
+        wrapper = self.worker.moriio_wrapper
+        with wrapper.lock:
+            return wrapper._is_transfer_terminal_locked(transfer_id)
+
+    def _mark_request_done(self, transfer_id: str) -> None:
+        """Mark a request done so its blocks are freed, even on transfer failure."""
+        wrapper = self.worker.moriio_wrapper
+        with wrapper.lock:
+            wrapper.done_req_ids.append(MoRIIOTransferAck(transfer_id))
+            wrapper.done_remote_allocate_req_dict.pop(transfer_id, None)
+            wrapper._mark_transfer_terminal_locked(transfer_id)
+        self._clear_transfer_state(transfer_id)
+
+    def _is_remote_ready(self, task: WriteTask) -> bool:
+        """Check if remote blocks are allocated for this task.
+
+        Args:
+            task: The write task
+
+        Returns:
+            True if remote blocks are ready
+        """
+        return (
+            task.transfer_id in self.worker.moriio_wrapper.done_remote_allocate_req_dict
+        )
+
+    def _get_remote_alloc_info(self, transfer_id: str) -> RemoteAllocInfo:
+        """Get remote allocation info for a request.
+
+        Args:
+            transfer_id:TransferId The request ID
+
+        Returns:
+            Remote allocation information
+
+        Raises:
+            KeyError: If allocation info is missing
+        """
+        try:
+            return self.worker.moriio_wrapper.done_remote_allocate_req_dict[transfer_id]
+        except KeyError as e:
+            raise KeyError(
+                f"Remote allocation info missing for transfer {transfer_id}"
+            ) from e
+
+    def _execute_write_task(self, task: WriteTask) -> None:
+        """Execute a single write task.
+
+        Args:
+            task: The write task to execute
+
+        """
+        # Get remote allocation info
+        request_info = self._get_remote_alloc_info(task.transfer_id)
+        with self._write_state_lock:
+            request_info.completion_request_id = task.request_id
+            request_info.completion_remote_notify_port = task.remote_notify_port
+            request_info.completion_remote_ip = task.remote_ip
+            if task.transfer_id in self._sealed_writes:
+                request_info.writes_expected = self._sealed_writes[task.transfer_id]
+
+        if request_info.block_ids is None:
+            logger.debug(
+                "Request remote block IDs not ready:request_id = %s, transfer_id = %s",
+                task.request_id,
+                task.transfer_id,
+            )
+            return
+
+        # Wait for CUDA event
+        # The attention computation of the current layer cannot
+        # overlap with the kv transfer task,
+        # otherwise it will cause precision issues.
+        # This event is used to synchronize the kv transfer and computation tasks.
+        task.event.synchronize()
+
+        # Update engine ID with DP rank
+        task.dst_engine_id = self.worker.get_engine_name_with_dp(
+            task.dst_engine_id, request_info.decode_dp_rank
+        )
+
+        # Get or create sessions
+        sessions, remote_moriio_meta = self.worker._get_built_session(
+            task.dst_engine_id
+        )
+
+        # Prepare transfer plan
+        plan = self._prepare_transfer_plan(task, request_info, remote_moriio_meta)
+
+        # Execute transfer
+        transfer_statuses = self._do_layer_write(plan, sessions)
+        with self._write_state_lock:
+            request_info.transfer_statuses.extend(transfer_statuses)
+
+        # Finalize if all layers complete
+        self._mark_write_done(task.transfer_id, request_info)
+
+    def _prepare_transfer_plan(
+        self,
+        task: WriteTask,
+        request_info: RemoteAllocInfo,
+        remote_moriio_meta: MoRIIOAgentMetadata,
+    ) -> LayerTransferPlan:
+        """Prepare the transfer plan for a layer.
+
+        Args:
+            task: The write task
+            request_info: Remote allocation information
+
+        Returns:
+            The transfer plan
+        """
+        layer_cache = self.worker.kv_caches[task.layer_name]
+        geometry_key = _get_write_geometry_key(layer_cache)
+        offsets = request_info.transfer_offsets.get(geometry_key)
+        if offsets is None:
+            offsets = self.worker._compute_block_transfer_offsets(
+                task.layer_name,
+                task.local_block_ids,
+                request_info.block_ids,
+                remote_moriio_meta,
+            )
+            request_info.transfer_offsets[geometry_key] = offsets
+
+        # Get session index
+        layer_names = list(self.worker.layer_name_to_local_kv_cache_metadata.keys())
+        sess_idx = layer_names.index(task.layer_name)
+
+        local_off, remote_off, sizes = offsets
+
+        return LayerTransferPlan(
+            request_id=task.request_id,
+            transfer_id=task.transfer_id,
+            layer_name=task.layer_name,
+            sess_idx=sess_idx,
+            transfer_local_offsets=local_off,
+            transfer_remote_offsets=remote_off,
+            transfer_sizes=sizes,
+            use_batch=True,
+        )
+
+    def _do_layer_write(self, plan: LayerTransferPlan, sessions: list) -> list[Any]:
+        """Perform the actual layer write.
+
+        Args:
+            plan: The transfer plan
+            sessions: List of transfer sessions
+        """
+        if plan.use_batch:
+            return [
+                self.worker.moriio_wrapper.write_remote_data(
+                    plan.transfer_sizes,
+                    plan.transfer_local_offsets,
+                    plan.transfer_remote_offsets,
+                    sessions[plan.sess_idx],
+                )
+            ]
+
+        transfer_statuses: list[Any] = []
+        for i in range(len(plan.transfer_local_offsets)):
+            transfer_statuses.append(
+                self.worker.moriio_wrapper.write_remote_data_single(
+                    plan.transfer_sizes[i],
+                    plan.transfer_local_offsets[i],
+                    plan.transfer_remote_offsets[i],
+                    plan.sess_idx,
+                )
+            )
+        return transfer_statuses
+
+    def _mark_write_done(
+        self, transfer_id: TransferId, request_info: RemoteAllocInfo
+    ) -> None:
+        """Record one completed WRITE task and finalize if sealed."""
+        with self._write_state_lock:
+            request_info.writes_done += 1
+        self._finalize_if_complete(transfer_id, request_info)
+
+    def _finalize_if_complete(
+        self, transfer_id: TransferId, request_info: RemoteAllocInfo
+    ) -> None:
+        """Finalize transfer if all scheduled writes are complete."""
+        with self._write_state_lock:
+            expected = request_info.writes_expected
+            if expected is None or request_info.writes_done < expected:
+                return
+            if request_info.completion_notified:
+                return
+            request_id = request_info.completion_request_id
+            remote_notify_port = request_info.completion_remote_notify_port
+            remote_ip = request_info.completion_remote_ip
+            if request_id is None or remote_notify_port is None or remote_ip is None:
+                return
+            transfer_statuses = list(request_info.transfer_statuses)
+            request_info.transfer_statuses.clear()
+            request_info.completion_notified = True
+
+        # Wait for this request's transfers to complete.
+        self.worker.moriio_wrapper.waiting_for_transfer_complete(transfer_statuses)
+
+        remote_port = remote_notify_port + get_port_offset(
+            request_info.decode_dp_rank, self.worker.tp_rank
+        )
+        # Consider using RDMA immediate data in decode side
+        # to eliminate the need for this notification.
+        # Consider including the first gen token from prefill in the notification
+
+        # Send completion notification
+        self.worker.moriio_wrapper.send_notify(
+            transfer_id, remote_ip, remote_port, message_type="write_done"
+        )
+        # mark request as done, then we can free the blocks
+        with self.worker.moriio_wrapper.lock:
+            self.worker.moriio_wrapper.done_req_ids.append(
+                MoRIIOTransferAck(transfer_id)
+            )
+            self.worker.moriio_wrapper.done_remote_allocate_req_dict.pop(
+                transfer_id, None
+            )
+            self.worker.moriio_wrapper._mark_transfer_terminal_locked(transfer_id)
+        self._clear_transfer_state(transfer_id)
+        logger.debug(
+            "Completed transfer for (request, transfer) %s, %s, notified port %d",
+            request_id,
+            transfer_id,
+            remote_port,
+        )
+
+
+class MoRIIOWrapper:
+    """Wrapper for MoRIIO engine operations.
+
+    Handles both producer and consumer roles for KV cache transfers.
+
+    Args:
+        moriio_engine:  MoRIIO engine instance
+        tp_rank: Tensor parallel rank
+        dp_rank: Data parallel rank
+    """
+
+    def __init__(
+        self,
+        moriio_engine: "IOEngine | None" = None,
+        tp_rank: int = 0,
+        dp_rank: int = 0,
+        transfer_timeout: float = MoRIIOConstants.DEFAULT_TRANSFER_TIMEOUT,
+    ):
+        self.tp_rank = tp_rank
+        self.dp_rank = dp_rank
+        self.moriio_engine = moriio_engine
+        self.remote_memory_metadata = None
+        self.local_memory_registered = False
+        self.local_memory_metadata = None
+        self.transfer_status: list[Any] = []
+        self.remote_engine_ip: str | None = None
+        self.notify_port: int | None = None
+        self.lock = threading.Lock()
+        self.done_req_ids: list[MoRIIOTransferAck] = []
+        self.done_remote_allocate_req_dict: dict[TransferId, RemoteAllocInfo] = {}
+        self.done_write_cache_req_ids: list[str] = []
+        self._terminal_transfer_ids: OrderedDict[TransferId, None] = OrderedDict()
+        self._transfer_timeout = transfer_timeout
+        self.notify_thread: threading.Thread | None = None
+        self.sessions: list[IOEngine.Session] = []
+        self.paths: dict[str, zmq.Socket] = {}
+
+    def set_moriio_engine(self, moriio_engine):
+        assert moriio_engine is not None, (
+            "You Cannot pass None engine to MoRIIOWrapper!"
+        )
+        self.moriio_engine = moriio_engine
+
+    def set_backend_type(
+        self,
+        backend_type: "BackendType",
+        qp_per_transfer: int = 1,
+        post_batch_size: int = -1,
+        num_workers: int = 1,
+    ) -> None:
+        assert self.moriio_engine is not None, "MoRIIO engine must be set first"
+        if backend_type == BackendType.XGMI:
+            logger.info("Using MoRIIO backend: XGMI")
+            self.moriio_engine.create_backend(backend_type, XgmiBackendConfig())
+        else:
+            logger.info(
+                "Using MoRIIO backend: RDMA "
+                "(qp_per_transfer=%d, post_batch_size=%d, num_workers=%d)",
+                qp_per_transfer,
+                post_batch_size,
+                num_workers,
+            )
+            rdma_cfg = RdmaBackendConfig(
+                qp_per_transfer,
+                post_batch_size,
+                num_workers,
+                PollCqMode.POLLING,
+                # vLLM uses ZMQ for completion signaling
+                # and never calls PopInboundTransferStatus.
+                # With notifications enabled, ibv_post_send
+                # ENOMEM at high concurrency permanently
+                # poisons TransferStatus before the data WR
+                # completes, hanging requests in
+                # WAITING_FOR_REMOTE_KVS indefinitely.
+                enable_notification=False,
+            )
+            self.moriio_engine.create_backend(backend_type, rdma_cfg)
+
+    def get_agent_metadata(self):
+        assert self.moriio_engine is not None, "MoRIIO engine must be set first"
+        engine_metadata = self.moriio_engine.get_engine_desc()
+        engine_metadata_packed = engine_metadata.pack()
+        return engine_metadata_packed
+
+    def register_remote_engine(self, remote_packed_engine_metadata):
+        assert self.moriio_engine is not None, "MoRIIO engine must be set first"
+        consumer_engine_metadata = EngineDesc.unpack(remote_packed_engine_metadata)
+        self.moriio_engine.register_remote_engine(consumer_engine_metadata)
+        return consumer_engine_metadata.key
+
+    def register_local_tensor(self, tensor: torch.Tensor):
+        assert self.moriio_engine is not None, "MoRIIO engine must be set first"
+        try:
+            self.local_memory_metadata = self.moriio_engine.register_torch_tensor(
+                tensor
+            )
+            assert self.local_memory_metadata is not None, (
+                "register_torch_tensor returned None"
+            )
+            local_memory_metadata_packed = self.local_memory_metadata.pack()
+        except Exception as e:
+            raise MoRIIOError(f"Failed to register local memory: {e}") from e
+        self.local_memory_registered = True
+        return local_memory_metadata_packed
+
+    def get_unpack_memory_metadata(self, packed_memory_metadata):
+        return MemoryDesc.unpack(packed_memory_metadata)
+
+    def build_session(self, local_memory_metadata, remote_memory_metadata):
+        assert self.moriio_engine is not None, "MoRIIO engine must be set first"
+        return self.moriio_engine.create_session(
+            local_memory_metadata, remote_memory_metadata
+        )
+
+    def read_remote_data(
+        self, transfer_size_byte, local_offset=0, remote_offset=0, session=None
+    ):
+        assert self.local_memory_registered, "You have not register local memory data!"
+        assert self.moriio_engine is not None, "MoRIIO engine must be set first"
+        transfer_status = session.batch_read(
+            local_offset,
+            remote_offset,
+            transfer_size_byte,
+            self.moriio_engine.allocate_transfer_uid(),
+        )
+
+        return transfer_status
+
+    def write_remote_data(
+        self, transfer_size_byte, local_offset=0, remote_offset=0, session=None
+    ):
+        assert self.local_memory_registered, "You have not register local memory data!"
+        assert self.moriio_engine is not None, "MoRIIO engine must be set first"
+        write_uid = self.moriio_engine.allocate_transfer_uid()
+
+        transfer_status = session.batch_write(
+            local_offset, remote_offset, transfer_size_byte, write_uid
+        )
+        return transfer_status
+
+    def write_remote_data_single(
+        self, transfer_size_byte, local_offset=0, remote_offset=0, sess_idx=0
+    ):
+        assert self.local_memory_registered, "You have not register local memory data!"
+        assert self.moriio_engine is not None, "MoRIIO engine must be set first"
+        transfer_status = self.sessions[sess_idx].write(
+            local_offset,
+            remote_offset,
+            transfer_size_byte,
+            self.moriio_engine.allocate_transfer_uid(),
+        )
+        return transfer_status
+
+    def waiting_for_transfer_complete(self, transfer_statuses: list[Any] | None = None):
+        if transfer_statuses is None:
+            with self.lock:
+                transfers_to_wait = self.transfer_status[:]
+                self.transfer_status.clear()
+        else:
+            transfers_to_wait = list(transfer_statuses)
+
+        if not transfers_to_wait:
+            return
+
+        timeout = self._transfer_timeout
+        deadline = time.monotonic() + timeout
+        remaining = list(transfers_to_wait)
+        errors: list[str] = []
+
+        while remaining:
+            timed_out = time.monotonic() > deadline
+            still_waiting = []
+            for status in remaining:
+                if status.Succeeded():
+                    continue
+                if status.Failed():
+                    errors.append(
+                        f"RDMA transfer failed: {status.Message()} "
+                        f"(code={status.Code()})"
+                    )
+                    continue
+                if timed_out:
+                    errors.append(
+                        f"RDMA transfer timed out after {timeout:.0f}s. "
+                        f"Check NIC queue depth or reduce concurrency; "
+                        f"adjust with kv_connector_extra_config.transfer_timeout."
+                    )
+                else:
+                    still_waiting.append(status)
+            remaining = still_waiting
+            if remaining:
+                time.sleep(0.001)
+        if errors:
+            raise TransferError(
+                f"{len(errors)}/{len(transfers_to_wait)} transfers failed:\n"
+                + "\n".join(errors)
+            )
+
+    def async_wait_reqid(self):
+        assert self.notify_port is not None, "Notify port cannot be None"
+
+        if self.notify_thread is not None:
+            return
+
+        def _async_wait():
+            host = "*"
+            path = make_zmq_path("tcp", host, self.notify_port)
+            logger.info("Node starting to listen notify from path = %s", path)
+
+            with zmq_ctx(zmq.ROUTER, path) as sock:
+                while True:
+                    try:
+                        identity, msg = sock.recv_multipart()
+                        self._handle_message(msg)
+                    except Exception as e:
+                        logger.error("Error processing message: %s", e)
+                        raise HandshakeError(f"Error processing message: {e}") from e
+
+        self.notify_thread = threading.Thread(
+            target=_async_wait, daemon=True, name="moriio-notify-listener"
+        )
+        self.notify_thread.start()
+
+    def _handle_message(self, msg: bytes):
+        """Handles incoming messages from remote nodes."""
+        # Handles incoming remote messages:
+        # Prefill Role:
+        #   [write] mode: receives block information (allocation)
+        #   [read]  mode: receives block release messages from decode side
+        # Decode Role:
+        #   [write] mode: receives KV cache write completion notifications
+        msg_str = repr(msg)
+        handled = False
+        try:
+            data = msgpack.loads(msg)
+            if isinstance(data, dict):
+                self._handle_structured_message(data)
+                return
+        except (
+            msgpack.exceptions.ExtraData,
+            msgpack.exceptions.UnpackException,
+            ValueError,
+        ):
+            logger.debug("Failed to decode msgpack message, will try as string")
+            pass
+
+        try:
+            msg_str = msg.decode("UTF-8")
+            if msg_str:
+                self._handle_completion_message(msg_str)
+                handled = True
+        except UnicodeDecodeError:
+            logger.warning("Received non-UTF8 message: %r", msg)
+        if not handled:
+            raise MoRIIOError(f"Unhandled message format: {msg_str}")
+
+    def _handle_structured_message(self, data: dict):
+        message_type = data.get("type")
+        if message_type is None and "req_id" in data:
+            message_type = "remote_blocks"
+
+        if message_type == "remote_blocks":
+            self._handle_remote_blocks_message(data)
+        elif message_type == "write_done":
+            self._handle_write_done_message(data)
+        elif message_type == "release":
+            self._handle_release_message(data)
+        else:
+            raise MoRIIOError(f"Unhandled structured message type: {message_type}")
+
+    def _handle_remote_blocks_message(self, data: dict):
+        assert get_role() == ROLE.PRODUCER, "Only prefill can get block messages"
+        transfer_id = data["transfer_id"]
+        block_notify_list = data.get("block_notify_list", [])
+        decode_dp_rank = data.get("decode_rank", 0)
+        if not block_notify_list:
+            raise MoRIIOError(
+                "block_notify_list cannot be empty in remote allocate message"
+            )
+
+        with self.lock:
+            if self._is_transfer_terminal_locked(transfer_id):
+                logger.debug(
+                    "Ignoring remote allocation for terminal transfer %s",
+                    transfer_id,
+                )
+                return
+            self.done_remote_allocate_req_dict[transfer_id] = RemoteAllocInfo(
+                block_ids=block_notify_list, decode_dp_rank=decode_dp_rank
+            )
+
+    def _handle_write_done_message(self, data: dict):
+        assert get_role() != ROLE.PRODUCER, (
+            "Only decode can get WRITE completion messages"
+        )
+        transfer_id = data["transfer_id"]
+        with self.lock:
+            self.done_write_cache_req_ids.append(transfer_id)
+
+    def _handle_release_message(self, data: dict):
+        assert get_role() == ROLE.PRODUCER, (
+            "Only prefill can get transfer release messages"
+        )
+        transfer_id = data["transfer_id"]
+        consumer_tp_size = int(data.get("consumer_tp_size", 1))
+        if consumer_tp_size <= 0:
+            raise MoRIIOError(
+                f"Invalid consumer_tp_size in release message: {consumer_tp_size}"
+            )
+        with self.lock:
+            self.done_req_ids.append(MoRIIOTransferAck(transfer_id, consumer_tp_size))
+            self.done_remote_allocate_req_dict.pop(transfer_id, None)
+            self._mark_transfer_terminal_locked(transfer_id)
+
+    def _handle_completion_message(self, msg: str):
+        with self.lock:
+            if get_role() == ROLE.PRODUCER:
+                self.done_req_ids.append(MoRIIOTransferAck(msg))
+                self.done_remote_allocate_req_dict.pop(msg, None)
+                self._mark_transfer_terminal_locked(msg)
+            else:
+                self.done_write_cache_req_ids.append(msg)
+
+    def _is_transfer_terminal_locked(self, transfer_id: TransferId) -> bool:
+        return transfer_id in self._terminal_transfer_ids
+
+    def _mark_transfer_terminal_locked(self, transfer_id: TransferId) -> None:
+        self._terminal_transfer_ids[transfer_id] = None
+        self._terminal_transfer_ids.move_to_end(transfer_id)
+        while len(self._terminal_transfer_ids) > _MAX_TERMINAL_TRANSFER_IDS:
+            self._terminal_transfer_ids.popitem(last=False)
+
+    def send_notify(
+        self,
+        req_ids,
+        remote_ip,
+        remote_port,
+        message_type: str | None = None,
+        message_fields: dict[str, Any] | None = None,
+    ):
+        if not remote_ip or not remote_port:
+            logger.warning("Missing remote_ip or remote_port for notification")
+            return
+
+        path = make_zmq_path("tcp", remote_ip, remote_port)
+
+        if path not in self.paths:
+            ctx = zmq.Context.instance()
+            sock = make_zmq_socket(
+                ctx=ctx, path=path, socket_type=zmq.DEALER, bind=False
+            )
+            self.paths[path] = sock
+
+        req_list = req_ids if isinstance(req_ids, list) else [req_ids]
+
+        sock = self.paths[path]
+        try:
+            for req_id in req_list:
+                if not isinstance(req_id, str):
+                    logger.warning(
+                        "Invalid req_id type: %s, expected str", type(req_id)
+                    )
+                    continue
+                if message_type is None:
+                    sock.send(req_id.encode("utf-8"))
+                else:
+                    payload = {"type": message_type, "transfer_id": req_id}
+                    if message_fields:
+                        payload.update(message_fields)
+                    sock.send(msgpack.dumps(payload))
+        except Exception as e:
+            logger.error("Failed to send notification to %s: %s", path, e)
+            self.paths.pop(path, None)
+            raise
+
+    def pop_finished_req_ids(self):
+        # Producer invocation: return every completion message since the last
+        # call. Do not dedupe: heterogeneous TP can produce multiple release
+        # ACKs for the same transfer_id and the caller must count each one.
+        with self.lock:
+            done_send = list(self.done_req_ids)
+            self.done_req_ids = []
+        return done_send
+
+    def pop_finished_write_req_ids(self):
+        # Call the consumer in write mode to get the collection after write completion
+        with self.lock:
+            done_write_cache = set(self.done_write_cache_req_ids)
+            self.done_write_cache_req_ids = []
+        return done_write_cache
+
+    def shutdown(self):
+        logger.debug("Closing MoRIIOWrapper and cleaning up ZMQ sockets")
+        for path, sock in self.paths.items():
+            try:
+                sock.close(linger=0)
+                logger.debug("Closed ZMQ socket for path: %s", path)
+            except Exception as e:
+                logger.warning("Error closing ZMQ socket for path %s: %s", path, e)
+        self.paths.clear()

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py

@@ -0,0 +1,347 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Callable, Mapping
+from typing import NamedTuple
+
+import torch
+
+from vllm.v1.kv_cache_interface import (
+    AttentionSpec,
+    KVCacheConfig,
+    KVCacheSpec,
+    MLAAttentionSpec,
+    SlidingWindowMLASpec,
+    UniformTypeKVCacheSpecs,
+)
+
+
+class LayerTransferGeometry(NamedTuple):
+    num_blocks: int
+    block_size: int
+    block_len: int
+    slot_size_bytes: int
+    block_stride: int
+    local_kv_stride: int | None
+    remote_kv_stride: int | None
+    transfers_per_block: int
+    regions_per_block: int
+    split_kv_regions: bool
+
+
+def build_layer_to_spec(kv_cache_config: KVCacheConfig) -> dict[str, KVCacheSpec]:
+    layer_to_spec: dict[str, KVCacheSpec] = {}
+    for group in kv_cache_config.kv_cache_groups:
+        group_spec = group.kv_cache_spec
+        if isinstance(group_spec, UniformTypeKVCacheSpecs):
+            layer_to_spec.update(
+                {
+                    layer_name: group_spec.kv_cache_specs[layer_name]
+                    for layer_name in group.layer_names
+                }
+            )
+        else:
+            layer_to_spec.update(
+                {layer_name: group_spec for layer_name in group.layer_names}
+            )
+    return layer_to_spec
+
+
+def is_mla_cache_layer(
+    layer_to_spec: Mapping[str, KVCacheSpec], layer_name: str
+) -> bool:
+    try:
+        spec = layer_to_spec[layer_name]
+    except KeyError as e:
+        raise ValueError(f"Missing KV cache spec for layer {layer_name}") from e
+    return isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec))
+
+
+def _content_packed_dim(spec: AttentionSpec) -> int:
+    head_size_v = getattr(spec, "head_size_v", spec.head_size)
+    return spec.head_size + head_size_v
+
+
+def _spec_dim_matches(value: int, expected: int | None) -> bool:
+    return expected is None or value == expected
+
+
+def _kernel_layout_matches(
+    spec: KVCacheSpec, kernel_block_size: int, num_kv_heads: int, head_dim: int
+) -> bool:
+    if kernel_block_size <= 0 or spec.block_size % kernel_block_size != 0:
+        return False
+    return _spec_dim_matches(
+        num_kv_heads, getattr(spec, "num_kv_heads", None)
+    ) and _spec_dim_matches(head_dim, getattr(spec, "head_size", None))
+
+
+def _select_kernel_block_layout(
+    layer_name: str, shape: torch.Size, spec: KVCacheSpec
+) -> tuple[int, int, int]:
+    axis2_matches = _kernel_layout_matches(spec, shape[2], shape[3], shape[4])
+    axis3_matches = _kernel_layout_matches(spec, shape[3], shape[2], shape[4])
+
+    if axis2_matches and axis3_matches and shape[2] != shape[3]:
+        raise ValueError(
+            f"Ambiguous MoRIIO kernel-block K/V cache shape for layer "
+            f"{layer_name}: {tuple(shape)}"
+        )
+    if axis2_matches:
+        return shape[2], shape[3], shape[4]
+    if axis3_matches:
+        return shape[3], shape[2], shape[4]
+
+    raise ValueError(
+        f"Unsupported MoRIIO K/V cache shape for layer {layer_name}: "
+        f"{tuple(shape)} does not contain block size {spec.block_size}"
+    )
+
+
+def get_layer_transfer_geometry(
+    layer_name: str,
+    kv_cache: torch.Tensor,
+    layer_to_spec: Mapping[str, KVCacheSpec],
+    remote_num_blocks: int | None = None,
+) -> LayerTransferGeometry:
+    shape = kv_cache.shape
+    stride = kv_cache.stride()
+    element_size = kv_cache.element_size()
+    spec = layer_to_spec[layer_name]
+    is_mla_cache = is_mla_cache_layer(layer_to_spec, layer_name)
+
+    if is_mla_cache and len(shape) == 3:
+        num_blocks, block_size, latent_dim = shape
+        slot_size_bytes = latent_dim * element_size
+        block_len = block_size * slot_size_bytes
+        return LayerTransferGeometry(
+            num_blocks=num_blocks,
+            block_size=block_size,
+            block_len=block_len,
+            slot_size_bytes=slot_size_bytes,
+            block_stride=stride[0],
+            local_kv_stride=None,
+            remote_kv_stride=None,
+            transfers_per_block=1,
+            regions_per_block=1,
+            split_kv_regions=False,
+        )
+
+    if not is_mla_cache and len(shape) == 5 and shape[0] == 2:
+        _, num_blocks = shape[:2]
+        kernel_blocks_per_block = 1
+        if shape[2] == spec.block_size:
+            block_size, num_kv_heads, head_dim = shape[2:]
+        elif shape[3] == spec.block_size:
+            num_kv_heads, block_size, head_dim = shape[2:]
+        else:
+            kernel_num_blocks = num_blocks
+            kernel_block_size, num_kv_heads, head_dim = _select_kernel_block_layout(
+                layer_name, shape, spec
+            )
+            kernel_blocks_per_block = spec.block_size // kernel_block_size
+            if kernel_num_blocks % kernel_blocks_per_block != 0:
+                raise ValueError(
+                    f"Unsupported MoRIIO K/V cache shape for layer {layer_name}: "
+                    f"{tuple(shape)} has {kernel_num_blocks} kernel blocks, "
+                    f"not divisible by {kernel_blocks_per_block}"
+                )
+            num_blocks = kernel_num_blocks // kernel_blocks_per_block
+            block_size = spec.block_size
+        slot_size_bytes = num_kv_heads * head_dim * element_size
+        block_len = block_size * slot_size_bytes
+        return LayerTransferGeometry(
+            num_blocks=num_blocks,
+            block_size=block_size,
+            block_len=block_len,
+            slot_size_bytes=slot_size_bytes,
+            block_stride=stride[1] * kernel_blocks_per_block,
+            local_kv_stride=stride[0],
+            remote_kv_stride=(
+                stride[1] * kernel_blocks_per_block * (remote_num_blocks or num_blocks)
+            ),
+            transfers_per_block=2,
+            regions_per_block=1,
+            split_kv_regions=True,
+        )
+
+    if not is_mla_cache and len(shape) == 5 and shape[1] == 2:
+        num_blocks = shape[0]
+        if shape[2] == spec.block_size:
+            block_size, num_kv_heads, head_dim = shape[2:]
+            slot_size_bytes = num_kv_heads * head_dim * element_size
+            block_len = block_size * slot_size_bytes
+            return LayerTransferGeometry(
+                num_blocks=num_blocks,
+                block_size=block_size,
+                block_len=block_len,
+                slot_size_bytes=slot_size_bytes,
+                block_stride=stride[0],
+                local_kv_stride=stride[1],
+                remote_kv_stride=stride[1],
+                transfers_per_block=2,
+                regions_per_block=2,
+                split_kv_regions=False,
+            )
+        elif shape[3] == spec.block_size:
+            num_kv_heads, block_size, head_dim = shape[2:]
+        else:
+            kernel_num_blocks = num_blocks
+            kernel_block_size, _, _ = _select_kernel_block_layout(
+                layer_name, shape, spec
+            )
+            kernel_blocks_per_block = spec.block_size // kernel_block_size
+            if kernel_num_blocks % kernel_blocks_per_block != 0:
+                raise ValueError(
+                    f"Unsupported MoRIIO K/V cache shape for layer {layer_name}: "
+                    f"{tuple(shape)} has {kernel_num_blocks} kernel blocks, "
+                    f"not divisible by {kernel_blocks_per_block}"
+                )
+            num_blocks = kernel_num_blocks // kernel_blocks_per_block
+            block_size = spec.block_size
+            block_stride = stride[0] * kernel_blocks_per_block
+            block_len = block_stride * element_size
+            slot_size_bytes = block_len // block_size
+            return LayerTransferGeometry(
+                num_blocks=num_blocks,
+                block_size=block_size,
+                block_len=block_len,
+                slot_size_bytes=slot_size_bytes,
+                block_stride=block_stride,
+                local_kv_stride=None,
+                remote_kv_stride=None,
+                transfers_per_block=1,
+                regions_per_block=1,
+                split_kv_regions=False,
+            )
+        slot_size_bytes = num_kv_heads * head_dim * element_size
+        block_len = block_size * slot_size_bytes
+        return LayerTransferGeometry(
+            num_blocks=num_blocks,
+            block_size=block_size,
+            block_len=block_len,
+            slot_size_bytes=slot_size_bytes,
+            block_stride=stride[0],
+            local_kv_stride=stride[1],
+            remote_kv_stride=stride[1],
+            transfers_per_block=2,
+            regions_per_block=2,
+            split_kv_regions=False,
+        )
+
+    if (
+        not is_mla_cache
+        and isinstance(spec, AttentionSpec)
+        and len(shape) == 4
+        and shape[1] == spec.num_kv_heads
+        and shape[2] == spec.block_size
+        and shape[3] == _content_packed_dim(spec)
+    ):
+        num_blocks, num_kv_heads, block_size, packed_dim = shape
+        slot_size_bytes = num_kv_heads * packed_dim * element_size
+        block_len = block_size * slot_size_bytes
+        return LayerTransferGeometry(
+            num_blocks=num_blocks,
+            block_size=block_size,
+            block_len=block_len,
+            slot_size_bytes=slot_size_bytes,
+            block_stride=stride[0],
+            local_kv_stride=None,
+            remote_kv_stride=None,
+            transfers_per_block=1,
+            regions_per_block=1,
+            split_kv_regions=False,
+        )
+
+    cache_kind = "MLA" if is_mla_cache else "K/V"
+    raise ValueError(
+        f"Unsupported MoRIIO {cache_kind} cache shape for layer "
+        f"{layer_name}: {tuple(shape)}"
+    )
+
+
+def iter_layer_registration_regions(
+    layer_name: str,
+    kv_cache: torch.Tensor,
+    layer_to_spec: Mapping[str, KVCacheSpec],
+) -> list[tuple[torch.Tensor, int]]:
+    geometry = get_layer_transfer_geometry(layer_name, kv_cache, layer_to_spec)
+    region_len = geometry.num_blocks * geometry.regions_per_block * geometry.block_len
+    if geometry.split_kv_regions:
+        return [(cache, region_len) for cache in kv_cache]
+    return [(kv_cache, region_len)]
+
+
+def merge_contiguous_offsets(
+    offsets_local: list[int],
+    offsets_remote: list[int],
+    sizes: list[int],
+) -> tuple[list[int], list[int], list[int]]:
+    if not offsets_local:
+        return [], [], []
+    if not (len(offsets_local) == len(offsets_remote) == len(sizes)):
+        raise ValueError("Input list lengths mismatch")
+
+    rows = sorted(zip(offsets_local, offsets_remote, sizes), key=lambda row: row[0])
+    merged: list[list[int]] = []
+    for local, remote, size in rows:
+        if (
+            merged
+            and local == merged[-1][0] + merged[-1][2]
+            and remote == merged[-1][1] + merged[-1][2]
+        ):
+            merged[-1][2] += size
+        else:
+            merged.append([local, remote, size])
+
+    return (
+        [row[0] for row in merged],
+        [row[1] for row in merged],
+        [row[2] for row in merged],
+    )
+
+
+def compute_block_transfer_offsets(
+    layer_name: str,
+    kv_cache: torch.Tensor,
+    layer_to_spec: Mapping[str, KVCacheSpec],
+    local_block_ids: list[int],
+    remote_block_ids: list[int],
+    remote_num_blocks: int,
+    merge_fn: Callable[
+        [list[int], list[int], list[int]], tuple[list[int], list[int], list[int]]
+    ] = merge_contiguous_offsets,
+) -> tuple[list[int], list[int], list[int]]:
+    if len(local_block_ids) != len(remote_block_ids):
+        raise ValueError(
+            "local_block_ids and remote_block_ids must have the same length: "
+            f"{len(local_block_ids)} != {len(remote_block_ids)}"
+        )
+    geometry = get_layer_transfer_geometry(
+        layer_name, kv_cache, layer_to_spec, remote_num_blocks
+    )
+    element_size = kv_cache.element_size()
+    transfer_size_byte = geometry.block_len
+    per_block = geometry.transfers_per_block
+    total = len(local_block_ids) * per_block
+    offset_local = [0] * total
+    offset_remote = [0] * total
+    sizes = [transfer_size_byte] * total
+
+    w = 0
+    for lb, rb in zip(local_block_ids, remote_block_ids):
+        offset_local[w] = element_size * (lb * geometry.block_stride)
+        offset_remote[w] = element_size * (rb * geometry.block_stride)
+        w += 1
+        if per_block == 2:
+            assert geometry.local_kv_stride is not None
+            assert geometry.remote_kv_stride is not None
+            offset_local[w] = element_size * (
+                geometry.local_kv_stride + lb * geometry.block_stride
+            )
+            offset_remote[w] = element_size * (
+                geometry.remote_kv_stride + rb * geometry.block_stride
+            )
+            w += 1
+
+    return merge_fn(offset_local, offset_remote, sizes)

vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py

@@ -0,0 +1,667 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import copy
+from collections.abc import Callable, Iterable
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, cast
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.config.kv_transfer import KVTransferConfig
+from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBaseType
+from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    CopyBlocksOp,
+    KVConnectorBase_V1,
+    KVConnectorHandshakeMetadata,
+    KVConnectorMetadata,
+    KVConnectorRole,
+    KVConnectorWorkerMetadata,
+    SupportsHMA,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
+    KVConnectorPromMetrics,
+    KVConnectorStats,
+    PromMetric,
+    PromMetricT,
+)
+from vllm.logger import init_logger
+from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.outputs import KVConnectorOutput
+
+if TYPE_CHECKING:
+    from vllm.distributed.kv_events import KVCacheEvent
+    from vllm.forward_context import ForwardContext
+    from vllm.v1.core.block_pool import BlockPool
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+    from vllm.v1.request import Request
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class MultiKVConnectorMetadata(KVConnectorMetadata):
+    metadata: tuple[KVConnectorMetadata, ...]
+    extra_async_saves: dict[str, int] | None = None
+
+
+@dataclass
+class MultiKVConnectorWorkerMetadata(KVConnectorWorkerMetadata):
+    metadata: tuple[KVConnectorWorkerMetadata | None, ...]
+
+    def aggregate(self, other: KVConnectorWorkerMetadata) -> KVConnectorWorkerMetadata:
+        assert isinstance(other, MultiKVConnectorWorkerMetadata)
+
+        assert len(self.metadata) == len(other.metadata)
+        metadata_list = []
+        for metadata1, metadata2 in zip(self.metadata, other.metadata):
+            if metadata1 is None:
+                metadata_list.append(metadata2)
+            elif metadata2 is None:
+                metadata_list.append(metadata1)
+            else:
+                metadata_list.append(metadata1.aggregate(metadata2))
+
+        return MultiKVConnectorWorkerMetadata(metadata=tuple(metadata_list))
+
+
+@dataclass
+class MultiKVConnectorStats(KVConnectorStats):
+    """
+    Maintain a dict of KVConnectorStats objects, one for each connector.
+    This is used to aggregate the stats from all connectors separately.
+    """
+
+    def aggregate(self, other: KVConnectorStats) -> KVConnectorStats:
+        for connector_id, stats in other.data.items():
+            if connector_id not in self.data:
+                self[connector_id] = stats
+            else:
+                assert isinstance(stats, type(self.data[connector_id]))
+                self[connector_id] = self[connector_id].aggregate(stats)
+        return self
+
+    def reset(self):
+        for stats in self.data.values():
+            stats.reset()
+
+    def reduce(self) -> dict[str, Any]:
+        # TODO (NickLucche) Adjust for logging on separate lines
+        return {
+            connector_id: stats.reduce() for connector_id, stats in self.data.items()
+        }
+
+    def is_empty(self) -> bool:
+        return all(stats.is_empty() for stats in self.data.values())
+
+    def __getitem__(self, connector_id: str) -> KVConnectorStats:
+        return self.data[connector_id]
+
+    def __setitem__(self, connector_id: str, stats: KVConnectorStats):
+        self.data[connector_id] = stats
+
+
+class MultiKVConnectorPromMetrics(KVConnectorPromMetrics):
+    def __init__(
+        self,
+        vllm_config: "VllmConfig",
+        metric_types: dict[type[PromMetric], type[PromMetricT]],
+        labelnames: list[str],
+        per_engine_labelvalues: dict[int, list[object]],
+        prom_metrics: dict[str, KVConnectorPromMetrics],
+    ):
+        super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues)
+        self._prom_metrics = prom_metrics
+
+    def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
+        for connector_id, stats_data in transfer_stats_data.items():
+            assert connector_id in self._prom_metrics, (
+                f"{connector_id} is not contained in the list of registered connectors "
+                f"with Prometheus metrics support: {self._prom_metrics.keys()}"
+            )
+            self._prom_metrics[connector_id].observe(stats_data["data"], engine_idx)
+
+
+class MultiConnector(KVConnectorBase_V1, SupportsHMA):
+    """
+    A wrapper for using multiple KVConnectors at the same time.
+
+    The current logic is:
+    - Load KV from the first connector that advertises available tokens from
+      get_num_new_matched_tokens(), based on the order in the config.
+    - Save to all connectors.
+    """
+
+    @classmethod
+    def requires_piecewise_for_cudagraph(cls, extra_config: dict[str, Any]) -> bool:
+        """
+        MultiConnector requires PIECEWISE CUDA graph mode if any of its
+        child connectors require it.
+        """
+        connectors_config = extra_config.get("connectors", [])
+        for conn_config in connectors_config:
+            temp_ktc = KVTransferConfig(**conn_config)
+            connector_cls = KVConnectorFactory.get_connector_class(temp_ktc)
+            child_extra_config = conn_config.get("kv_connector_extra_config", {})
+            if connector_cls.requires_piecewise_for_cudagraph(child_extra_config):
+                return True
+        return False
+
+    @classmethod
+    def all_children_support_hma(cls, kv_transfer_config: "KVTransferConfig") -> bool:
+        """Return True only if every configured child connector supports HMA."""
+        connectors_config = kv_transfer_config.kv_connector_extra_config.get(
+            "connectors", []
+        )
+        if not connectors_config:
+            return False
+        for conn_config in connectors_config:
+            child_config = KVTransferConfig(
+                **{"engine_id": kv_transfer_config.engine_id, **conn_config}
+            )
+            if not KVConnectorFactory.supports_hma_config(child_config):
+                return False
+        return True
+
+    def __init__(
+        self,
+        vllm_config: "VllmConfig",
+        role: KVConnectorRole,
+        kv_cache_config: "KVCacheConfig",
+    ):
+        super().__init__(
+            vllm_config=vllm_config, role=role, kv_cache_config=kv_cache_config
+        )
+
+        self._connectors: list[KVConnectorBase_V1] = []
+        self._ktc_kv_transfer_config = []
+        for connector_cls, temp_config in self._get_connector_classes_and_configs(
+            vllm_config
+        ):
+            self._connectors.append(connector_cls(temp_config, role, kv_cache_config))
+            self._ktc_kv_transfer_config.append(temp_config.kv_transfer_config)
+
+        assert vllm_config.kv_transfer_config is not None
+        self._all_support_hma = MultiConnector.all_children_support_hma(
+            vllm_config.kv_transfer_config
+        )
+        assert (
+            vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
+            or self._all_support_hma
+        ), "HMA should not be enabled unless all sub-connectors support it"
+
+        # A mapping from request id to the index of the connector chosen to
+        # load the request from (if any).
+        self._requests_to_connector: dict[str, int] = {}
+
+        # Keeps track of *additional* remaining async saves (beyond 1) to be
+        # finished per request. Not needed for async loads since we only allow
+        # a single connector to load.
+        # Propagated from scheduler to worker side via the connector metadata.
+        self._extra_async_saves: dict[str, int] = {}
+
+    @property
+    def prefer_cross_layer_blocks(self) -> bool:
+        if not self._connectors:
+            return False
+        return all(c.prefer_cross_layer_blocks for c in self._connectors)
+
+    @classmethod
+    def _get_connector_classes_and_configs(
+        cls, vllm_config: "VllmConfig"
+    ) -> list[tuple[type[KVConnectorBaseType], "VllmConfig"]]:
+        assert vllm_config.kv_transfer_config is not None
+        ktcs = vllm_config.kv_transfer_config.kv_connector_extra_config.get(
+            "connectors"
+        )
+        assert ktcs is not None
+        ret: list[tuple[type[KVConnectorBaseType], VllmConfig]] = []
+        for ktc in ktcs:
+            temp_config = copy.copy(vllm_config)
+            engine_id = ktc.get("engine_id", vllm_config.kv_transfer_config.engine_id)
+            temp_config.kv_transfer_config = KVTransferConfig(
+                **ktc, engine_id=engine_id
+            )
+            ret.append(
+                (
+                    KVConnectorFactory.get_connector_class(
+                        temp_config.kv_transfer_config
+                    ),
+                    temp_config,
+                )
+            )
+        return ret
+
+    def register_cross_layers_kv_cache(
+        self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
+    ):
+        # Register on all connectors
+        for c in self._connectors:
+            c.register_cross_layers_kv_cache(kv_cache, attn_backend)
+
+    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
+        for c in self._connectors:
+            c.register_kv_caches(kv_caches)
+
+    def bind_gpu_block_pool(self, gpu_block_pool: "BlockPool") -> None:
+        for c in self._connectors:
+            c.bind_gpu_block_pool(gpu_block_pool)
+
+    # We must override the base class method here because we need to bind
+    # the metadata to each connector in the order of the connectors in the
+    # MultiKVConnectorMetadata.
+    #
+    # Note: Call the base class method to ensure metadata is also set on the
+    # MultiConnector instance itself; otherwise, `has_connector_metadata()` will
+    # always return False.
+    def bind_connector_metadata(self, connector_metadata: KVConnectorMetadata) -> None:
+        assert isinstance(connector_metadata, MultiKVConnectorMetadata)
+        if connector_metadata.extra_async_saves:
+            self._extra_async_saves.update(connector_metadata.extra_async_saves)
+        for c, cm in zip(self._connectors, connector_metadata.metadata):
+            c.bind_connector_metadata(cm)
+        super().bind_connector_metadata(connector_metadata)
+
+    def clear_connector_metadata(self) -> None:
+        for c in self._connectors:
+            c.clear_connector_metadata()
+        super().clear_connector_metadata()
+
+    def shutdown(self):
+        exception: Exception | None = None
+        for c in self._connectors:
+            try:
+                c.shutdown()
+            except Exception as e:
+                logger.exception(
+                    "Exception during connector %s shutdown.", c.__class__.__name__
+                )
+                exception = e
+        if exception:
+            raise exception
+
+    # ==============================
+    # Worker-side methods
+    # ==============================
+    def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
+        for c in self._connectors:
+            c.start_load_kv(forward_context, **kwargs)
+
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        for c in self._connectors:
+            c.wait_for_layer_load(layer_name)
+
+    def save_kv_layer(
+        self,
+        layer_name: str,
+        kv_layer: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        **kwargs,
+    ) -> None:
+        for c in self._connectors:
+            c.save_kv_layer(layer_name, kv_layer, attn_metadata, **kwargs)
+
+    def wait_for_save(self):
+        for c in self._connectors:
+            c.wait_for_save()
+
+    def get_finished(
+        self, finished_req_ids: set[str]
+    ) -> tuple[set[str] | None, set[str] | None]:
+        finished_sending: set[str] = set()
+        finished_recving: set[str] = set()
+        for c in self._connectors:
+            sending, recving = c.get_finished(finished_req_ids)
+            if not recving and not sending:
+                continue
+            # Aggregate finished recving request ids.
+            finished_recving.update(recving or ())
+            # Aggregate finished sending request ids - only include
+            # once we've drained the "extra" count (for cases where
+            # more than one connector is async-saving the same request).
+            for req_id in sending or ():
+                extra_pending = self._extra_async_saves.get(req_id)
+                if extra_pending is None:
+                    finished_sending.add(req_id)
+                    continue
+                assert extra_pending > 0
+                if extra_pending == 1:
+                    del self._extra_async_saves[req_id]
+                else:
+                    self._extra_async_saves[req_id] = extra_pending - 1
+
+        return finished_sending or None, finished_recving or None
+
+    def get_block_ids_with_load_errors(self) -> set[int]:
+        agg_block_ids: set[int] = set()
+        for c in self._connectors:
+            agg_block_ids |= c.get_block_ids_with_load_errors()
+        return agg_block_ids
+
+    def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
+        """Set xPU-specific copy ops for all sub-connectors."""
+        for c in self._connectors:
+            c.set_host_xfer_buffer_ops(copy_operation)
+
+    def handle_preemptions(self, kv_connector_metadata: KVConnectorMetadata):
+        """Handle preempted requests for all sub-connectors."""
+        assert isinstance(kv_connector_metadata, MultiKVConnectorMetadata)
+        for c, cm in zip(self._connectors, kv_connector_metadata.metadata):
+            c.handle_preemptions(cm)
+
+    def get_finished_count(self) -> int | None:
+        # TODO(https://github.com/vllm-project/vllm/issues/33400)
+        # Currently no connectors return non-None
+        return None
+
+    def build_connector_worker_meta(self) -> KVConnectorWorkerMetadata | None:
+        metadata_list: list[KVConnectorWorkerMetadata | None] | None = None
+        for i, c in enumerate(self._connectors):
+            kv_connector_worker_meta = c.build_connector_worker_meta()
+            if metadata_list is None and kv_connector_worker_meta is not None:
+                metadata_list = [None] * i
+            if metadata_list is not None:
+                metadata_list.append(kv_connector_worker_meta)
+        if metadata_list is None:
+            return None
+        return MultiKVConnectorWorkerMetadata(metadata=tuple(metadata_list))
+
+    # TODO: Add a generic implementation of 'get_kv_connector_kv_cache_events'
+    # method for the MultiConnector. It should be able to get events from
+    # multiple connectors, handling the case where only a subset of the
+    # requested connectors implements the 'get_kv_connector_kv_cache_events'
+    # WIP: https://github.com/vllm-project/vllm/pull/31811
+
+    # ==============================
+    # Scheduler-side methods
+    # ==============================
+    def get_num_new_matched_tokens(
+        self,
+        request: "Request",
+        num_computed_tokens: int,
+    ) -> tuple[int | None, bool]:
+        to_return = (0, False)
+        for i, c in enumerate(self._connectors):
+            toks, load_async = c.get_num_new_matched_tokens(
+                request, num_computed_tokens
+            )
+            # If there is a connector still looking up the matches,
+            # we return None to indicate that we are not done yet.
+            if toks is None:
+                return (None, False)
+            # The first connector that has new matched tokens will be assigned
+            # to this request.
+            if to_return[0] == 0 and toks > 0:
+                self._requests_to_connector[request.request_id] = i
+                to_return = (toks, load_async)
+        return to_return
+
+    def update_state_after_alloc(
+        self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
+    ):
+        chosen_connector = self._requests_to_connector.get(request.request_id, -1)
+        for i, c in enumerate(self._connectors):
+            if i == chosen_connector:
+                # Forward call to the chosen connector (if any).
+                c.update_state_after_alloc(request, blocks, num_external_tokens)
+            else:
+                # Other connectors still receive the request's real blocks
+                c.update_state_after_alloc(request, blocks, 0)
+
+    def on_new_request(self, request: "Request") -> None:
+        for c in self._connectors:
+            c.on_new_request(request)
+
+    def build_connector_meta(
+        self, scheduler_output: SchedulerOutput
+    ) -> MultiKVConnectorMetadata:
+        metadata = MultiKVConnectorMetadata(
+            metadata=tuple(
+                c.build_connector_meta(scheduler_output) for c in self._connectors
+            )
+        )
+        if self._extra_async_saves:
+            metadata.extra_async_saves = self._extra_async_saves
+            self._extra_async_saves = {}
+        return metadata
+
+    def update_connector_output(self, connector_output: KVConnectorOutput):
+        multi_connector_worker_meta: MultiKVConnectorWorkerMetadata | None = None
+        if connector_output.kv_connector_worker_meta is not None:
+            assert isinstance(
+                connector_output.kv_connector_worker_meta,
+                MultiKVConnectorWorkerMetadata,
+            )
+            multi_connector_worker_meta = connector_output.kv_connector_worker_meta
+
+        try:
+            for i, c in enumerate(self._connectors):
+                if multi_connector_worker_meta is not None:
+                    # set the connector-specific worker metadata
+                    connector_output.kv_connector_worker_meta = (
+                        multi_connector_worker_meta.metadata[i]
+                    )
+                c.update_connector_output(connector_output)
+        finally:
+            # restore kv_connector_worker_meta
+            connector_output.kv_connector_worker_meta = multi_connector_worker_meta
+
+    def get_handshake_metadata(self) -> KVConnectorHandshakeMetadata | None:
+        """
+        Get the KVConnector handshake metadata from sub-connectors.
+        Returns the first non-None metadata from sub-connectors.
+        """
+        for c in self._connectors:
+            metadata = c.get_handshake_metadata()
+            if metadata is not None:
+                return metadata
+        return None
+
+    def set_xfer_handshake_metadata(
+        self, metadata: dict[int, KVConnectorHandshakeMetadata]
+    ) -> None:
+        """
+        Set the KV connector handshake metadata for all sub-connectors.
+        This is needed to start the NIXL listener thread for NixlConnector.
+        """
+        for c in self._connectors:
+            c.set_xfer_handshake_metadata(metadata)
+
+    def set_xfer_handshake_metadata_pp_aware(
+        self, metadata: dict[tuple[int, int], KVConnectorHandshakeMetadata]
+    ) -> None:
+        for c in self._connectors:
+            c.set_xfer_handshake_metadata_pp_aware(metadata)
+
+    def _aggregate_request_finished(
+        self,
+        request: "Request",
+        per_connector_fn: Callable[
+            [KVConnectorBase_V1], tuple[bool, dict[str, Any] | None]
+        ],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        async_saves = 0
+        kv_txfer_params = None
+        for c in self._connectors:
+            async_save, txfer_params = per_connector_fn(c)
+            if async_save:
+                async_saves += 1
+            if txfer_params is not None:
+                if kv_txfer_params is not None:
+                    clashes = set(kv_txfer_params) & set(txfer_params)
+                    if clashes:
+                        raise RuntimeError(
+                            "Key clash in kv_transfer_params from multiple "
+                            f"connectors: {clashes}"
+                        )
+                    kv_txfer_params.update(txfer_params)
+                else:
+                    kv_txfer_params = txfer_params
+        if async_saves > 1:
+            self._extra_async_saves[request.request_id] = async_saves - 1
+
+        self._requests_to_connector.pop(request.request_id, None)
+
+        return async_saves > 0, kv_txfer_params
+
+    def request_finished(
+        self,
+        request: "Request",
+        blocks: list[int],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        return self._aggregate_request_finished(
+            request,
+            lambda c: c.request_finished(request, blocks),
+        )
+
+    def request_finished_all_groups(
+        self,
+        request: "Request",
+        block_ids: tuple[list[int], ...],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        if not self._all_support_hma:
+            assert len(block_ids) == 1, (
+                "HMA with multiple kv_cache_groups requires all "
+                "sub-connectors to support HMA"
+            )
+            return self.request_finished(request, block_ids[0])
+
+        return self._aggregate_request_finished(
+            request,
+            lambda c: cast(SupportsHMA, c).request_finished_all_groups(
+                request, block_ids
+            ),
+        )
+
+    def take_events(self) -> Iterable["KVCacheEvent"]:
+        for c in self._connectors:
+            yield from c.take_events()
+
+    def has_pending_push_work(self) -> bool:
+        return any(c.has_pending_push_work() for c in self._connectors)
+
+    @classmethod
+    def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None:
+        """
+        Get the required KV cache layout for this connector.
+        Args:
+            vllm_config (VllmConfig): the vllm config.
+
+        Returns:
+            str: the required KV cache layout. e.g. HND, or NHD.
+            None if the connector does not require a specific layout.
+        """
+        assert vllm_config.kv_transfer_config is not None
+        layouts: set[str] = set()
+        for connector_cls, temp_config in cls._get_connector_classes_and_configs(
+            vllm_config
+        ):
+            required_kvcache_layout = connector_cls.get_required_kvcache_layout(
+                temp_config
+            )
+            if required_kvcache_layout is not None:
+                layouts.add(required_kvcache_layout)
+
+        if len(layouts) > 1:
+            raise ValueError(
+                f"KV cache layout mismatch: "
+                f"found {len(layouts)} different layouts "
+                f"({', '.join(layouts)})."
+                f"All connectors must use the same layout."
+            )
+        return next(iter(layouts), None)
+
+    @classmethod
+    def build_kv_connector_stats(
+        cls, data: dict[str, Any] | None = None
+    ) -> KVConnectorStats | None:
+        if data is None:
+            return MultiKVConnectorStats()
+
+        # data is a dict mapping connector name to their stats data.
+        # The stats data can be either:
+        # 1. Already-instantiated KVConnectorStats objects (same process)
+        # 2. Serialized dicts (cross-process after serialization)
+        # We need to reconstruct proper KVConnectorStats objects from dicts
+        reconstructed_data = {}
+        for connector_name, stats_value in data.items():
+            # If already a KVConnectorStats object, use it directly
+            if isinstance(stats_value, KVConnectorStats):
+                reconstructed_data[connector_name] = stats_value
+                continue
+
+            # Otherwise, reconstruct from serialized dict
+            # Get the connector class to reconstruct its stats
+            connector_cls = KVConnectorFactory.get_connector_class_by_name(
+                connector_name
+            )
+
+            # stats_value is the serialized dataclass which contains {'data': {...}}
+            # We need to extract the inner 'data' field to avoid double-nesting
+            assert isinstance(stats_value, dict) and "data" in stats_value, (
+                f"Expected a dict with a 'data' field, got {stats_value}"
+            )
+            inner_data = stats_value["data"]
+
+            # Use the connector's build_kv_connector_stats to reconstruct
+            if reconstructed_stats := connector_cls.build_kv_connector_stats(
+                data=inner_data
+            ):
+                reconstructed_data[connector_name] = reconstructed_stats
+
+        return MultiKVConnectorStats(data=reconstructed_data)
+
+    def get_kv_connector_stats(self) -> MultiKVConnectorStats | None:
+        # Group connector stats by connector type.
+        stats_by_connector: MultiKVConnectorStats | None = None
+        for c in self._connectors:
+            stats = c.get_kv_connector_stats()
+            if stats is None:
+                continue
+            if stats_by_connector is None:
+                # Lazy init to allow optional return value.
+                stats_by_connector = MultiKVConnectorStats()
+            connector_id = c.__class__.__name__
+            if connector_id in stats_by_connector.data:
+                stats_by_connector[connector_id] = stats_by_connector[
+                    connector_id
+                ].aggregate(stats)
+            else:
+                stats_by_connector[connector_id] = stats
+        return stats_by_connector
+
+    @classmethod
+    def build_prom_metrics(
+        cls,
+        vllm_config: "VllmConfig",
+        metric_types: dict[type["PromMetric"], type["PromMetricT"]],
+        labelnames: list[str],
+        per_engine_labelvalues: dict[int, list[object]],
+    ) -> KVConnectorPromMetrics:
+        prom_metrics: dict[str, KVConnectorPromMetrics] = {}
+        seen_classes: set[type] = set()
+        for connector_cls, temp_config in cls._get_connector_classes_and_configs(
+            vllm_config
+        ):
+            if connector_cls in seen_classes:
+                continue
+            seen_classes.add(connector_cls)
+            connector_prom = connector_cls.build_prom_metrics(
+                temp_config, metric_types, labelnames, per_engine_labelvalues
+            )
+            if connector_prom is not None:
+                prom_metrics[connector_cls.__name__] = connector_prom
+        return MultiKVConnectorPromMetrics(
+            vllm_config,
+            metric_types,
+            labelnames,
+            per_engine_labelvalues,
+            prom_metrics,
+        )
+
+    def reset_cache(self) -> bool:
+        results = [c.reset_cache() is not False for c in self._connectors]
+        return all(results)

vllm/distributed/kv_transfer/kv_connector/v1/nixl/__init__.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NIXL KV-cache transfer connector (disaggregated prefill / decode)."""
+
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_scheduler import (
+    NixlBaseConnectorScheduler,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import (
+    NixlBaseConnectorWorker,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.connector import (
+    NixlBaseConnector,
+    NixlConnector,
+    NixlPullConnector,
+    NixlPushConnector,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
+    NixlAgentMetadata,
+    NixlConnectorMetadata,
+    NixlHandshakePayload,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import (
+    NixlPullConnectorScheduler,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_worker import (
+    NixlPullConnectorWorker,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_scheduler import (
+    NixlPushConnectorScheduler,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import (
+    NixlPushConnectorWorker,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.scheduler import (
+    NixlConnectorScheduler,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import (
+    NixlKVConnectorStats,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
+    NixlConnectorWorker,
+)
+
+__all__ = [
+    "NixlAgentMetadata",
+    "NixlBaseConnector",
+    "NixlBaseConnectorScheduler",
+    "NixlBaseConnectorWorker",
+    "NixlConnector",
+    "NixlConnectorMetadata",
+    "NixlConnectorScheduler",
+    "NixlConnectorWorker",
+    "NixlHandshakePayload",
+    "NixlKVConnectorStats",
+    "NixlPullConnector",
+    "NixlPullConnectorScheduler",
+    "NixlPullConnectorWorker",
+    "NixlPushConnector",
+    "NixlPushConnectorScheduler",
+    "NixlPushConnectorWorker",
+]

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py

@@ -0,0 +1,461 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Base scheduler-side logic for the NIXL connector."""
+
+import threading
+import time
+from typing import TYPE_CHECKING, Any
+
+import msgspec
+import zmq
+
+from vllm import envs
+from vllm.distributed.kv_transfer.kv_connector.utils import (
+    BlockIds,
+    EngineId,
+    yield_req_data,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorHandshakeMetadata,
+    KVConnectorMetadata,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
+    GET_META_MSG,
+    HeartbeatInfo,
+    NixlConnectorMetadata,
+    NixlHandshakePayload,
+    ReqId,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import zmq_ctx
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+from vllm.utils.math_utils import cdiv
+from vllm.utils.network_utils import make_zmq_path
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    MambaSpec,
+    SlidingWindowSpec,
+)
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+    from vllm.v1.outputs import KVConnectorOutput
+    from vllm.v1.request import Request
+
+logger = init_logger(__name__)
+
+
+class NixlBaseConnectorScheduler:
+    """Base implementation of Scheduler side methods shared by pull and push."""
+
+    def __init__(
+        self,
+        vllm_config: "VllmConfig",
+        engine_id: str,
+        kv_cache_config: "KVCacheConfig",
+    ):
+        self.vllm_config = vllm_config
+        self.block_size = vllm_config.cache_config.block_size
+        self.engine_id: EngineId = engine_id
+        self.kv_cache_config = kv_cache_config
+        self.side_channel_host = envs.VLLM_NIXL_SIDE_CHANNEL_HOST
+        self.side_channel_port = (
+            envs.VLLM_NIXL_SIDE_CHANNEL_PORT
+            + vllm_config.parallel_config.data_parallel_index
+        )
+        assert vllm_config.kv_transfer_config is not None
+        self._kv_lease_duration: int = (
+            vllm_config.kv_transfer_config.get_from_extra_config(
+                "kv_lease_duration", 30
+            )
+        )
+        # NOTE (NickLucche): For now we use a hardcoded value for a simpler interface.
+        self._heartbeat_interval = self._kv_lease_duration // 6
+        if current_platform.device_type == "cpu":
+            self.use_host_buffer = False
+        else:
+            self.use_host_buffer = (
+                vllm_config.kv_transfer_config.kv_buffer_device == "cpu"
+            )
+        self._is_hma_required = (
+            not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
+            # Also handle unlikely SW-only model case instead of checking num_groups>1.
+            and any(
+                not isinstance(g.kv_cache_spec, FullAttentionSpec)
+                for g in kv_cache_config.kv_cache_groups
+            )
+        )
+        self._has_mamba = any(
+            isinstance(g.kv_cache_spec, MambaSpec)
+            for g in kv_cache_config.kv_cache_groups
+        )
+
+        logger.info("Initializing NIXL Scheduler %s", engine_id)
+        if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager:
+            logger.info("Hybrid Memory Allocator is enabled with NIXL")
+
+        # Background thread for handling new handshake requests.
+        self._nixl_handshake_listener_t: threading.Thread | None = None
+        self._stop_event = threading.Event()
+
+        # Requests that need to start recv/send.
+        # New requests are added by update_state_after_alloc in
+        # the scheduler. Used to make metadata passed to Worker.
+        self._reqs_need_recv: dict[ReqId, tuple[Request, BlockIds]] = {}
+        self._reqs_need_save: dict[ReqId, Request] = {}
+        # Reqs to send and their expiration time
+        self._reqs_need_send: dict[ReqId, float] = {}
+        self._reqs_in_batch: set[ReqId] = set()
+        # Reqs to remove from processed set because they're not to send after
+        # remote prefill or aborted.
+        self._reqs_not_processed: set[ReqId] = set()
+
+        # Heartbeat tracking: requests needing periodic lease-renewal heartbeats to
+        # remote P-side, stored as ready-to-send HeartbeatInfo grouped by remote engine
+        self._heartbeat_by_engine: dict[EngineId, HeartbeatInfo] = {}
+        # Reverse lookup: local req_id -> (engine_id, remote_req_id) for O(1) removal
+        self._heartbeat_req_engine: dict[ReqId, tuple[EngineId, ReqId]] = {}
+        self._last_heartbeat_time: float = 0.0
+
+        # Gather Sliding Window sizes for each kv cache group (if any) in number of
+        # blocks per KV cache group. This is used to clip the local attention window.
+        sw_sizes_tokens: list[tuple[int, int]] = [
+            (g.kv_cache_spec.sliding_window, g.kv_cache_spec.block_size)
+            if isinstance(g.kv_cache_spec, SlidingWindowSpec)
+            else (0, self.block_size)
+            for g in kv_cache_config.kv_cache_groups
+        ]
+        # cdiv(n_tokens, block_size) gives blocks/window; add 1 to conservatively
+        # account for boundary overlap eg window isn't fully aligned with blocks.
+        self.blocks_per_sw = [
+            cdiv(n_tokens, block_size) + 1 if n_tokens else 0
+            for n_tokens, block_size in sw_sizes_tokens
+        ]
+
+        # Threshold to decide whether to compute kv cache locally
+        # or pull from a remote node: minimum number of remote
+        # tokens to amortize the xfer latencies
+        self.kv_recompute_threshold: int = int(
+            vllm_config.kv_transfer_config.get_from_extra_config(
+                "kv_recompute_threshold", 64
+            )
+        )
+
+        # Bi-directional KV transfer feature supports KV block
+        # transfers from D node to P node
+        self.is_bidirectional_kv_xfer_enabled = (
+            vllm_config.kv_transfer_config.get_from_extra_config(
+                "bidirectional_kv_xfer", False
+            )
+        )
+        self.decoder_kv_blocks_ttl = (
+            vllm_config.kv_transfer_config.get_from_extra_config(
+                "decoder_kv_blocks_ttl", 480
+            )
+        )
+
+        if self.is_bidirectional_kv_xfer_enabled and self.kv_recompute_threshold > 0:
+            logger.info(
+                "Bidirectional KV transfer is enabled and the kv "
+                "recompute threshold is set to %d tokens."
+                "KV blocks on D are released after a TTL of %d seconds.",
+                self.kv_recompute_threshold,
+                self.decoder_kv_blocks_ttl,
+            )
+
+    def shutdown(self):
+        self._stop_event.set()
+        if self._nixl_handshake_listener_t is not None:
+            self._nixl_handshake_listener_t.join()
+            self._nixl_handshake_listener_t = None
+
+    def on_new_request(self, request: "Request") -> None:
+        """Track a request that may need heartbeats."""
+        params = request.kv_transfer_params
+        # NOTE (NickLucche) This excludes request meant for P, ie heartbeats are
+        # effectively disabled for Bidirectional KV transfer.
+        if params is None or not params.get("do_remote_prefill"):
+            return
+        # Only track if all required remote fields are present.
+        remote_engine_id = params.get("remote_engine_id")
+        remote_request_id = params.get("remote_request_id")
+        host = params.get("remote_host")
+        port = params.get("remote_port")
+        tp_size = params.get("tp_size")
+        pp_size = params.get("pp_size", 1)
+        if (
+            remote_engine_id is None
+            or remote_request_id is None
+            or host is None
+            or port is None
+            or tp_size is None
+        ):
+            return
+        if remote_engine_id not in self._heartbeat_by_engine:
+            self._heartbeat_by_engine[remote_engine_id] = HeartbeatInfo(
+                req_ids=set(),
+                host=host,
+                port=port,
+                tp_size=tp_size,
+                pp_size=pp_size,
+            )
+        self._heartbeat_by_engine[remote_engine_id].req_ids.add(remote_request_id)
+        self._heartbeat_req_engine[request.request_id] = (
+            remote_engine_id,
+            remote_request_id,
+        )
+
+    def _stop_heartbeat(self, req_id: ReqId) -> None:
+        """Remove *req_id* from heartbeat tracking (if tracked)."""
+        if key := self._heartbeat_req_engine.pop(req_id, None):
+            engine_id, remote_id = key
+            if info := self._heartbeat_by_engine.get(engine_id):
+                info.req_ids.discard(remote_id)
+                if not info.req_ids:
+                    # Clean up empty engines so we don't leak a key when remote dies.
+                    del self._heartbeat_by_engine[engine_id]
+
+    def get_sw_clipped_blocks(self, block_ids: BlockIds) -> BlockIds:
+        """
+        Clip the number of blocks to the sliding window size for each kv cache group
+        that employs SWA.
+        This is necessary because the KV Cache manager initially allocates blocks for
+        the entire sequence length, and successively cleans up blocks that are outside
+        the window prior to the `request_finished_all_groups` hook.
+        """
+        if len(block_ids) == 0 or not self._is_hma_required:
+            # No blocks to clip eg Full prefix cache hit or not a hybrid model.
+            return block_ids
+        # NOTE (NickLucche) This logic is currently handled at the connector level
+        # because offloading connectors might want to receive the whole sequence even
+        # for SWA groups. We will abstract this logic once the interface is more stable
+        assert len(block_ids) == len(self.blocks_per_sw), (
+            "Number of KV cache groups must match"
+        )
+        # For non-SWA groups, blocks_per_sw is 0 so we return all block_ids unchanged
+        return tuple(
+            [
+                blocks[-self.blocks_per_sw[i] :]
+                if self.blocks_per_sw[i] > 0
+                else blocks
+                for i, blocks in enumerate(block_ids)
+            ]
+        )
+
+    def set_xfer_handshake_metadata(
+        self, metadata: dict[tuple[int, int], KVConnectorHandshakeMetadata]
+    ) -> None:
+        """
+        Set the KV connector handshake metadata for this connector.
+
+        Args:
+            metadata (dict): the handshake metadata to set.
+        """
+        encoded_data: dict[tuple[int, int], bytes] = {}
+        encoder = msgspec.msgpack.Encoder()
+        for (pp_rank, tp_rank), rank_metadata in metadata.items():
+            if not isinstance(rank_metadata, NixlHandshakePayload):
+                raise ValueError(
+                    "NixlConnectorScheduler expects NixlHandshakePayload for "
+                    "handshake metadata."
+                )
+            encoded_data[(pp_rank, tp_rank)] = encoder.encode(rank_metadata)
+            logger.debug(
+                "PP rank %d, TP rank %d: encoded NixlHandshakePayload size: %s bytes",
+                pp_rank,
+                tp_rank,
+                str(len(encoded_data[(pp_rank, tp_rank)])),
+            )
+
+        # Only start the listener when we have metadata to serve.
+        if self._nixl_handshake_listener_t is None:
+            ready_event = threading.Event()
+            self._nixl_handshake_listener_t = threading.Thread(
+                target=self._nixl_handshake_listener,
+                args=(
+                    encoded_data,
+                    ready_event,
+                    self._stop_event,
+                    self.side_channel_host,
+                    self.side_channel_port,
+                ),
+                daemon=True,
+                name="nixl_handshake_listener",
+            )
+            self._nixl_handshake_listener_t.start()
+            ready_event.wait()  # Wait for listener ZMQ socket to be ready.
+
+    @staticmethod
+    def _nixl_handshake_listener(
+        encoded_data: dict[tuple[int, int], Any],
+        ready_event: threading.Event,
+        stop_event: threading.Event,
+        host: str,
+        port: int,
+    ):
+        """Background thread for getting new NIXL handshakes."""
+        # NOTE(rob): this is a simple implementation. We will move
+        # to a better approach via HTTP endpoint soon.
+
+        # Listen for new requests for metadata.
+        path = make_zmq_path("tcp", host, port)
+        logger.debug("Starting listening on path: %s", path)
+        with zmq_ctx(zmq.ROUTER, path) as sock:
+            sock.setsockopt(zmq.RCVTIMEO, 1000)
+            ready_event.set()
+            while True:
+                try:
+                    identity, _, msg = sock.recv_multipart()
+                except zmq.Again:
+                    if stop_event.is_set():
+                        break
+                    continue
+                # Decode (GET_META_MSG, pp_rank, tp_rank).
+                msg, target_pp_rank, target_tp_rank = msgspec.msgpack.decode(msg)
+                logger.debug(
+                    "Received message for pp rank %s, tp rank %s",
+                    target_pp_rank,
+                    target_tp_rank,
+                )
+                if msg != GET_META_MSG:
+                    logger.warning("Connection listener got unexpected message %s", msg)
+                sock.send_multipart(
+                    (identity, b"", encoded_data[(target_pp_rank, target_tp_rank)])
+                )
+
+    def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int:
+        """D-side only. Returns N-1 for Mamba models since the decoder
+        always recomputes the last token and must start from h(N-1)."""
+        if self._has_mamba and num_prompt_tokens > 1:
+            return num_prompt_tokens - 1
+        return num_prompt_tokens
+
+    def _truncate_mamba_request_for_prefill(self, request: "Request") -> None:
+        """P-side only: drop the last prompt token so the prefiller computes
+        h(N-1) instead of h(N). The decoder recomputes the last token to
+        derive h(N) correctly.
+
+        Guarded by ``_p_side_truncated`` to avoid repeated truncation if the
+        request is preempted and rescheduled."""
+        params = request.kv_transfer_params
+        if (
+            params is not None
+            # Guard against repeated truncation after preemption/reschedule.
+            and not params.get("_p_side_truncated")
+            and request.num_prompt_tokens > 1
+        ):
+            if request.prompt_token_ids is not None:
+                request.prompt_token_ids.pop()
+            elif request.prompt_embeds is not None:
+                request.prompt_embeds = request.prompt_embeds[:-1]
+            else:
+                return
+
+            request._all_token_ids.pop()
+            request.num_prompt_tokens -= 1
+            request.max_tokens = 1
+            params["_p_side_truncated"] = True
+
+    def _build_save_meta(
+        self,
+        meta: NixlConnectorMetadata,
+        scheduler_output: SchedulerOutput,
+    ) -> None:
+        # only called when use_host_buffer is True to build the save metadata
+
+        # NOTE: For the prefill side, there might be a chance that an early added
+        # request is a chunked prefill, so we need to check if new blocks are added
+        for req_id, new_block_id_groups, _ in yield_req_data(scheduler_output):
+            req_to_save = self._reqs_need_save.get(req_id)
+            if req_to_save is None or new_block_id_groups is None:
+                continue
+            req = req_to_save
+
+            assert req.kv_transfer_params is not None
+            clipped_block_id_groups = self.get_sw_clipped_blocks(new_block_id_groups)
+            meta.add_new_req_to_save(
+                request_id=req_id,
+                local_block_ids=clipped_block_id_groups,
+                kv_transfer_params=req.kv_transfer_params,
+            )
+            assert scheduler_output.num_scheduled_tokens is not None
+            num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
+            is_partial = (
+                req.num_computed_tokens + num_scheduled_tokens
+            ) < req.num_prompt_tokens
+            if not is_partial:
+                # For non-partial prefills, once new req_meta is scheduled, it
+                # can be removed from _reqs_need_save.
+                # For partial prefill case, we will retain the request in
+                # _reqs_need_save until all blocks are scheduled with req_meta.
+                # Therefore, only pop if `not is_partial`.
+                self._reqs_need_save.pop(req_id)
+
+    def build_connector_meta(
+        self,
+        scheduler_output: SchedulerOutput,
+    ) -> KVConnectorMetadata:
+        meta = NixlConnectorMetadata()
+
+        # Loop through scheduled reqs and convert to ReqMeta.
+        for req_id, (req, block_ids) in self._reqs_need_recv.items():
+            assert req.kv_transfer_params is not None
+            meta.add_new_req_to_recv(
+                request_id=req_id,
+                local_block_ids=block_ids,
+                kv_transfer_params=req.kv_transfer_params,
+            )
+
+        if self.use_host_buffer:
+            self._build_save_meta(meta, scheduler_output)
+
+        meta.reqs_to_send = self._reqs_need_send
+        meta.reqs_in_batch = self._reqs_in_batch
+        meta.reqs_not_processed = self._reqs_not_processed
+
+        # Package heartbeats, throttled by heartbeat_interval.
+        if self._heartbeat_by_engine:
+            now = time.perf_counter()
+            if now - self._last_heartbeat_time >= self._heartbeat_interval:
+                self._last_heartbeat_time = now
+                meta.heartbeat_by_engine = self._heartbeat_by_engine
+
+        # Clear the list once workers start the transfers
+        self._reqs_need_recv.clear()
+        self._reqs_in_batch = set()
+        self._reqs_not_processed = set()
+        self._reqs_need_send = {}
+
+        return meta
+
+    def update_connector_output(self, connector_output: "KVConnectorOutput") -> None:
+        """Stop heartbeating for requests whose KV transfer completed."""
+        for req_id in connector_output.finished_recving or ():
+            self._stop_heartbeat(req_id)
+
+    def has_pending_push_work(self) -> bool:
+        return False
+
+    ############################################################
+    # Abstract methods that subclasses must implement
+    ############################################################
+
+    def get_num_new_matched_tokens(
+        self, request: "Request", num_computed_tokens: int
+    ) -> tuple[int, bool]:
+        raise NotImplementedError
+
+    def update_state_after_alloc(
+        self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
+    ):
+        raise NotImplementedError
+
+    def request_finished(
+        self,
+        request: "Request",
+        block_ids: BlockIds,
+    ) -> tuple[bool, dict[str, Any] | None]:
+        raise NotImplementedError

vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py

@@ -0,0 +1,395 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NIXL connector facades.
+
+This module hosts the thin facade classes that vLLM's KV-connector layer
+instantiates. Almost all the real work lives in the per-mode scheduler
+and worker classes; the connector classes here only forward calls.
+
+* :class:`NixlBaseConnector` – common logic shared by pull and push.
+* :class:`NixlPullConnector` – pull-based (READ) KV transfer.
+* :class:`NixlPushConnector` – push-based (WRITE) KV transfer.
+* ``NixlConnector`` – backward-compatible alias for :class:`NixlPullConnector`.
+"""
+
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.distributed.kv_transfer.kv_connector.utils import (
+    EngineId,
+    get_current_attn_backend,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    CopyBlocksOp,
+    KVConnectorBase_V1,
+    KVConnectorHandshakeMetadata,
+    KVConnectorMetadata,
+    KVConnectorRole,
+    SupportsHMA,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
+    KVConnectorPromMetrics,
+    KVConnectorStats,
+    PromMetric,
+    PromMetricT,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
+    NixlConnectorMetadata,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import (
+    NixlPullConnectorScheduler,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_worker import (
+    NixlPullConnectorWorker,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_scheduler import (
+    NixlPushConnectorScheduler,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import (
+    NixlPushConnectorWorker,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import (
+    NixlKVConnectorStats,
+    NixlPromMetrics,
+)
+from vllm.forward_context import ForwardContext
+from vllm.logger import init_logger
+from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
+from vllm.v1.attention.backends.utils import get_kv_cache_layout
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.kv_cache_interface import MambaSpec
+from vllm.v1.outputs import KVConnectorOutput
+
+if TYPE_CHECKING:
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_scheduler import (
+        NixlBaseConnectorScheduler,
+    )
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import (
+        NixlBaseConnectorWorker,
+    )
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+    from vllm.v1.request import Request
+
+logger = init_logger(__name__)
+
+
+class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):
+    """Base connector with common logic shared by pull and push modes."""
+
+    @property
+    def prefer_cross_layer_blocks(self) -> bool:
+        if any(
+            [
+                isinstance(group.kv_cache_spec, MambaSpec)
+                for group in self.kv_cache_config.kv_cache_groups
+            ]
+        ):
+            # Hybrid SSM models do not yet support cross-layer layout
+            return False
+
+        backend = get_current_attn_backend(self._vllm_config)
+        if backend.get_name() not in (
+            "FLASH_ATTN",
+            "FLASHINFER",
+            "TRITON_ATTN",
+        ):
+            return False
+
+        # For now there is no benefit to run cross layers when backend
+        # does not support on HND
+        if get_kv_cache_layout() != "HND":
+            return False
+
+        extra_config = self.kv_transfer_config.kv_connector_extra_config
+        return (
+            str(extra_config.get("enable_cross_layers_blocks", "False")).lower()
+            == "true"
+        )
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        role: KVConnectorRole,
+        kv_cache_config: "KVCacheConfig",
+    ):
+        super().__init__(vllm_config, role, kv_cache_config)
+        assert vllm_config.kv_transfer_config is not None
+        assert vllm_config.kv_transfer_config.engine_id is not None
+
+        if vllm_config.kv_transfer_config.kv_role == "kv_both":
+            logger.warning_once(
+                "Using kv_role='kv_both' with NixlConnector is deprecated "
+                "and will be removed in a future release. Please set "
+                "kv_role='kv_producer' for prefill instances and "
+                "kv_role='kv_consumer' for decode instances. "
+            )
+
+        self.kv_cache_config = kv_cache_config
+        self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id
+        self.kv_transfer_config = vllm_config.kv_transfer_config
+        # Subclasses must set self.connector_scheduler and self.connector_worker
+        self.connector_scheduler: NixlBaseConnectorScheduler | None = None
+        self.connector_worker: NixlBaseConnectorWorker | None = None
+
+    ############################################################
+    # Class Methods
+    ############################################################
+    @classmethod
+    def get_required_kvcache_layout(cls, vllm_config: VllmConfig):
+        if vllm_config.model_config is None:
+            logger.warning_once(
+                "Unable to detect current VLLM config. "
+                "Fallback to default kv cache layout."
+            )
+            return None
+        use_mla = vllm_config.model_config.use_mla
+        if use_mla:
+            # return None when we have mla
+            # as the layout should not matter in that case,
+            # which fallback to the default behavior.
+            return None
+        logger.info_once(
+            "NixlConnector setting KV cache layout to HND for better xfer performance."
+        )
+        return "HND"
+
+    ############################################################
+    # Scheduler Side Methods
+    ############################################################
+
+    def get_num_new_matched_tokens(
+        self, request: "Request", num_computed_tokens: int
+    ) -> tuple[int | None, bool]:
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.get_num_new_matched_tokens(
+            request, num_computed_tokens
+        )
+
+    def update_state_after_alloc(
+        self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
+    ):
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.update_state_after_alloc(
+            request, blocks, num_external_tokens
+        )
+
+    def build_connector_meta(
+        self,
+        scheduler_output: SchedulerOutput,
+    ) -> KVConnectorMetadata:
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.build_connector_meta(scheduler_output)
+
+    def on_new_request(self, request: "Request") -> None:
+        assert self.connector_scheduler is not None
+        self.connector_scheduler.on_new_request(request)
+
+    def update_connector_output(self, connector_output: KVConnectorOutput):
+        assert self.connector_scheduler is not None
+        self.connector_scheduler.update_connector_output(connector_output)
+
+    def request_finished(
+        self,
+        request: "Request",
+        block_ids: list[int],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.request_finished(request, (block_ids,))
+
+    def request_finished_all_groups(
+        self,
+        request: "Request",
+        block_ids: tuple[list[int], ...],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.request_finished(request, block_ids)
+
+    def set_xfer_handshake_metadata_pp_aware(
+        self, metadata: dict[tuple[int, int], KVConnectorHandshakeMetadata]
+    ) -> None:
+        """
+        Set handshake metadata keyed by (pp_rank, tp_rank) so the side
+        channel can serve every PP stage's agent metadata.
+
+        Args:
+            metadata (dict): the handshake metadata to set.
+        """
+        assert self.connector_scheduler is not None
+        self.connector_scheduler.set_xfer_handshake_metadata(metadata)
+
+    ############################################################
+    # Worker Side Methods
+    ############################################################
+    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
+        assert self.connector_worker is not None
+        self.connector_worker.register_kv_caches(kv_caches)
+
+    def register_cross_layers_kv_cache(
+        self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
+    ):
+        assert self.connector_worker is not None
+        self.connector_worker.register_cross_layers_kv_caches(kv_cache)
+
+    def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
+        assert self.connector_worker is not None
+        self.connector_worker.set_host_xfer_buffer_ops(copy_operation)
+
+    def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
+        """Get the finished recving and sending requests."""
+        assert self.connector_worker is not None
+        return self.connector_worker.get_finished()
+
+    def get_block_ids_with_load_errors(self) -> set[int]:
+        """Get block IDs that failed to load via NIXL."""
+        assert self.connector_worker is not None
+        return self.connector_worker.get_block_ids_with_load_errors()
+
+    def get_kv_connector_stats(self) -> KVConnectorStats | None:
+        if self.connector_worker is None:
+            return None
+        return self.connector_worker.get_kv_connector_stats()
+
+    @classmethod
+    def build_kv_connector_stats(
+        cls, data: dict[str, Any] | None = None
+    ) -> KVConnectorStats | None:
+        return (
+            NixlKVConnectorStats(data=data)
+            if data is not None
+            else NixlKVConnectorStats()
+        )
+
+    @classmethod
+    def build_prom_metrics(
+        cls,
+        vllm_config: VllmConfig,
+        metric_types: dict[type[PromMetric], type[PromMetricT]],
+        labelnames: list[str],
+        per_engine_labelvalues: dict[int, list[object]],
+    ) -> KVConnectorPromMetrics:
+        return NixlPromMetrics(
+            vllm_config, metric_types, labelnames, per_engine_labelvalues
+        )
+
+    def wait_for_layer_load(self, layer_name: str) -> None:
+        """NixlConnector does not do layerwise saving."""
+        pass
+
+    def save_kv_layer(
+        self,
+        layer_name: str,
+        kv_layer: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        **kwargs,
+    ) -> None:
+        """NixlConnector does not save explicitly."""
+        pass
+
+    def wait_for_save(self):
+        assert self.connector_worker is not None
+        assert isinstance(self._connector_metadata, NixlConnectorMetadata)
+        if self.connector_worker.use_host_buffer and self.connector_worker.copy_blocks:
+            self.connector_worker.save_kv_to_host(self._connector_metadata)
+
+    def has_pending_push_work(self) -> bool:
+        if self.connector_scheduler is not None:
+            return self.connector_scheduler.has_pending_push_work()
+        return False
+
+    def shutdown(self):
+        if self.connector_worker is not None:
+            self.connector_worker.shutdown()
+        if self.connector_scheduler is not None:
+            self.connector_scheduler.shutdown()
+
+    def get_handshake_metadata(self) -> KVConnectorHandshakeMetadata | None:
+        """
+        Get the KVConnector handshake metadata for this connector.
+        This metadata is used for out-of-band connector handshake
+        between P/D workers.
+
+        Returns:
+            KVConnectorHandshakeMetadata: the handshake metadata.
+            None if no handshake metadata is available.
+        """
+        assert self.connector_worker is not None
+        return self.connector_worker.xfer_handshake_metadata
+
+
+class NixlPullConnector(NixlBaseConnector):
+    """Pull-based (READ) NIXL KV transfer connector."""
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        role: KVConnectorRole,
+        kv_cache_config: "KVCacheConfig",
+    ):
+        super().__init__(vllm_config, role, kv_cache_config)
+        if role == KVConnectorRole.SCHEDULER:
+            self.connector_scheduler = NixlPullConnectorScheduler(
+                vllm_config, self.engine_id, kv_cache_config
+            )
+            self.connector_worker = None
+        elif role == KVConnectorRole.WORKER:
+            self.connector_scheduler = None
+            self.connector_worker = NixlPullConnectorWorker(
+                vllm_config, self.engine_id, kv_cache_config
+            )
+
+    def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
+        assert self.connector_worker is not None
+        assert isinstance(self.connector_worker, NixlPullConnectorWorker)
+        assert isinstance(self._connector_metadata, NixlConnectorMetadata)
+        self.connector_worker.start_load_kv(self._connector_metadata)
+
+
+class NixlPushConnector(NixlBaseConnector):
+    """Push-based (WRITE) NIXL KV transfer connector."""
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        role: KVConnectorRole,
+        kv_cache_config: "KVCacheConfig",
+    ):
+        super().__init__(vllm_config, role, kv_cache_config)
+        self.connector_scheduler: NixlPushConnectorScheduler | None = None
+        self.connector_worker: NixlPushConnectorWorker | None = None
+        if role == KVConnectorRole.SCHEDULER:
+            self.connector_scheduler = NixlPushConnectorScheduler(
+                vllm_config, self.engine_id, kv_cache_config
+            )
+        elif role == KVConnectorRole.WORKER:
+            self.connector_worker = NixlPushConnectorWorker(
+                vllm_config, self.engine_id, kv_cache_config
+            )
+        else:
+            raise ValueError(f"Unsupported KVConnectorRole: {role}")
+
+    def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
+        """Drive push processing on the worker.
+
+        The worker enqueues registrations / finished blocks for the
+        background ``nixl-push-writer`` thread; the writer issues the
+        WRITE transfers and polls NIXL notifs without further
+        engine-thread involvement.
+        """
+        assert self.connector_worker is not None
+        assert isinstance(self._connector_metadata, NixlConnectorMetadata)
+        self.connector_worker.start_load_kv(self._connector_metadata)
+
+
+# Backward compatibility: NixlConnector is the pull-based connector.
+NixlConnector = NixlPullConnector
+
+
+__all__ = [
+    "NixlBaseConnector",
+    "NixlConnector",
+    "NixlPullConnector",
+    "NixlPushConnector",
+]

vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py

@@ -0,0 +1,229 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Metadata dataclasses and helpers for the NIXL connector."""
+
+from dataclasses import dataclass
+from typing import Any
+
+from vllm.config import VllmConfig
+from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds, EngineId
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorHandshakeMetadata,
+    KVConnectorMetadata,
+)
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+TransferHandle = int
+ReqId = str
+
+GET_META_MSG = b"get_meta_msg"
+
+# Push-mode (WRITE-based) registration notification.
+# Sent worker-to-worker over NIXL: D worker -> P worker, encoded as
+# PUSH_REG_NOTIF_PREFIX + msgpack(registration_data).
+PUSH_REG_NOTIF_PREFIX = b"PUSH_REG:"
+#
+# NIXL Connector Version
+#
+# Increment this version whenever there is an incompatible change to:
+#   - NixlAgentMetadata schema
+#   - kv_transfer_params schema or semantics
+#   - NIXL transfer protocol or wire format
+#   - KV cache memory layout or block organization
+#   - Any other change that breaks P/D interoperability
+#
+# Version History:
+#   1: Initial version with compatibility checking
+#   2: Add remote_request_id to kv_transfer_params
+#   3: Add physical_blocks_per_logical_kv_block to NixlAgentMetadata
+#   4: Add KV block lease renewal through heartbeats
+#
+NIXL_CONNECTOR_VERSION: int = 4
+
+
+@dataclass
+class NixlAgentMetadata:
+    engine_id: str
+    agent_metadata: bytes
+    kv_caches_base_addr: list[int]
+    device_id: int
+    num_blocks: int
+    block_lens: list[int]
+    kv_cache_layout: str
+    block_size: int
+    ssm_sizes: tuple[int, int]
+    attn_backend_name: str
+    physical_blocks_per_logical_kv_block: int
+
+
+@dataclass
+class NixlHandshakePayload(KVConnectorHandshakeMetadata):
+    """
+    Wrapper for NIXL handshake sent over the wire.
+
+    Enables two-phase decoding for graceful compatibility checking:
+    1. Decode NixlHandshakePayload to get compatibility_hash
+    2. Compute local hash and compare
+    3. Only if hashes match, decode agent_metadata_bytes
+
+    This prevents decoder errors when NixlAgentMetadata schema is
+    incompatible, allowing graceful failure with clear error message.
+    """
+
+    compatibility_hash: str
+    agent_metadata_bytes: bytes  # NixlAgentMetadata encoded
+
+
+def compute_nixl_compatibility_hash(
+    vllm_config: VllmConfig, attn_backend_name: str, cross_layers_blocks: bool
+) -> str:
+    """
+    Compute compatibility hash for NIXL KV transfer.
+
+    Hash only the factors that affect whether two NIXL instances can
+    successfully transfer KV cache data.
+
+    Factors included:
+    - vLLM version and NIXL connector version
+    - Model architecture (name, dtype, KV heads, layers)
+    - KV cache format (dtype, sliding window)
+    - Attention backend
+
+    Note: Factors like tensor_parallel_size, block_size, and kv_cache_layout
+    are validated at runtime in _validate_remote_agent_handshake and are not
+    included in this hash to support heterogeneous deployments.
+
+    Note - the set of factors are likely to evolve significantly over
+    time to be more or less permissive.
+
+    Returns:
+        SHA-256 hex digest
+    """
+    from vllm import __version__ as vllm_version
+    from vllm.config.utils import hash_factors
+
+    model_config = vllm_config.model_config
+    cache_config = vllm_config.cache_config
+    is_hma_enabled = not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
+
+    factors = {
+        # Version compatibility
+        "vllm_version": vllm_version,
+        "nixl_connector_version": NIXL_CONNECTOR_VERSION,
+        # Model architecture - affects KV cache shape
+        "model": model_config.model,
+        "dtype": str(model_config.dtype),
+        "num_kv_heads": model_config.get_total_num_kv_heads(),
+        "head_size": model_config.get_head_size(),
+        "num_hidden_layers": model_config.get_total_num_hidden_layers(),
+        # Attention backend and KV cache dtype affect memory layout
+        "attn_backend_name": attn_backend_name,
+        "cache_dtype": str(cache_config.cache_dtype),
+        "cross_layers_blocks": cross_layers_blocks,
+        "is_hma_enabled": is_hma_enabled,
+    }
+
+    compat_hash = hash_factors(factors)
+    logger.debug(
+        "NIXL compatibility hash: %s (model=%s, dtype=%s, num_kv_heads=%d, "
+        "cache_dtype=%s, attn_backend=%s)",
+        compat_hash,
+        factors["model"],
+        factors["dtype"],
+        factors["num_kv_heads"],
+        factors["cache_dtype"],
+        attn_backend_name,
+    )
+    return compat_hash
+
+
+@dataclass
+class HeartbeatInfo:
+    """Heartbeat data for a single remote engine, sent from D worker to P."""
+
+    req_ids: set[ReqId]
+    host: str
+    port: int
+    tp_size: int
+    pp_size: int = 1
+
+
+@dataclass
+class RemoteMeta:
+    block_ids: BlockIds
+    host: str
+    port: int
+    engine_id: str
+    request_id: str
+
+
+@dataclass
+class ReqMeta:
+    local_block_ids: BlockIds
+    # To be used when logical block size does not match the kernel block size
+    local_physical_block_ids: BlockIds
+    tp_size: int
+    remote: RemoteMeta | None = None
+    # Remote block size, discovered during NIXL handshake (push mode).
+    remote_block_size: int | None = None
+    # Remote producer pipeline-parallel size (push mode, D side).
+    pp_size: int = 1
+
+
+class NixlConnectorMetadata(KVConnectorMetadata):
+    def __init__(self):
+        self.reqs_to_recv: dict[ReqId, ReqMeta] = {}
+        self.reqs_to_save: dict[ReqId, ReqMeta] = {}
+        self.reqs_to_send: dict[ReqId, float] = {}
+        self.reqs_in_batch: set[ReqId] = set()
+        self.reqs_not_processed: set[ReqId] = set()
+        # Heartbeat data grouped by remote engine, sent by D worker to P.
+        self.heartbeat_by_engine: dict[EngineId, HeartbeatInfo] = {}
+        # Push mode (D side): registration data the D worker should send to
+        # P workers via NIXL notification on this step.
+        self.push_registrations: dict[ReqId, dict[str, Any]] = {}
+        # Push mode (P side): newly finished request blocks to be matched
+        # against pending D registrations on the P worker.
+        self.push_finished_blocks: dict[ReqId, BlockIds] = {}
+
+    def _add_new_req(
+        self,
+        local_block_ids: BlockIds,
+        kv_transfer_params: dict[str, Any],
+    ) -> ReqMeta:
+        return ReqMeta(
+            local_block_ids=local_block_ids,
+            local_physical_block_ids=local_block_ids,
+            # P workers don't need to receive tp_size from proxy here.
+            tp_size=kv_transfer_params.get("tp_size", 1),
+            remote_block_size=kv_transfer_params.get("remote_block_size"),
+            pp_size=kv_transfer_params.get("pp_size", 1),
+        )
+
+    def add_new_req_to_save(
+        self,
+        request_id: ReqId,
+        local_block_ids: BlockIds,
+        kv_transfer_params: dict[str, Any],
+    ):
+        self.reqs_to_save[request_id] = self._add_new_req(
+            local_block_ids, kv_transfer_params
+        )
+
+    def add_new_req_to_recv(
+        self,
+        request_id: ReqId,
+        local_block_ids: BlockIds,
+        kv_transfer_params: dict[str, Any],
+    ):
+        req = self._add_new_req(local_block_ids, kv_transfer_params)
+        req.remote = RemoteMeta(
+            block_ids=kv_transfer_params["remote_block_ids"],
+            engine_id=kv_transfer_params["remote_engine_id"],
+            request_id=kv_transfer_params["remote_request_id"],
+            host=kv_transfer_params["remote_host"],
+            port=kv_transfer_params["remote_port"],
+        )
+        self.reqs_to_recv[request_id] = req

vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py

@@ -0,0 +1,275 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Pull-specific scheduler-side logic for the NIXL connector."""
+
+import time
+from typing import TYPE_CHECKING, Any
+
+from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_scheduler import (
+    NixlBaseConnectorScheduler,
+)
+from vllm.logger import init_logger
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+    from vllm.v1.request import Request
+
+logger = init_logger(__name__)
+
+
+class NixlPullConnectorScheduler(NixlBaseConnectorScheduler):
+    """Pull-specific scheduler logic (READ-based KV transfer)."""
+
+    def __init__(
+        self,
+        vllm_config: "VllmConfig",
+        engine_id: str,
+        kv_cache_config: "KVCacheConfig",
+    ):
+        super().__init__(vllm_config, engine_id, kv_cache_config)
+
+    def get_num_new_matched_tokens(
+        self, request: "Request", num_computed_tokens: int
+    ) -> tuple[int, bool]:
+        """
+        For remote prefill, pull all prompt blocks from remote
+        asynchronously relative to engine execution.
+
+        Args:
+            request (Request): the request object.
+            num_computed_tokens (int): the number of locally
+                computed tokens for this request
+        Returns:
+            * the number of tokens that can be loaded from the
+              external KV cache beyond what is already computed.
+            * true if the external KV cache tokens will be loaded
+              asynchronously (between scheduler steps).
+        """
+
+        params = request.kv_transfer_params
+        logger.debug(
+            "NIXLConnector get_num_new_matched_tokens: "
+            "num_computed_tokens=%s, kv_transfer_params=%s",
+            num_computed_tokens,
+            params,
+        )
+
+        if params is not None and params.get("do_remote_prefill"):
+            # Remote prefill: get all prompt blocks from remote.
+            token_ids = request.prompt_token_ids or []
+            actual = self._get_remote_prefill_token_count(len(token_ids))
+            count = actual - num_computed_tokens
+            if count > 0:
+                return count, True
+
+        if params is not None and params.get("do_remote_decode") and self._has_mamba:
+            self._truncate_mamba_request_for_prefill(request)
+
+        if (
+            params is not None
+            and params.get("do_remote_decode")
+            and params.get("remote_block_ids")
+            and all(
+                p in params
+                for p in (
+                    "remote_engine_id",
+                    "remote_request_id",
+                    "remote_host",
+                    "remote_port",
+                )
+            )
+        ):
+            # Decode node has kv blocks for part of prefill request, so, provide them
+            # as an external token count to scheduler.
+            # The tokens will be loaded if not already present
+            # in the prefill node local cache
+            remote_num_tokens = params.get("remote_num_tokens") or 0
+            count = (
+                min(remote_num_tokens, request.num_prompt_tokens) - num_computed_tokens
+            )
+            if count > 0:
+                # Check kv_recompute_threshold: skip pull if
+                # remote tokens are below the threshold.
+                if (
+                    self.kv_recompute_threshold > 0
+                    and count < self.kv_recompute_threshold
+                ):
+                    logger.debug(
+                        "Skipping remote pull for %s: %d remote tokens < threshold %d",
+                        request.request_id,
+                        count,
+                        self.kv_recompute_threshold,
+                    )
+                    return 0, False
+                return count, True
+
+        # No remote prefill for this request.
+        return 0, False
+
+    def update_state_after_alloc(
+        self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
+    ):
+        params = request.kv_transfer_params
+        logger.debug(
+            "NIXLConnector update_state_after_alloc: "
+            "num_external_tokens=%s, kv_transfer_params=%s",
+            num_external_tokens,
+            params,
+        )
+
+        if not params:
+            return
+
+        if params.get("do_remote_decode") or (
+            params.get("do_remote_prefill") and self.is_bidirectional_kv_xfer_enabled
+        ):
+            self._reqs_in_batch.add(request.request_id)
+        if self.use_host_buffer and params.get("do_remote_decode"):
+            # NOTE: when accelerator is not directly supported by Nixl,
+            # prefilled blocks need to be saved to host memory before transfer.
+            self._reqs_need_save[request.request_id] = request
+        elif params.get("do_remote_prefill") or (
+            params.get("do_remote_decode")
+            and self.is_bidirectional_kv_xfer_enabled
+            and not params.get("_remote_blocks_processed")
+        ):
+            if params.get("remote_block_ids"):
+                if all(
+                    p in params
+                    for p in (
+                        "remote_engine_id",
+                        "remote_request_id",
+                        "remote_host",
+                        "remote_port",
+                    )
+                ):
+                    # If remote_blocks and num_external_tokens = 0, we have
+                    # a full prefix cache hit on the local node. We need to call
+                    # send_notif in _read_blocks to free the memory on the remote node.
+
+                    unhashed_local_block_ids: BlockIds = (
+                        blocks.get_unhashed_block_ids_all_groups()
+                        if num_external_tokens > 0
+                        else ()
+                    )
+                    local_block_ids = self.get_sw_clipped_blocks(
+                        unhashed_local_block_ids
+                    )
+
+                    # Get unhashed blocks to pull from remote. Mind that a full prefix
+                    # cache hit is indicated with an empty list.
+                    self._reqs_need_recv[request.request_id] = (
+                        request,
+                        local_block_ids,
+                    )
+
+                else:
+                    logger.warning(
+                        "Got invalid KVTransferParams: %s. This "
+                        "request will not utilize KVTransfer",
+                        params,
+                    )
+            else:
+                assert num_external_tokens == 0
+            # Only trigger 1 KV transfer per request.
+            params["do_remote_prefill"] = False
+            params["_remote_blocks_processed"] = True
+
+    def request_finished(
+        self,
+        request: "Request",
+        block_ids: "BlockIds",
+    ) -> tuple[bool, dict[str, Any] | None]:
+        """
+        Once a request is finished, determine whether request blocks
+        should be freed now or will be sent asynchronously and freed later.
+        """
+        from vllm.v1.request import RequestStatus
+
+        params = request.kv_transfer_params
+        logger.debug(
+            "NIXLConnector request_finished(%s), request_status=%s, "
+            "kv_transfer_params=%s",
+            request.request_id,
+            request.status,
+            params,
+        )
+        if not params:
+            return False, None
+
+        is_p_node = bool(params.get("do_remote_decode"))
+        is_d_node = not is_p_node
+
+        # Stop heartbeating for aborted requests that never reached finished_recving:
+        # normal path cleans up in update_connector_output.
+        self._stop_heartbeat(request.request_id)
+
+        if params.get("do_remote_prefill"):
+            # If do_remote_prefill is still True when the request is finished,
+            # update_state_after_alloc must not have been called (the request
+            # must have been aborted before it was scheduled, e.g. via the
+            # abort_immediately path used to clean up KV-transfer requests
+            # rejected at the D-side serving layer).
+            # To avoid stranding the prefill blocks in the prefill instance,
+            # we must add empty block_ids to _reqs_need_recv so that our
+            # worker side will notify and free blocks in the prefill instance.
+            self._reqs_need_recv[request.request_id] = (request, [])
+            params["do_remote_prefill"] = False
+            return False, None
+
+        if is_d_node and not self.is_bidirectional_kv_xfer_enabled:
+            return False, None
+
+        if request.status not in (
+            RequestStatus.FINISHED_LENGTH_CAPPED,
+            RequestStatus.FINISHED_STOPPED,
+        ):
+            # Also include the case of a P/D Prefill request with immediate
+            # block free (eg abort). Stop tracking this request.
+            self._reqs_not_processed.add(request.request_id)
+            # Clear _reqs_need_save if a request is aborted as partial prefill.
+            self._reqs_need_save.pop(request.request_id, None)
+            return False, None
+
+        # TODO: check whether block_ids actually ever be 0. If not we could
+        # remove the conditional below
+        delay_free_blocks = any(len(group) > 0 for group in block_ids)
+        remote_num_tokens = 0
+        if delay_free_blocks:
+            # Prefill request on remote. It will be read from D upon completion
+            request_kv_blocks_ttl = self._kv_lease_duration
+            if is_d_node:
+                # For blocks pinned on D, use a simpler timeout for now instead of a
+                # lease mechanism as turn2 request is client-driven.
+                request_kv_blocks_ttl = self.decoder_kv_blocks_ttl
+            logger.debug(
+                "NIXLConnector request_finished(%s) waiting for %d seconds "
+                "before releasing blocks",
+                request.request_id,
+                request_kv_blocks_ttl,
+            )
+            self._reqs_need_send[request.request_id] = (
+                time.perf_counter() + request_kv_blocks_ttl
+            )
+            # NOTE HMA will "mark" empty/null blocks in groups with 0s (eg SWA ones),
+            # trimming down after allocating for the whole sequence length. Empty
+            # blocks are always at the start of the list.
+            # Here we "unpad" blocks to send the actual remote blocks to be read.
+            block_ids = self.get_sw_clipped_blocks(block_ids)
+
+            remote_num_tokens = request.num_computed_tokens
+
+        return delay_free_blocks, dict(
+            do_remote_prefill=is_p_node,
+            do_remote_decode=is_d_node,
+            remote_block_ids=block_ids,
+            remote_engine_id=self.engine_id,
+            remote_request_id=request.request_id,
+            remote_host=self.side_channel_host,
+            remote_port=self.side_channel_port,
+            tp_size=self.vllm_config.parallel_config.tensor_parallel_size,
+            remote_num_tokens=remote_num_tokens,
+        )

vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py

@@ -0,0 +1,382 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Pull-specific (READ) worker-side logic for the NIXL connector."""
+
+import time
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import (
+    NixlBaseConnectorWorker,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
+    NixlConnectorMetadata,
+    ReqMeta,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (
+    ReadSpec,
+)
+from vllm.logger import init_logger
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+
+logger = init_logger(__name__)
+
+
+class NixlPullConnectorWorker(NixlBaseConnectorWorker):
+    """Pull-specific (READ) worker logic."""
+
+    def __init__(
+        self,
+        vllm_config: "VllmConfig",
+        engine_id: str,
+        kv_cache_config: "KVCacheConfig",
+    ):
+        super().__init__(vllm_config, engine_id, kv_cache_config)
+
+    def start_load_kv(self, metadata: NixlConnectorMetadata):
+        """
+        Start loading by triggering non-blocking nixl_xfer.
+        We check for these trnxs to complete in each step().
+        """
+        for req_id, meta in metadata.reqs_to_recv.items():
+            meta.local_physical_block_ids = self._logical_to_kernel_block_ids(
+                meta.local_block_ids
+            )
+            assert meta.remote is not None
+            # Remote block IDs are kept logical here; expanded in
+            # _read_blocks_for_req using the remote engine's phys ratio.
+            remote_engine_id = meta.remote.engine_id
+            logger.debug(
+                "start_load_kv for request %s from remote engine %s. "
+                "Num local_block_ids: %s. Num remote_block_ids: %s. ",
+                req_id,
+                remote_engine_id,
+                len(meta.local_physical_block_ids),
+                len(meta.remote.block_ids),
+            )
+            # always store metadata for failure recovery
+            self._recving_metadata[req_id] = meta
+            if remote_engine_id not in self._remote_agents:
+                # Initiate handshake with remote engine to exchange metadata.
+                with self._handshake_lock:
+                    if remote_engine_id not in self._remote_agents:
+                        self._background_nixl_handshake(req_id, remote_engine_id, meta)
+                        continue
+
+            # Handshake already completed, start async read xfer.
+            self._read_blocks_for_req(req_id, meta)
+
+        # Start transfers for requests whose handshakes have now finished.
+        while not self._ready_requests.empty():
+            self._read_blocks_for_req(*self._ready_requests.get_nowait())
+
+        # Keep around the requests that have been part of a batch. This is
+        # needed because async scheduling pushes the misalignment between the
+        # moment in which requests expiration is set (P side) and the moment in
+        # which blocks are read from D. As P can now more easily lag behind D
+        # while processing the next batch, we make sure to only set an
+        # expiration for requests that have not been read from D yet.
+        for req_id in metadata.reqs_in_batch:
+            self._reqs_to_process.add(req_id)
+
+        # Remove all requests that are not to be processed (eg aborted).
+        for req_id in metadata.reqs_not_processed:
+            self._reqs_to_process.discard(req_id)
+            # We should never get an abort after setting an expiry timer
+            assert req_id not in self._reqs_to_send
+
+        # Add to requests that are waiting to be read and track expiration.
+        for req_id, expiration_time in metadata.reqs_to_send.items():
+            if req_id in self._reqs_to_process:
+                self._reqs_to_send[req_id] = expiration_time
+
+        # Send heartbeats to P-side engines to keep KV blocks alive while
+        # requests sit in the D scheduler WAITING queue.
+        self._send_heartbeats(metadata)
+
+    def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
+        assert meta.remote is not None and self.transfer_topo is not None
+        engine_id = meta.remote.engine_id
+        # Update last activity from this remote. Mind that cleanup is done on main
+        # thread (this one), so we don't race on this structure.
+        self._engine_last_active[engine_id] = time.perf_counter()
+        plan = self.tp_mappings[engine_id]
+        remote_info = self.transfer_topo.get_engine_info(engine_id)
+        tp_ratio = self.transfer_topo.tp_ratio(remote_info.remote_tp_size)
+
+        meta.remote.block_ids = self._logical_to_remote_kernel_block_ids(
+            meta.remote.block_ids,
+            remote_info.remote_physical_blocks_per_logical,
+        )
+        remote_block_ids = meta.remote.block_ids
+        local_block_ids = meta.local_physical_block_ids
+        num_groups = len(local_block_ids)
+        read_specs = [
+            ReadSpec(
+                remote_rank=rank,
+                local_block_ids=[
+                    list(local_block_ids[g])
+                    if rank in plan.source_ranks_per_group[g]
+                    else []
+                    for g in range(num_groups)
+                ],
+                remote_block_ids=[
+                    list(remote_block_ids[g])
+                    if rank in plan.source_ranks_per_group[g]
+                    else []
+                    for g in range(num_groups)
+                ],
+            )
+            for rank in plan.all_source_ranks
+        ]
+
+        # D may have to perform multiple reads from different remote ranks.
+        # MLA opt: when P TP > D TP, only a single read is executed for
+        # the first remote rank (cache is duplicated)..
+        if self.use_mla and tp_ratio < 0:
+            assert len(read_specs) == 1
+
+        for i, spec in enumerate(read_specs):
+            remote_block_size = remote_info.remote_block_size
+            logger.debug(
+                "Remote agent %s available, calling _read_blocks"
+                " on remote rank %s with remote block size %s for req %s",
+                meta.remote.engine_id,
+                spec.remote_rank,
+                remote_block_size,
+                req_id,
+            )
+            # Get side handles.
+            if tp_ratio < 0 and not self.use_mla:
+                assert remote_block_size == self.block_size
+                # Remote tp_size > local tp_size: we must perform multiple
+                # reads. Get the memory chunk onto which we will write to.
+                local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[tp_ratio][i]
+            else:
+                # Single read from remote, we write to the whole memory region.
+                # Also handle remote block size different from local block size.
+                local_xfer_side_handle = self.src_xfer_handles_by_block_size[
+                    remote_block_size
+                ]
+
+            # Destination handle: remote_engine_id -> remote_rank -> handle.
+            remote_xfer_side_handle = self.dst_xfer_side_handles[meta.remote.engine_id][
+                spec.remote_rank
+            ]
+
+            self._read_blocks(
+                read_spec=spec,
+                request_id=req_id,
+                dst_engine_id=meta.remote.engine_id,
+                remote_request_id=meta.remote.request_id,
+                local_xfer_side_handle=local_xfer_side_handle,
+                remote_xfer_side_handle=remote_xfer_side_handle,
+            )
+
+        if self.use_mla and tp_ratio < 0 and read_specs:
+            # ..but we still need to notify the other remote ranks that we
+            # have the blocks we need so they can update the request state.
+            notif_id = f"{meta.remote.request_id}:{self.world_size}".encode()
+            remote_agents = self._remote_agents[meta.remote.engine_id]
+            for rank_to_notify, agent in remote_agents.items():
+                if rank_to_notify != (0, read_specs[0].remote_rank):
+                    self.nixl_wrapper.send_notif(agent, notif_msg=notif_id)
+
+    def _read_blocks(
+        self,
+        read_spec: ReadSpec,
+        dst_engine_id: str,
+        request_id: str,
+        remote_request_id: str,
+        local_xfer_side_handle: int,
+        remote_xfer_side_handle: int,
+    ):
+        """
+        Post a READ point-to-point xfer request from a single local worker to
+        a single remote worker.
+        """
+        assert self.transfer_topo is not None
+        remote_rank = read_spec.remote_rank
+        local_block_ids = read_spec.local_block_ids
+        remote_block_ids = read_spec.remote_block_ids
+
+        remote_info = self.transfer_topo.get_engine_info(dst_engine_id)
+        block_size_ratio = self.transfer_topo.block_size_ratio(
+            remote_info.remote_block_size
+        )
+        if block_size_ratio > 1:
+            # TODO (NickLucche) assume HMA is off. Change to handle multiple KV groups.
+            assert not self._is_hma_required
+            local_block_ids0 = local_block_ids[0] if local_block_ids else []
+            remote_block_ids0 = remote_block_ids[0]
+            local_block_ids_mapped = self.get_mapped_blocks(
+                np.asarray(local_block_ids0), block_size_ratio
+            ).tolist()
+            if len(local_block_ids_mapped) > len(remote_block_ids0):
+                # NOTE:
+                # get_mapped_blocks will always expand block_ids for n times.
+                # ex:
+                # prefill block_ids with block_size as 4:
+                # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+                # Local decode block_ids with block_size as 16: [1, 2, 3]
+                # expanded decode block_ids with get_mapped_blocks from [1, 2, 3] to
+                # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
+                # Then we clip local to align with prefill
+                # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] to
+                # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+                local_block_ids_mapped = local_block_ids_mapped[
+                    : len(remote_block_ids0)
+                ]
+            local_block_ids = [local_block_ids_mapped] if local_block_ids_mapped else []
+            remote_block_ids = [remote_block_ids0]
+        # NOTE(rob): having the staging blocks be on the READER side is
+        # not going to work well (since we will have to call rearrange tensors).
+        # after we detect the txn is complete (which means we cannot make the
+        # read trxn async easily). If we want to make "READ" happen cleanly,
+        # then we will need to have the staging blocks on the remote side.
+
+        # NOTE(rob): according to nvidia the staging blocks are used to
+        # saturate IB with heterogeneous TP sizes.
+
+        # Number of D TP workers that will read from dst P. Propagate info
+        # on notification so that dst worker can wait before freeing blocks.
+        notif_id = f"{remote_request_id}:{self.world_size}".encode()
+
+        # Full prefix cache hit: do not need to read remote blocks,
+        # just notify P worker that we have the blocks we need.
+        if len(local_block_ids) == 0:
+            # A full prefix cache hit is indicated with an empty list.
+            agent_name = self._remote_agents[dst_engine_id][(0, remote_rank)]
+            try:
+                self.nixl_wrapper.send_notif(agent_name, notif_msg=notif_id)
+            except Exception as e:
+                self._log_failure(
+                    failure_type="notification_failed",
+                    msg="P worker blocks will be freed after timeout. "
+                    "This may indicate network issues.",
+                    req_id=request_id,
+                    error=e,
+                    dst_engine_id=dst_engine_id,
+                    remote_rank=remote_rank,
+                    remote_agent_name=agent_name,
+                )
+                self.xfer_stats.record_failed_notification()
+            return
+
+        assert (
+            len(remote_block_ids)
+            == len(local_block_ids)
+            == len(self.kv_cache_config.kv_cache_groups)
+        )
+        remote_physical_per_logical = remote_info.remote_physical_blocks_per_logical
+        local_block_ids, remote_block_ids = self._apply_prefix_caching(
+            local_block_ids, remote_block_ids, remote_physical_per_logical
+        )
+
+        # NOTE (nicolo) With homogeneous TP, each TP worker loads KV from
+        # corresponding rank. With heterogeneous TP, fixing D>P, the D tp
+        # workers will issue xfers to parts of the P worker remote kv caches.
+
+        # Get descs ids.
+        remote_block_descs_ids = self._compute_desc_ids(
+            block_ids=remote_block_ids,
+            dst_num_blocks=self.dst_num_blocks[dst_engine_id],
+            block_size_ratio=None,
+            physical_blocks_per_logical=remote_info.remote_physical_blocks_per_logical,
+        )
+        local_block_descs_ids = self._compute_desc_ids(
+            block_ids=local_block_ids,
+            dst_num_blocks=self.dst_num_blocks[self.engine_id],
+            block_size_ratio=block_size_ratio,
+            physical_blocks_per_logical=self._physical_blocks_per_logical_kv_block,
+        )
+
+        assert len(local_block_descs_ids) == len(remote_block_descs_ids)
+
+        # Prepare transfer with Nixl.
+        handle = None
+        try:
+            handle = self.nixl_wrapper.make_prepped_xfer(
+                "READ",
+                local_xfer_side_handle,
+                local_block_descs_ids,
+                remote_xfer_side_handle,
+                remote_block_descs_ids,
+                notif_msg=notif_id,
+            )
+
+            # Begin async xfer.
+            self.nixl_wrapper.transfer(handle)
+
+            # Use handle to check completion in future step().
+            self._recving_transfers[request_id].append(handle)
+        except Exception as e:
+            # mark all (logical) blocks for this request as invalid
+            self._log_failure(
+                failure_type="transfer_setup_failed",
+                req_id=request_id,
+                msg="Marking blocks as invalid",
+                error=e,
+                dst_engine_id=dst_engine_id,
+                remote_rank=remote_rank,
+            )
+            self._handle_failed_transfer(request_id, handle)
+
+    def _get_new_notifs(self) -> set[str]:
+        """
+        Get req_ids which got a remote xfer message. When multiple consumers
+        are reading from the same producer (heterogeneous TP scenario), wait
+        for all consumers to be done pulling.
+
+        Also handles heartbeat notifications ("HB:req1,req2,...") by
+        extending the lease on the referenced requests.
+        """
+        assert self.transfer_topo is not None
+        notified_req_ids: set[str] = set()
+        for notifs in self.nixl_wrapper.get_new_notifs().values():
+            for notif in notifs:
+                msg = notif.decode("utf-8")
+
+                # Handle heartbeat messages from D-side.
+                if msg.startswith("HB:"):
+                    self._handle_heartbeat(msg[3:])
+                    continue
+
+                req_id, tp_size = msg.rsplit(":", 1)
+                if (
+                    req_id not in self._reqs_to_send
+                    and req_id not in self._reqs_to_process
+                ):
+                    logger.error(
+                        "Potentially invalid KV blocks for "
+                        "unrecognized request %s were retrieved by "
+                        "a decode worker. They may have expired.",
+                        req_id,
+                    )
+                    continue
+
+                # NOTE: `tp_ratio` is the opposite when swapping local<>remote
+                n_consumers = int(tp_size)
+                tp_ratio = self.transfer_topo.tp_ratio(n_consumers)
+
+                # Number of reads *per producer* to wait for.
+                # When remote D TP > local P TP we expect `tp_ratio` reads.
+                consumers_per_producer = (
+                    -tp_ratio if n_consumers > self.world_size else 1
+                )
+
+                self.consumer_notification_counts_by_req[req_id] += 1
+                # Wait all consumers (D) to be done reading before freeing.
+                if (
+                    self.consumer_notification_counts_by_req[req_id]
+                    == consumers_per_producer
+                ):
+                    notified_req_ids.add(req_id)
+                    del self.consumer_notification_counts_by_req[req_id]
+                    self._reqs_to_process.remove(req_id)
+                    self._reqs_to_send.pop(req_id, None)
+        return notified_req_ids

vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py

@@ -0,0 +1,356 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Push-specific scheduler-side logic for the NIXL connector.
+
+In push mode, scheduler-side responsibilities are:
+
+* D side (decode): on ``update_state_after_alloc``, stash registration data
+  (D's identity + locally allocated block IDs) into
+  ``_push_pending_registrations``. The D worker drains it from
+  ``meta.push_registrations`` next step and sends a NIXL notification to the
+  P worker (no scheduler-level networking).
+* P side (prefill): on ``request_finished``, stash the finished block IDs
+  into ``_finished_request_blocks`` for the lease, and into
+  ``_newly_finished_push_blocks`` so the P worker picks them up via
+  ``meta.push_finished_blocks`` and matches against any D registrations
+  it already received via NIXL notifications.
+* Both sides: ``has_pending_push_work`` keeps the engine main loop stepping
+  while pushes are in flight. ``update_connector_output`` cleans up
+  ``_finished_request_blocks`` once the WRITE completes.
+
+A soft per-registration watchdog on the D scheduler fails requests that have
+been registered but not fulfilled within a configurable timeout.
+"""
+
+from __future__ import annotations
+
+import time
+from typing import TYPE_CHECKING, Any
+
+from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
+    KVConnectorMetadata,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_scheduler import (
+    NixlBaseConnectorScheduler,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
+    NixlConnectorMetadata,
+    ReqId,
+)
+from vllm.logger import init_logger
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+    from vllm.v1.core.sched.output import SchedulerOutput
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+    from vllm.v1.outputs import KVConnectorOutput
+    from vllm.v1.request import Request
+
+logger = init_logger(__name__)
+
+
+class NixlPushConnectorScheduler(NixlBaseConnectorScheduler):
+    """Push-specific scheduler logic (WRITE-based KV transfer).
+
+    All P2P communication is deferred to the worker level via NIXL
+    notifications. The scheduler communicates with workers only through
+    the standard ``build_connector_meta`` / ``update_connector_output``
+    hooks.
+    """
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        engine_id: str,
+        kv_cache_config: KVCacheConfig,
+    ):
+        super().__init__(vllm_config, engine_id, kv_cache_config)
+        if self.is_bidirectional_kv_xfer_enabled:
+            raise NotImplementedError(
+                "Bidirectional KV transfer is not supported for NIXL push connector."
+            )
+
+        # D-side: registration data to pass to D workers via metadata on
+        # the next ``build_connector_meta`` call.
+        self._push_pending_registrations: dict[ReqId, dict[str, Any]] = {}
+
+        # D-side: track the wall-clock deadline for each registered request
+        # to detect "registered but never fulfilled" failures (e.g. the P
+        # node disappeared after registration). Keyed by D request_id.
+        self._push_registration_deadlines: dict[ReqId, float] = {}
+
+        # P-side: block IDs for finished requests, kept for the lease and
+        # used to drive ``has_pending_push_work``.
+        self._finished_request_blocks: dict[ReqId, BlockIds] = {}
+        # P-side: newly finished blocks to ship to P workers on next step.
+        self._newly_finished_push_blocks: dict[ReqId, BlockIds] = {}
+
+        # Soft watchdog timeout (seconds) for D-side registrations that
+        # never receive a push completion. Defaults to the existing
+        # decoder KV blocks TTL so behaviour matches the lease.
+        assert vllm_config.kv_transfer_config is not None
+        self._push_registration_timeout: float = float(
+            vllm_config.kv_transfer_config.get_from_extra_config(
+                "push_registration_timeout",
+                self.decoder_kv_blocks_ttl,
+            )
+        )
+
+    def get_num_new_matched_tokens(
+        self, request: Request, num_computed_tokens: int
+    ) -> tuple[int, bool]:
+        """In push mode, D doesn't pull — it registers blocks and waits.
+
+        However, we still need to handle the do_remote_prefill case where D
+        needs to know how many tokens will be pushed.
+        """
+        params = request.kv_transfer_params
+        logger.debug(
+            "NixlPushConnector get_num_new_matched_tokens: "
+            "num_computed_tokens=%s, kv_transfer_params=%s",
+            num_computed_tokens,
+            params,
+        )
+
+        if params is not None and params.get("do_remote_prefill"):
+            token_ids = request.prompt_token_ids or []
+            actual = self._get_remote_prefill_token_count(len(token_ids))
+            count = actual - num_computed_tokens
+            if count > 0:
+                return count, True
+
+        if params is not None and params.get("do_remote_decode") and self._has_mamba:
+            self._truncate_mamba_request_for_prefill(request)
+
+        return 0, False
+
+    def update_state_after_alloc(
+        self, request: Request, blocks: KVCacheBlocks, num_external_tokens: int
+    ):
+        """In push mode, D stores registration data for the worker to send
+        to P via NIXL notification (deferred to ``build_connector_meta``).
+        """
+        params = request.kv_transfer_params
+        logger.debug(
+            "NixlPushConnector update_state_after_alloc: "
+            "num_external_tokens=%s, kv_transfer_params=%s",
+            num_external_tokens,
+            params,
+        )
+
+        if not params:
+            return
+
+        # P side: track the request as in-batch so the lease accounting
+        # matches what the worker expects on the next step.
+        if params.get("do_remote_decode"):
+            self._reqs_in_batch.add(request.request_id)
+
+        # P side with host-buffer offload: defer save to the worker.
+        if self.use_host_buffer and params.get("do_remote_decode"):
+            self._reqs_need_save[request.request_id] = request
+            return
+
+        # D side: only act on the first call (``do_remote_prefill`` is
+        # unset on re-entry by the marker below).
+        if not params.get("do_remote_prefill"):
+            return
+
+        if num_external_tokens <= 0:
+            # Nothing to receive: full prefix-cache hit on D, no
+            # registration to stage.
+            return
+
+        # First-pass D path: stash registration data the worker will
+        # ship to P on the next ``build_connector_meta`` cycle.
+        logger.debug(
+            "KV PUSH mode: D node storing registration for request %s",
+            request.request_id,
+        )
+        local_block_ids: BlockIds = blocks.get_unhashed_block_ids_all_groups()
+        local_block_ids = self.get_sw_clipped_blocks(local_block_ids)
+
+        # ``remote_*`` fields are P's coordinates (from D's perspective).
+        # ``decode_*`` fields are D's own info that P needs for the
+        # reverse handshake before WRITE-ing.
+        self._push_pending_registrations[request.request_id] = {
+            "request_id": request.request_id,
+            "decode_engine_id": self.engine_id,
+            "decode_host": self.side_channel_host,
+            "decode_port": self.side_channel_port,
+            "decode_tp_size": (self.vllm_config.parallel_config.tensor_parallel_size),
+            "local_block_ids": local_block_ids,
+            "remote_engine_id": params["remote_engine_id"],
+            "remote_host": params["remote_host"],
+            "remote_port": params["remote_port"],
+            "remote_tp_size": params["tp_size"],
+            "remote_pp_size": params.get("pp_size", 1),
+        }
+        self._push_registration_deadlines[request.request_id] = (
+            time.perf_counter() + self._push_registration_timeout
+        )
+        # In push mode D doesn't know P's blocks; P determines them
+        # from the registration. We still track the request as
+        # needing recv so the engine waits for P's WRITE completion.
+        # ``remote_block_ids`` is also seeded to an empty tuple so the
+        # base scheduler's ``add_new_req_to_recv`` can build the
+        # ReqMeta without a KeyError — the actual remote block IDs are
+        # learned by P over the NIXL handshake at WRITE time.
+        params["remote_block_ids"] = ()
+        self._reqs_need_recv[request.request_id] = (request, local_block_ids)
+
+        # Mark as processed so a re-entry (e.g. preemption + reschedule)
+        # doesn't re-stage the registration.
+        params["do_remote_prefill"] = False
+
+    def request_finished(
+        self,
+        request: Request,
+        block_ids: BlockIds,
+    ) -> tuple[bool, dict[str, Any] | None]:
+        """Push-mode request_finished: stores blocks for workers."""
+        from vllm.v1.request import RequestStatus
+
+        params = request.kv_transfer_params
+        logger.debug(
+            "NixlPushConnector request_finished(%s), request_status=%s, "
+            "kv_transfer_params=%s",
+            request.request_id,
+            request.status,
+            params,
+        )
+        if not params:
+            return False, None
+
+        is_p_node = bool(params.get("do_remote_decode"))
+
+        self._stop_heartbeat(request.request_id)
+        # Drop any pending registration deadline; the request either
+        # completed or was cancelled.
+        self._push_registration_deadlines.pop(request.request_id, None)
+
+        if params.get("do_remote_prefill"):
+            # ``do_remote_prefill`` is still set, which means
+            # ``update_state_after_alloc`` never ran (it would have
+            # flipped this flag to False). The request was aborted
+            # before it could be scheduled — e.g. rejected at the D
+            # serving layer via abort_immediately. To keep P from
+            # stranding the prefill blocks, we still register an empty
+            # recv so the worker emits a notif that lets P free them.
+            # Seed remote_block_ids so add_new_req_to_recv won't KeyError.
+            params["remote_block_ids"] = ()
+            self._reqs_need_recv[request.request_id] = (request, [])
+            params["do_remote_prefill"] = False
+            return False, None
+
+        # Push connector only acts on the P-side terminal path; D-side
+        # finishing without a remote prefill is a no-op.
+        if not is_p_node:
+            return False, None
+
+        if request.status not in (
+            RequestStatus.FINISHED_LENGTH_CAPPED,
+            RequestStatus.FINISHED_STOPPED,
+        ):
+            self._reqs_not_processed.add(request.request_id)
+            self._reqs_need_save.pop(request.request_id, None)
+            return False, None
+
+        delay_free_blocks = any(len(group) > 0 for group in block_ids)
+        remote_num_tokens = 0
+        if delay_free_blocks:
+            logger.debug(
+                "NixlPushConnector request_finished(%s) waiting for %d seconds "
+                "before releasing blocks",
+                request.request_id,
+                self._kv_lease_duration,
+            )
+            self._reqs_need_send[request.request_id] = (
+                time.perf_counter() + self._kv_lease_duration
+            )
+
+            block_ids = self.get_sw_clipped_blocks(block_ids)
+            remote_num_tokens = request.num_computed_tokens
+
+            # Store finished blocks for worker-level matching with D
+            # registrations (via NIXL notifications).
+            self._finished_request_blocks[request.request_id] = block_ids
+            self._newly_finished_push_blocks[request.request_id] = block_ids
+
+        return delay_free_blocks, dict(
+            do_remote_prefill=True,
+            do_remote_decode=False,
+            remote_block_ids=block_ids,
+            remote_engine_id=self.engine_id,
+            remote_request_id=request.request_id,
+            remote_host=self.side_channel_host,
+            remote_port=self.side_channel_port,
+            tp_size=self.vllm_config.parallel_config.tensor_parallel_size,
+            pp_size=self.vllm_config.parallel_config.pipeline_parallel_size,
+            remote_num_tokens=remote_num_tokens,
+        )
+
+    def build_connector_meta(
+        self,
+        scheduler_output: SchedulerOutput,
+    ) -> KVConnectorMetadata:
+        meta = super().build_connector_meta(scheduler_output)
+        assert isinstance(meta, NixlConnectorMetadata)
+
+        # Watchdog: any D-side registration whose deadline has passed without
+        # a corresponding push completion is treated as failed and cleaned up.
+        # The corresponding request is already tracked via _reqs_need_recv;
+        # the engine layer will eventually time it out via the lease, but we
+        # at least drop the stale registration so we don't keep retrying.
+        now = time.perf_counter()
+        # Deadlines are inserted in non-decreasing order (monotonic clock +
+        # constant timeout, armed once per request), and dict insertion order
+        # is preserved across key deletions, so we can stop at the first
+        # not-yet-expired entry instead of scanning the whole dict.
+        expired = []
+        for rid, deadline in self._push_registration_deadlines.items():
+            if deadline > now:
+                break
+            expired.append(rid)
+        for rid in expired:
+            self._push_registration_deadlines.pop(rid, None)
+            # Avoid resending a registration that already timed out.
+            self._push_pending_registrations.pop(rid, None)
+            logger.warning(
+                "NixlPushConnector: registration for request %s timed out "
+                "after %.1fs without a push completion",
+                rid,
+                self._push_registration_timeout,
+            )
+
+        # D side: package pending registrations for D workers to send out.
+        if self._push_pending_registrations:
+            meta.push_registrations = dict(self._push_pending_registrations)
+            self._push_pending_registrations.clear()
+
+        # P side: package newly finished blocks for P workers to match against
+        # any D registrations they have received via NIXL notifications.
+        if self._newly_finished_push_blocks:
+            meta.push_finished_blocks = dict(self._newly_finished_push_blocks)
+            self._newly_finished_push_blocks.clear()
+
+        return meta
+
+    def has_pending_push_work(self) -> bool:
+        # Keep the engine main loop alive while we have:
+        # - finished P blocks awaiting WRITE completion, or
+        # - pending D registrations the worker has not yet shipped, or
+        # - newly finished blocks not yet shipped to P workers.
+        return bool(self._finished_request_blocks or self._push_pending_registrations)
+
+    def update_connector_output(self, connector_output: KVConnectorOutput) -> None:
+        """Clean up finished request blocks after push completes."""
+        super().update_connector_output(connector_output)
+        for req_id in connector_output.finished_sending or ():
+            self._finished_request_blocks.pop(req_id, None)
+        # On D side, finished_recving means the push completed; clear the
+        # watchdog so we don't trip an expiration on a fulfilled request.
+        for req_id in connector_output.finished_recving or ():
+            self._push_registration_deadlines.pop(req_id, None)

vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py

@@ -0,0 +1,780 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Push-specific (WRITE) worker-side logic for the NIXL connector.
+
+A dedicated ``nixl-push-writer`` thread owns all push-related NIXL ops:
+calls ``get_new_notifs`` (routing PUSH_REG internally; HB / completion
+notifs are forwarded to the engine main thread), sends PUSH_REG via
+``send_notif``, matches D registrations with P finished blocks, and
+issues WRITE transfers via ``make_prepped_xfer`` / ``transfer``.
+
+The engine main thread feeds the writer through three queues:
+``_reg_send_inbox`` (D-side regs to send), ``_finished_blocks_inbox``
+(P-side blocks from metadata) and ``_pending_completion_notifs``
+(non-PUSH_REG notifs forwarded back for HB / completion accounting).
+
+Wake model: the writer self-polls every
+``_PUSH_WRITER_POLL_INTERVAL_MS`` only while it has unmatched
+``_push_finished_blocks`` (i.e. P-side blocks waiting for a D PUSH_REG
+notif that has no other wake source). All other progress is
+event-driven: the engine main thread sets ``_push_writer_wake`` from
+``start_load_kv`` (when handing it new work) and from ``get_finished``
+(so each engine step gives the writer a chance to drain NIXL notifs);
+the handshake-completion callback sets the same event after a deferred
+PUSH_REG send has been queued. When a request's lease expires (the base
+worker reports it via ``done_sending``) or the WRITE completes,
+``get_finished`` enqueues an eviction onto ``_evict_finished_inbox`` so
+the writer drops any leftover ``_push_finished_blocks`` /
+``_pending_d_registrations`` and stops self-polling.
+"""
+
+import queue
+import threading
+import time
+from collections import defaultdict
+from concurrent.futures import Future
+from typing import TYPE_CHECKING, Any
+
+import msgspec
+import numpy as np
+
+from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker import (
+    NixlBaseConnectorWorker,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
+    PUSH_REG_NOTIF_PREFIX,
+    NixlConnectorMetadata,
+    RemoteMeta,
+    ReqId,
+    ReqMeta,
+    TransferHandle,
+)
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import ReadSpec
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import get_base_request_id
+from vllm.logger import init_logger
+
+if TYPE_CHECKING:
+    import torch
+
+    from vllm.config import VllmConfig
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+
+logger = init_logger(__name__)
+
+# Writer-thread poll cadence while there is in-flight push state. When
+# fully idle, the writer blocks on a wake event signalled by the engine
+# main thread (start_load_kv / get_finished). Smaller -> lower latency
+# while active, slightly more CPU.
+_PUSH_WRITER_POLL_INTERVAL_MS = 1.0
+
+
+class NixlPushConnectorWorker(NixlBaseConnectorWorker):
+    """Push-specific (WRITE) worker logic. See module docstring."""
+
+    def __init__(
+        self,
+        vllm_config: "VllmConfig",
+        engine_id: str,
+        kv_cache_config: "KVCacheConfig",
+    ):
+        super().__init__(vllm_config, engine_id, kv_cache_config)
+
+        # Heartbeat handshakes to a PP-sharded producer must be notif-only,
+        # like the PUSH_REG path.
+        self._hb_handshake_notif_only = True
+
+        # Push-specific state.
+        # P-side: outgoing WRITE handles awaiting completion, keyed by
+        # request_id. Mutated by writer (submit) and main thread
+        # (``_pop_done_transfers``); guarded by
+        # ``_sending_transfers_lock``.
+        self._sending_transfers = defaultdict[ReqId, list[TransferHandle]](list)
+        self._sending_transfers_lock = threading.Lock()
+
+        # Writer-thread owned matching state.
+        # P-side: finished request blocks received from scheduler metadata
+        # that have not yet been matched with an incoming D registration.
+        self._push_finished_blocks: dict[ReqId, BlockIds] = {}
+        # P-side: D registrations received via NIXL notification that have
+        # not yet been matched with a finished P request.
+        self._pending_d_registrations: dict[ReqId, dict[str, Any]] = {}
+
+        # Cross-thread channels.
+        self._reg_send_inbox: queue.Queue[tuple[str, dict[str, Any]]] = queue.Queue()
+        self._finished_blocks_inbox: queue.Queue[tuple[str, BlockIds]] = queue.Queue()
+        self._pending_completion_notifs: queue.Queue[bytes] = queue.Queue()
+        # Main thread → writer: req_ids whose lease has expired or whose
+        # WRITE has completed. Writer drops them from
+        # ``_push_finished_blocks`` so an unmatched entry doesn't keep the
+        # writer busy-polling forever.
+        self._evict_finished_inbox: queue.Queue[str] = queue.Queue()
+
+        # Wake signal from engine main thread (start_load_kv / get_finished).
+        # Writer self-polls at _PUSH_WRITER_POLL_INTERVAL_MS while it has
+        # active in-flight state; otherwise it blocks until signalled.
+        self._push_writer_wake = threading.Event()
+
+        self._push_writer_stop = threading.Event()
+        self._push_writer_thread: threading.Thread | None = None
+
+    # --- Lifecycle ----------------------------------------------------- #
+
+    def register_kv_caches(self, kv_caches: dict[str, "torch.Tensor"]):
+        super().register_kv_caches(kv_caches)
+        if self._push_writer_thread is None:
+            self._push_writer_thread = threading.Thread(
+                target=self._push_writer_loop,
+                daemon=True,
+                name="nixl-push-writer",
+            )
+            self._push_writer_thread.start()
+            logger.info("nixl-push-writer thread started (rank=%d)", self.tp_rank)
+
+    def shutdown(self):
+        self._push_writer_stop.set()
+        # Unblock the writer if it's waiting in the no-active-state branch.
+        self._push_writer_wake.set()
+        if self._push_writer_thread is not None:
+            self._push_writer_thread.join(timeout=2)
+            self._push_writer_thread = None
+        with self._sending_transfers_lock:
+            for handles in self._sending_transfers.values():
+                for handle in handles:
+                    self.nixl_wrapper.release_xfer_handle(handle)
+            self._sending_transfers.clear()
+        super().shutdown()
+
+    # --- Engine-main-thread entry point -------------------------------- #
+
+    def start_load_kv(self, metadata: NixlConnectorMetadata):
+        """Pre-process metadata; defer NIXL ops to the writer thread."""
+        # D-side: track reqs waiting for P to push.
+        for req_id, meta in metadata.reqs_to_recv.items():
+            meta.local_physical_block_ids = self._logical_to_kernel_block_ids(
+                meta.local_block_ids
+            )
+            assert meta.remote is not None
+            remote_engine_id = meta.remote.engine_id
+            logger.debug(
+                "start_load_kv (push) for request %s from remote engine %s. "
+                "Num local_block_ids: %s. Num remote_block_ids: %s. ",
+                req_id,
+                remote_engine_id,
+                len(meta.local_physical_block_ids),
+                len(meta.remote.block_ids),
+            )
+            self._recving_metadata[req_id] = meta
+
+        # --- D-side: registrations to send to P via NIXL ---
+        if metadata.push_registrations:
+            for req_id, reg_data in metadata.push_registrations.items():
+                self._reg_send_inbox.put((req_id, reg_data))
+            self._push_writer_wake.set()
+
+        # --- P-side: newly finished blocks awaiting a D registration match ---
+        if metadata.push_finished_blocks:
+            for req_id, block_ids in metadata.push_finished_blocks.items():
+                self._finished_blocks_inbox.put((req_id, block_ids))
+            self._push_writer_wake.set()
+
+        # Batch + lease tracking (same as pull).
+        for req_id in metadata.reqs_in_batch:
+            self._reqs_to_process.add(req_id)
+        for req_id in metadata.reqs_not_processed:
+            self._reqs_to_process.discard(req_id)
+            assert req_id not in self._reqs_to_send
+        for req_id, expiration_time in metadata.reqs_to_send.items():
+            if req_id in self._reqs_to_process:
+                self._reqs_to_send[req_id] = expiration_time
+
+        # Heartbeats still leave from the main thread (base worker behaviour).
+        self._send_heartbeats(metadata)
+
+    # --- Writer thread ------------------------------------------------- #
+
+    def _push_writer_loop(self) -> None:
+        sleep_s = _PUSH_WRITER_POLL_INTERVAL_MS / 1000.0
+
+        while not self._push_writer_stop.is_set():
+            try:
+                # 1. D registrations to send.
+                while True:
+                    try:
+                        rid, rd = self._reg_send_inbox.get_nowait()
+                    except queue.Empty:
+                        break
+                    self._send_registration_to_p(rid, rd)
+
+                # 2. P-side finished blocks; match against pending regs.
+                while True:
+                    try:
+                        rid, blocks = self._finished_blocks_inbox.get_nowait()
+                    except queue.Empty:
+                        break
+                    matched = self._pop_matching_registration(rid)
+                    if matched is not None:
+                        self._do_start_push_kv(rid, blocks, matched)
+                    else:
+                        self._push_finished_blocks[rid] = blocks
+
+                # 2b. Evict finished blocks for requests that have either
+                # completed (WRITE acknowledged) or whose lease expired
+                # without a D registration.  Drop pending registrations
+                # for the same reason so we don't leak state.
+                while True:
+                    try:
+                        rid = self._evict_finished_inbox.get_nowait()
+                    except queue.Empty:
+                        break
+                    self._push_finished_blocks.pop(rid, None)
+                    self._pending_d_registrations.pop(rid, None)
+
+                # 3. NIXL notifs: route PUSH_REG; forward the rest.
+                for notifs in self.nixl_wrapper.get_new_notifs().values():
+                    for notif in notifs:
+                        if notif.startswith(PUSH_REG_NOTIF_PREFIX):
+                            self._handle_push_reg_notif(notif)
+                        else:
+                            self._pending_completion_notifs.put(notif)
+            except Exception:
+                logger.exception("nixl-push-writer error; continuing")
+
+            # Self-poll only while there is no other wake source: P-side
+            # finished blocks waiting for a D PUSH_REG match. All other
+            # progress is event-driven (see module docstring).
+            if self._push_finished_blocks:
+                self._push_writer_stop.wait(timeout=sleep_s)
+            else:
+                self._push_writer_wake.wait()
+                self._push_writer_wake.clear()
+
+    def _handle_push_reg_notif(self, notif: bytes) -> None:
+        try:
+            reg_data = msgspec.msgpack.decode(notif[len(PUSH_REG_NOTIF_PREFIX) :])
+        except Exception:
+            logger.exception("Failed to decode PUSH_REG notification payload")
+            return
+        rid = reg_data.get("request_id") if isinstance(reg_data, dict) else None
+        if not isinstance(rid, str):
+            logger.warning("PUSH_REG notif missing request_id; dropping")
+            return
+
+        match = self._pop_matching_finished_blocks(rid)
+        if match is not None:
+            fin_id, blocks = match
+            self._do_start_push_kv(fin_id, blocks, reg_data)
+        else:
+            self._pending_d_registrations[rid] = reg_data
+
+    # --- D-side registration send (writer thread) ---------------------- #
+
+    def _send_registration_to_p(
+        self,
+        req_id: str,
+        reg_data: dict[str, Any],
+    ) -> None:
+        """Handshake (if needed) then send PUSH_REG. ``send_notif`` always
+        executes on the writer; the handshake runs on the background executor
+        and the request is re-queued onto ``_reg_send_inbox`` once it
+        completes (at which point ``_ensure_handshake`` returns ``None`` and we
+        send directly)."""
+        remote_pp_size = reg_data.get("remote_pp_size", 1)
+        fut = self._ensure_handshake(
+            reg_data["remote_engine_id"],
+            reg_data["remote_host"],
+            reg_data["remote_port"],
+            reg_data["remote_tp_size"],
+            pp_size=remote_pp_size,
+            # D never addresses P memory in push mode; just load P's agents.
+            notif_agents_only=remote_pp_size > 1,
+        )
+        if fut is None:
+            self._do_send_reg_notif(req_id, reg_data)
+            return
+
+        def _on_handshake(
+            f: Future[dict[tuple[int, int], str]],
+            rid: str = req_id,
+            rd: dict[str, Any] = reg_data,
+        ) -> None:
+            try:
+                f.result()
+            except Exception as e:
+                self._log_failure(
+                    failure_type="push_reg_handshake_failed", req_id=rid, error=e
+                )
+                self._handle_failed_transfer(rid, None)
+                return
+            # Re-queue for the writer to send now that the handshake is done.
+            self._reg_send_inbox.put((rid, rd))
+            # Wake the writer so it sends the PUSH_REG promptly even if
+            # otherwise parked.
+            self._push_writer_wake.set()
+
+        fut.add_done_callback(_on_handshake)
+
+    def _do_send_reg_notif(self, req_id: str, reg_data: dict[str, Any]) -> None:
+        engine_id = reg_data["remote_engine_id"]
+        notif_msg = PUSH_REG_NOTIF_PREFIX + msgspec.msgpack.encode(reg_data)
+        agents = self._remote_agents.get(engine_id)
+        if not agents:
+            logger.error(
+                "No remote agents for engine %s; cannot send registration for %s",
+                engine_id,
+                req_id,
+            )
+            self._handle_failed_transfer(req_id, None)
+            return
+        for rank, agent_name in agents.items():
+            try:
+                self.nixl_wrapper.send_notif(agent_name, notif_msg=notif_msg)
+            except Exception as e:
+                self._log_failure(
+                    failure_type="push_reg_notif_failed",
+                    req_id=req_id,
+                    error=e,
+                    remote_rank=rank,
+                )
+        logger.debug(
+            "Sent PUSH_REG for %s to engine %s (%dB)", req_id, engine_id, len(notif_msg)
+        )
+
+    # --- Matching helpers --------------------------------------------- #
+
+    def _pop_matching_registration(self, request_id: str) -> dict[str, Any] | None:
+        """Pop the D-side registration matching *request_id*.
+
+        Exact key first, then a match after stripping the random suffix from
+        both sides. No match leaves the request unmatched (push not started).
+        """
+        data = self._pending_d_registrations.pop(request_id, None)
+        if data is not None:
+            return data
+        base_id = get_base_request_id(request_id)
+        for reg_id in list(self._pending_d_registrations):
+            if get_base_request_id(reg_id) == base_id:
+                return self._pending_d_registrations.pop(reg_id)
+        return None
+
+    def _pop_matching_finished_blocks(
+        self, request_id: str
+    ) -> tuple[str, BlockIds] | None:
+        """Pop the P-side finished blocks matching *request_id*.
+
+        Same lookup as ``_pop_matching_registration``: exact key, then a
+        match after stripping the random suffix from both sides.
+        """
+        blocks = self._push_finished_blocks.pop(request_id, None)
+        if blocks is not None:
+            return request_id, blocks
+        base_id = get_base_request_id(request_id)
+        for fin_id in list(self._push_finished_blocks):
+            if get_base_request_id(fin_id) == base_id:
+                return fin_id, self._push_finished_blocks.pop(fin_id)
+        return None
+
+    # --- WRITE transfer logic (writer thread) ------------------------- #
+
+    def _do_start_push_kv(
+        self,
+        request_id: str,
+        local_block_ids: BlockIds,
+        registration_data: dict[str, Any],
+    ) -> None:
+        """Start push-based KV transfer from P worker to D node.
+
+        ``local_block_ids`` are P's *logical* block IDs (from the P
+        scheduler's metadata). ``registration_data["local_block_ids"]``
+        are D's *logical* block IDs (from D's scheduler, sent over the
+        PUSH_REG notif). All conversion to physical block IDs is
+        deferred to ``_xfer_blocks_for_req`` so each side uses its own
+        physical-blocks-per-logical ratio (P uses
+        ``self._physical_blocks_per_logical_kv_block``; D's ratio is
+        learned during the NIXL handshake)."""
+        decode_engine_id = registration_data["decode_engine_id"]
+        remote_block_ids = registration_data["local_block_ids"]
+        decode_host = registration_data["decode_host"]
+        decode_port = registration_data["decode_port"]
+        decode_request_id = registration_data["request_id"]
+        if not local_block_ids:
+            logger.warning("No local blocks to push for request %s", request_id)
+            return
+
+        if not self._ensure_d_handshake(
+            decode_engine_id,
+            decode_host,
+            decode_port,
+            registration_data["decode_tp_size"],
+            request_id,
+        ):
+            return
+
+        # Both sides are kept in logical form here; ``_xfer_blocks_for_req``
+        # expands each side using the appropriate ratio.
+        logical_local = self._as_grouped_block_ids(local_block_ids)
+        logical_remote = self._as_grouped_block_ids(remote_block_ids)
+        physical_local = self._logical_to_kernel_block_ids(logical_local)
+
+        push_meta = ReqMeta(
+            local_block_ids=logical_local,
+            local_physical_block_ids=physical_local,
+            tp_size=self.world_size,
+            remote=RemoteMeta(
+                block_ids=logical_remote,
+                host="",
+                port=0,
+                engine_id=decode_engine_id,
+                request_id=decode_request_id,
+            ),
+        )
+
+        t0 = time.perf_counter()
+        self._xfer_blocks_for_req(req_id=request_id, meta=push_meta)
+        elapsed_ms = (time.perf_counter() - t0) * 1000.0
+        if elapsed_ms > 200.0:
+            logger.warning(
+                "_do_start_push_kv for %s took %.1fms (slow NIXL submission)",
+                request_id,
+                elapsed_ms,
+            )
+
+    def _ensure_d_handshake(
+        self,
+        decode_engine_id: str,
+        decode_host: str,
+        decode_port: int,
+        decode_tp_size: int,
+        request_id: str,
+    ) -> bool:
+        """First-time P→D handshake. Blocking call on the writer thread.
+
+        Returns True iff the handshake succeeded (or had already been
+        completed). Returns False if the handshake raised; the request is
+        skipped in that case (the engine layer will reschedule or fail it
+        via the standard lease/timeout path)."""
+        if decode_engine_id in self._remote_agents:
+            return True
+        try:
+            remote_agents = self._nixl_handshake(
+                decode_host,
+                decode_port,
+                decode_tp_size,
+                decode_engine_id,
+            )
+        except Exception:
+            logger.exception(
+                "Failed handshake to D %s for push %s",
+                decode_engine_id,
+                request_id,
+            )
+            return False
+        with self._handshake_lock:
+            self._remote_agents[decode_engine_id] = remote_agents
+        logger.info(
+            "Push handshake to D %s done (%d agents)",
+            decode_engine_id,
+            len(remote_agents),
+        )
+        return True
+
+    @staticmethod
+    def _as_grouped_block_ids(block_ids: BlockIds) -> BlockIds:
+        """Normalise a sequence of block IDs to a tuple-of-groups shape.
+
+        ``BlockIds`` is canonically a tuple of per-group lists, but some
+        registration payloads collapse a single-group case to a flat
+        list. Re-wrap that case so downstream group-aware helpers see a
+        consistent shape."""
+        if block_ids and not isinstance(block_ids[0], (list, tuple)):
+            return (list(block_ids),)
+        return block_ids
+
+    def _xfer_blocks_for_req(self, req_id: str, meta: ReqMeta):
+        """Issue WRITE transfers to one or more remote TP ranks."""
+        assert meta.remote is not None and self.transfer_topo is not None
+        engine_id = meta.remote.engine_id
+        plan = self.tp_mappings[engine_id]
+        remote_info = self.transfer_topo.get_engine_info(engine_id)
+        tp_ratio = self.transfer_topo.tp_ratio(remote_info.remote_tp_size)
+
+        # Expand D's logical IDs using the ratio learned during the
+        # NIXL handshake. ``meta`` is freshly built by
+        # ``_do_start_push_kv`` so mutating it here is safe.
+        meta.remote.block_ids = self._logical_to_remote_kernel_block_ids(
+            meta.remote.block_ids,
+            remote_info.remote_physical_blocks_per_logical,
+        )
+        remote_block_ids = meta.remote.block_ids
+        local_block_ids = meta.local_physical_block_ids
+        num_groups = len(local_block_ids)
+
+        if self.use_mla and tp_ratio < 0:
+            # MLA latent is replicated across D's TP ranks: the tp-mapping
+            # collapses to one rank (fine for reads), but push must WRITE every
+            # D rank or the rest decode stale KV; only the dst differs per rank.
+            assert len(plan.all_source_ranks) == 1
+            mla_local_ids = [list(ids) for ids in local_block_ids]
+            mla_remote_ids = [list(ids) for ids in remote_block_ids]
+            read_specs = [
+                ReadSpec(
+                    remote_rank=rank,
+                    local_block_ids=mla_local_ids,
+                    remote_block_ids=mla_remote_ids,
+                )
+                for rank in self.dst_xfer_side_handles[engine_id]
+            ]
+        else:
+            read_specs = [
+                ReadSpec(
+                    remote_rank=rank,
+                    local_block_ids=[
+                        list(local_block_ids[g])
+                        if rank in plan.source_ranks_per_group[g]
+                        else []
+                        for g in range(num_groups)
+                    ],
+                    remote_block_ids=[
+                        list(remote_block_ids[g])
+                        if rank in plan.source_ranks_per_group[g]
+                        else []
+                        for g in range(num_groups)
+                    ],
+                )
+                for rank in plan.all_source_ranks
+            ]
+
+        handles: list[int] = []
+        for i, spec in enumerate(read_specs):
+            remote_block_size = remote_info.remote_block_size
+            logger.debug(
+                "Remote agent %s available, calling _xfer_blocks"
+                " on remote rank %s with remote block size %s for req %s",
+                meta.remote.engine_id,
+                spec.remote_rank,
+                remote_block_size,
+                req_id,
+            )
+            if tp_ratio < 0 and not self.use_mla:
+                assert remote_block_size == self.block_size
+                local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[tp_ratio][i]
+            else:
+                local_xfer_side_handle = self.src_xfer_handles_by_block_size[
+                    remote_block_size
+                ]
+
+            remote_xfer_side_handle = self.dst_xfer_side_handles[meta.remote.engine_id][
+                spec.remote_rank
+            ]
+
+            handle = self._xfer_blocks(
+                read_spec=spec,
+                request_id=req_id,
+                dst_engine_id=meta.remote.engine_id,
+                remote_request_id=meta.remote.request_id,
+                local_xfer_side_handle=local_xfer_side_handle,
+                remote_xfer_side_handle=remote_xfer_side_handle,
+            )
+            if handle is not None:
+                handles.append(handle)
+
+        # Publish all the request's WRITE handles in one locked update: a
+        # partial set would let ``_pop_done_transfers`` finish the request
+        # early, then double-report it as the remaining writes land.
+        if handles:
+            with self._sending_transfers_lock:
+                self._sending_transfers[req_id].extend(handles)
+
+    def _xfer_blocks(
+        self,
+        read_spec: ReadSpec,
+        dst_engine_id: str,
+        request_id: str,
+        remote_request_id: str,
+        local_xfer_side_handle: int,
+        remote_xfer_side_handle: int,
+    ) -> int | None:
+        """Post a WRITE point-to-point xfer request.
+
+        Returns the in-flight transfer handle (so the caller can track all of
+        a request's handles atomically), or ``None`` if nothing was submitted.
+        """
+        assert self.transfer_topo is not None
+        remote_rank = read_spec.remote_rank
+        local_block_ids = read_spec.local_block_ids
+        remote_block_ids = read_spec.remote_block_ids
+
+        remote_info = self.transfer_topo.get_engine_info(dst_engine_id)
+        block_size_ratio = self.transfer_topo.block_size_ratio(
+            remote_info.remote_block_size
+        )
+        if block_size_ratio > 1:
+            assert not self._is_hma_required
+            local_block_ids0 = local_block_ids[0] if local_block_ids else []
+            remote_block_ids0 = remote_block_ids[0]
+            local_block_ids_mapped = self.get_mapped_blocks(
+                np.asarray(local_block_ids0), block_size_ratio
+            ).tolist()
+            if len(local_block_ids_mapped) > len(remote_block_ids0):
+                local_block_ids_mapped = local_block_ids_mapped[
+                    : len(remote_block_ids0)
+                ]
+            local_block_ids = [local_block_ids_mapped] if local_block_ids_mapped else []
+            remote_block_ids = [remote_block_ids0]
+
+        notif_id = f"{remote_request_id}:{self.world_size}".encode()
+
+        if len(local_block_ids) == 0:
+            logger.warning("No blocks to push for request %s", request_id)
+            return None
+
+        # Align per-group block counts for push.
+        local_block_ids = list(local_block_ids)
+        remote_block_ids = list(remote_block_ids)
+        for i in range(min(len(local_block_ids), len(remote_block_ids))):
+            num_local = len(local_block_ids[i])
+            num_remote = len(remote_block_ids[i])
+            if num_local > num_remote:
+                local_block_ids[i] = local_block_ids[i][:num_remote]
+            elif num_local < num_remote:
+                remote_block_ids[i] = remote_block_ids[i][:num_local]
+
+        # Get descs ids.
+        remote_block_descs_ids = self._compute_desc_ids(
+            block_ids=remote_block_ids,
+            dst_num_blocks=self.dst_num_blocks[dst_engine_id],
+            block_size_ratio=None,
+            physical_blocks_per_logical=remote_info.remote_physical_blocks_per_logical,
+        )
+        local_block_descs_ids = self._compute_desc_ids(
+            block_ids=local_block_ids,
+            dst_num_blocks=self.dst_num_blocks[self.engine_id],
+            block_size_ratio=block_size_ratio,
+            physical_blocks_per_logical=self._physical_blocks_per_logical_kv_block,
+        )
+
+        assert len(local_block_descs_ids) == len(remote_block_descs_ids)
+
+        handle = None
+        try:
+            handle = self.nixl_wrapper.make_prepped_xfer(
+                "WRITE",
+                local_xfer_side_handle,
+                local_block_descs_ids,
+                remote_xfer_side_handle,
+                remote_block_descs_ids,
+                notif_msg=notif_id,
+            )
+            self.nixl_wrapper.transfer(handle)
+            # Caller tracks the handle (atomically with the request's other
+            # writes) so P can free blocks once all of them are done.
+            return handle
+        except Exception as e:
+            self._log_failure(
+                failure_type="transfer_setup_failed",
+                req_id=request_id,
+                msg="Push WRITE submission failed; releasing handle",
+                error=e,
+                dst_engine_id=dst_engine_id,
+                remote_rank=remote_rank,
+            )
+            # On the P side this WRITE failure is purely outbound; we
+            # don't have a ``_recving_metadata`` entry to invalidate, so
+            # we just release the handle and let the engine reschedule
+            # via the lease / watchdog.
+            if handle is not None:
+                self.nixl_wrapper.release_xfer_handle(handle)
+            self.xfer_stats.record_failed_transfer()
+            return None
+
+    # --- Notification handling on engine main thread ------------------ #
+
+    def _get_new_notifs(self) -> set[str]:
+        """Drain HB / completion notifs forwarded by the writer thread.
+
+        The writer owns ``nixl_wrapper.get_new_notifs`` for push; PUSH_REG
+        notifs are handled there. Everything else is forwarded here for
+        existing accounting.
+        """
+        assert self.transfer_topo is not None
+        notified_req_ids: set[str] = set()
+        while True:
+            try:
+                notif = self._pending_completion_notifs.get_nowait()
+            except queue.Empty:
+                break
+
+            msg = notif.decode("utf-8")
+            if msg.startswith("HB:"):
+                self._handle_heartbeat(msg[3:])
+                continue
+
+            req_id, tp_size = msg.rsplit(":", 1)
+
+            # Not tracked as a P-side send/process for this notif.
+            if req_id not in self._reqs_to_send and req_id not in self._reqs_to_process:
+                if (meta := self._recving_metadata.get(req_id)) is not None:
+                    # Consumer waits for one notif per producer rank writing
+                    # here: pp_size stages * producers-per-consumer (>1 when
+                    # producer TP > consumer TP; tp_size is the producer TP).
+                    producers_per_consumer = max(1, int(tp_size) // self.world_size)
+                    expected_notifs = meta.pp_size * producers_per_consumer
+                    self.consumer_notification_counts_by_req[req_id] += 1
+                    notifs = self.consumer_notification_counts_by_req[req_id]
+                    if notifs < expected_notifs:
+                        continue
+                    del self.consumer_notification_counts_by_req[req_id]
+                    # P drove the transfer (we own no NIXL handle), so
+                    # materialise an empty ``_recving_transfers`` entry for
+                    # ``_pop_done_transfers`` to report done.
+                    self._recving_transfers.setdefault(req_id, [])
+                else:
+                    # Not tracked on either side (lease may have expired
+                    # before the notif arrived). Log and skip.
+                    logger.error(
+                        "Unrecognized request %s notif (may have expired).",
+                        req_id,
+                    )
+                continue
+
+            n_consumers = int(tp_size)
+            tp_ratio = self.transfer_topo.tp_ratio(n_consumers)
+            consumers_per_producer = -tp_ratio if n_consumers > self.world_size else 1
+            self.consumer_notification_counts_by_req[req_id] += 1
+            if (
+                self.consumer_notification_counts_by_req[req_id]
+                == consumers_per_producer
+            ):
+                notified_req_ids.add(req_id)
+                del self.consumer_notification_counts_by_req[req_id]
+                self._reqs_to_process.remove(req_id)
+                self._reqs_to_send.pop(req_id, None)
+        return notified_req_ids
+
+    def get_finished(self) -> tuple[set[str], set[str]]:
+        # Engine main thread asking for completions: also wake the writer
+        # so it gets a chance to drain NIXL notifs (heartbeats, completion
+        # notifs, late PUSH_REGs) even if it had been parked.
+        self._push_writer_wake.set()
+
+        done_sending, done_recving = super().get_finished()
+
+        # ``_pop_done_transfers`` mutates ``_sending_transfers``; the
+        # writer thread also appends to it, so guard the pop.
+        with self._sending_transfers_lock:
+            done_pushing = self._pop_done_transfers(self._sending_transfers)
+        for req_id in done_pushing:
+            self._reqs_to_send.pop(req_id, None)
+            self._reqs_to_process.discard(req_id)
+            self.consumer_notification_counts_by_req.pop(req_id, None)
+            done_sending.add(req_id)
+
+        # Tell the writer to drop any state it still holds for any
+        # request that just finished (push completed) or expired
+        # (lease ran out without a D registration ever arriving).
+        for req_id in done_sending:
+            self._evict_finished_inbox.put(req_id)
+        if done_sending:
+            self._push_writer_wake.set()
+
+        return done_sending, done_recving

vllm/distributed/kv_transfer/kv_connector/v1/nixl/scheduler.py

@@ -0,0 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Backward-compatible re-export of NixlPullConnectorScheduler."""
+
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import (
+    NixlPullConnectorScheduler,
+)
+
+# Backward compatibility: NixlConnectorScheduler is the pull-based scheduler.
+NixlConnectorScheduler = NixlPullConnectorScheduler
+
+__all__ = ["NixlConnectorScheduler", "NixlPullConnectorScheduler"]