[PD]适配v0.15.1

vllm/distributed/kv_transfer/kv_connector/factory.py

@@ -155,6 +155,11 @@ KVConnectorFactory.register_connector(
     "P2pNcclConnector",
 )
 
+KVConnectorFactory.register_connector(
+    "DuSwiftConnector",
+    "vllm.distributed.kv_transfer.kv_connector.v1.du.du_swift_connector",
+    "DuSwiftConnector")
+
 KVConnectorFactory.register_connector(
     "LMCacheConnectorV1",
     "vllm.distributed.kv_transfer.kv_connector.v1.lmcache_connector",

vllm/distributed/kv_transfer/kv_connector/v1/du/du_swift_connector.py

+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Optional
+
+import regex as re
+import torch
+import os
+from vllm import envs
+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.du.du_swift_engine import (
+    DuSwiftEngine, RemoteAddr)
+from vllm.distributed.parallel_state import get_world_group
+from vllm.forward_context import get_forward_context
+from vllm.logger import init_logger
+from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.distributed.parallel_state import get_pp_group, get_tp_group, get_dp_group
+
+if TYPE_CHECKING:
+    from vllm.v1.attention.backend import AttentionMetadata
+    from vllm.forward_context import ForwardContext
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+    from vllm.v1.request import Request
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class ReqMeta:
+    # Request Id
+    request_id: str
+    # Request tokens
+    token_ids: torch.Tensor
+    # Slot mappings, should have the same length as token_ids
+    slot_mapping: torch.Tensor
+    slot_mapping_device: torch.Tensor = None
+
+    @staticmethod
+    def make_meta(request_id: str, token_ids: list[int], block_ids: list[int],
+                  block_size: int) -> "ReqMeta":
+        valid_num_tokens = len(token_ids)
+        token_ids_tensor = torch.tensor(token_ids)
+        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(
+            request_id=request_id,
+            token_ids=token_ids_tensor,
+            slot_mapping=slot_mapping,
+        )
+
+        
+@dataclass
+class DuSwiftConnectorMetadata(KVConnectorMetadata):
+    requests: list[ReqMeta]
+
+    def __init__(self):
+        self.requests = []
+
+    def add_request(
+        self,
+        request_id: str,
+        token_ids: list[int],
+        block_ids: list[int],
+        block_size: int,
+    ) -> None:
+        self.requests.append(
+            ReqMeta.make_meta(request_id, token_ids, block_ids, block_size))
+
+
+class DuSwiftConnector(KVConnectorBase_V1):
+
+    def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole, kv_cache_config: Optional["KVCacheConfig"] = None):
+        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, Any] = {}
+        self.config = vllm_config.kv_transfer_config
+        self.is_producer = self.config.is_kv_producer
+        self.chunked_prefill: dict[str, Any] = {}
+
+        self._rank = get_world_group().rank \
+            if role == KVConnectorRole.WORKER else 0
+        self._local_rank = get_world_group().local_rank \
+            if role == KVConnectorRole.WORKER else 0
+        self._dp_rank = get_dp_group().rank_in_group \
+            if role == KVConnectorRole.WORKER else 0
+        self._pp_rank = get_pp_group().rank_in_group \
+            if role == KVConnectorRole.WORKER else 0
+        self._tp_rank = get_tp_group().rank_in_group \
+            if role == KVConnectorRole.WORKER else 0
+        self._dp_size = get_dp_group().world_size \
+            if role == KVConnectorRole.WORKER else 0
+        self._pp_size = get_pp_group().world_size \
+            if role == KVConnectorRole.WORKER else 0
+        self._tp_size = get_tp_group().world_size \
+            if role == KVConnectorRole.WORKER else 0
+
+        self.du_swift_engine = DuSwiftEngine(
+            local_rank=self._local_rank,
+            port_offset=self._rank,
+            config=self.config,
+            model_config=vllm_config.model_config,
+            dp_rank=self._dp_rank,
+            pp_rank=self._pp_rank,
+            tp_rank=self._tp_rank,
+            dp_size=self._dp_size,
+            pp_size=self._pp_size,
+            tp_size=self._tp_size
+        ) if role == KVConnectorRole.WORKER else None
+        
+        self.parallel_config = vllm_config.parallel_config
+        self.model_config = vllm_config.model_config
+        self.total_num_hidden_layers = getattr(self.model_config.hf_text_config,
+                                              "num_hidden_layers", 0)
+        self.pp_size = self.parallel_config.pipeline_parallel_size
+        self.tp_size = self.parallel_config.tensor_parallel_size
+        self.num_card = self.pp_size * self.tp_size
+
+        self.remote_tp_size = self.config.get_from_extra_config(
+            "remote_tp_size", self.tp_size)
+        self.remote_pp_size = self.config.get_from_extra_config(
+            "remote_pp_size", self.pp_size)
+        self.enable_asymmetric_p2p = self.config.get_from_extra_config(
+            "enable_asymmetric_p2p", False)
+        
+        self.remote_num_card = self.remote_tp_size * self.remote_pp_size
+        self.multiple_machines_d = 1 if self.remote_num_card > 8 else 0
+        self.multiple_machines_p = 1 if self.num_card > 8 else 0
+
+        if self.is_producer and self.multiple_machines_p == 1:
+            self.ip_map = {}
+            self.duplicate_keys = []
+            config_file = os.getenv('IP_CONFIG_FILE')
+            if not config_file:
+                print("Warning: Please set the IPVNet FILE environment variable for cross machine recognition of the second IP address")
+                return
+            try:
+                with open(config_file, 'r', encoding='utf-8') as file:
+                    for line_num, line in enumerate(file, 1):
+                        line = line.strip()
+                        if line and not line.startswith('#'):
+                            ips = line.split()
+                            if len(ips) == 2:
+                                first_ip, second_ip = ips
+                                if first_ip not in self.ip_map:
+                                    self.ip_map[first_ip] = second_ip
+                            else:
+                                print(f"warning: num {line_num} Incorrect format : {line}")
+            except Exception as e:
+                print(f"Error: Exception occurred while reading configuration file - {e}")
+
+
+    def get_ip_value(self, key):
+        return self.ip_map.get(key)
+        
+
+    # ==============================
+    # Worker-side methods
+    # ==============================
+
+    def start_load_kv(self, forward_context: "ForwardContext",
+                      **kwargs) -> 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.
+        """
+
+        # Only consumer/decode loads KV Cache
+        if self.is_producer:
+            return
+        assert self.du_swift_engine is not None
+
+        attn_metadata = forward_context.attn_metadata
+        if attn_metadata is None:
+            return
+        def inject_kv_into_layer(
+            dst_kv_cache_layer: torch.Tensor,
+            src_kv_cache: torch.Tensor,
+            slot_mapping: torch.Tensor,
+            request_id: str,
+        ) -> None:
+            """Inject the KV cache into the layer.
+
+            Args:
+                dst_kv_cache_layer (torch.Tensor): the destination KV cache
+                    layer. In shape [2, num_pages, page_size, xxx] if not
+                    using MLA, [num_pages, page_size, xxx] otherwise.
+                src_kv_cache (torch.Tensor): the source KV cache. In shape
+                    [2, num_tokens, xxx] if not using MLA, [num_tokens, xxx]
+                    otherwise.
+                slot_mapping (torch.Tensor): the slot mapping. In shape
+                    [num_tokens].
+                request_id (str): request id for log
+            """
+            dst_kv_cache_layer_shape = dst_kv_cache_layer.shape
+            if isinstance(attn_metadata, MLACommonMetadata) or all(isinstance(value, MLACommonMetadata) for value in attn_metadata.values()):
+                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)
+                self.check_tensors_except_dim(dst_kv_cache_layer, src_kv_cache,
+                                              0)
+                num_token = src_kv_cache.shape[0]
+                if len(slot_mapping) == num_token:
+                    dst_kv_cache_layer[slot_mapping, ...] = src_kv_cache
+                else:
+                    dst_kv_cache_layer[slot_mapping[:num_token],
+                                       ...] = src_kv_cache
+                    logger.warning(
+                        "🚧src_kv_cache does not match, num_slot:%d, "
+                        "num_token:%d, request_id:%s", len(slot_mapping),
+                        num_token, request_id)
+
+                dst_kv_cache_layer.reshape(dst_kv_cache_layer_shape)
+            else:
+                num_pages = dst_kv_cache_layer_shape[1]
+                page_size = dst_kv_cache_layer_shape[2]
+                dst_kv_cache_layer = dst_kv_cache_layer.reshape(
+                    2, num_pages * page_size, -1)
+                self.check_tensors_except_dim(dst_kv_cache_layer, src_kv_cache,
+                                              1)
+                num_token = src_kv_cache.shape[1]
+                if len(slot_mapping) == num_token:
+                    dst_kv_cache_layer[:, slot_mapping, ...] = src_kv_cache
+                else:
+                    dst_kv_cache_layer[:, slot_mapping[:num_token],
+                                       ...] = src_kv_cache
+                    logger.warning(
+                        "🚧src_kv_cache does not match, num_slot:%d, "
+                        "num_token:%d, request_id:%s", len(slot_mapping),
+                        num_token, request_id)
+
+                dst_kv_cache_layer.reshape(dst_kv_cache_layer_shape)
+
+        # Get the metadata
+        metadata: KVConnectorMetadata = \
+            self._get_connector_metadata()
+        assert isinstance(metadata, DuSwiftConnectorMetadata)
+
+        if metadata is None:
+            return
+
+        # Load the KV for each request each layer
+        for request in metadata.requests:
+            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
+                kv_cache = getattr(layer, 'kv_cache', None)
+                if kv_cache is None:
+                    continue
+
+                kv_cache_layer = kv_cache[ \
+                    forward_context.virtual_engine]
+
+                if not envs.VLLM_P2P_ASYNC:
+                    kv_cache = self.du_swift_engine.recv_tensor(
+                        request.request_id + "#" + layer_name)
+
+                    if kv_cache is None:
+                        logger.warning("🚧src_kv_cache is None, %s",
+                                    request.request_id)
+                        continue
+                    inject_kv_into_layer(kv_cache_layer, kv_cache,
+                                        request.slot_mapping, request.request_id)
+                    tensor_id = request.request_id + "#" + layer_name
+                    if tensor_id in self.du_swift_engine.recv_store:
+                        tensor = self.du_swift_engine.recv_store.pop(tensor_id, None)
+                        self.du_swift_engine.send_request_id_to_tensor_ids.pop(
+                                    request.request_id, None)
+                        self.du_swift_engine.recv_request_id_to_tensor_ids.pop(
+                                    request.request_id, None)
+                        addr = 0
+                        if isinstance(tensor, tuple):
+                            addr, _, _ = tensor
+                            self.du_swift_engine.pool.free(addr)
+                else:
+                    dst_kv_cache_layer_shape = kv_cache_layer.shape
+                    if isinstance(attn_metadata, MLACommonMetadata) or all(isinstance(value, MLACommonMetadata) for value in attn_metadata.values()):
+                        num_pages = dst_kv_cache_layer_shape[0]
+                        page_size = dst_kv_cache_layer_shape[1]
+                        assert kv_cache_layer.is_contiguous()
+                        dst_kv_cache_layer = kv_cache_layer.reshape(
+                            num_pages * page_size, -1)
+                    else:
+                        num_pages = dst_kv_cache_layer_shape[1]
+                        page_size = dst_kv_cache_layer_shape[2]
+                        assert kv_cache_layer.is_contiguous()
+                        dst_kv_cache_layer = kv_cache_layer.reshape(
+                            2, num_pages * page_size, -1)
+
+                    inject_start_index = 0
+                    for num in range(self.du_swift_engine.tensor_split_num):
+                        kv_cache = self.du_swift_engine.recv_tensor(
+                            request.request_id + "#" + layer_name + "#" + str(num))
+
+                        if kv_cache is None:
+                            logger.warning("🚧src_kv_cache is None, %s",
+                                        request.request_id)
+                            continue
+                        
+                        if isinstance(attn_metadata, MLACommonMetadata) or all(isinstance(value, MLACommonMetadata) for value in attn_metadata.values()):
+                            num_token = kv_cache.shape[0]
+                            if len(request.slot_mapping) == num_token:
+                                dst_kv_cache_layer[request.slot_mapping, ...] = kv_cache
+                            else:
+                                dst_kv_cache_layer[request.slot_mapping[inject_start_index:inject_start_index + num_token],
+                                                ...] = kv_cache
+                        else:
+                            num_token = kv_cache.shape[1]
+                            if len(request.slot_mapping) == num_token:
+                                dst_kv_cache_layer[:, request.slot_mapping, ...] = kv_cache
+                            else:
+                                dst_kv_cache_layer[:, request.slot_mapping[inject_start_index:inject_start_index + num_token],
+                                                ...] = kv_cache
+                        inject_start_index += num_token
+                        
+                        # inject_kv_into_layer(kv_cache_layer, kv_cache,
+                        #                     request.slot_mapping, request.request_id)
+                        tensor_id = request.request_id + "#" + layer_name + "#" + str(num)
+                        if tensor_id in self.du_swift_engine.recv_store:
+                            tensor = self.du_swift_engine.recv_store.pop(tensor_id, None)
+                            self.du_swift_engine.send_request_id_to_tensor_ids.pop(
+                                        request.request_id, None)
+                            self.du_swift_engine.recv_request_id_to_tensor_ids.pop(
+                                        request.request_id, None)
+                            addr = 0
+                            if isinstance(tensor, tuple):
+                                addr, _, _ = tensor
+                                self.du_swift_engine.pool.free(addr)
+                    dst_kv_cache_layer.reshape(dst_kv_cache_layer_shape)
+
+
+    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) -> 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.
+        """
+
+        # Only producer/prefill saves KV Cache
+        if not self.is_producer:
+            return
+
+        assert self.du_swift_engine is not None
+
+        is_mla = isinstance(attn_metadata, MLACommonMetadata)
+
+        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 (2, num_pages, page_size, xxx)
+            if MLA is not used, and (num_pages, 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,
+                                                                ...]
+            num_pages, page_size = layer.shape[1], layer.shape[2]
+            return layer.reshape(2, num_pages * page_size, -1)[:, slot_mapping,
+                                                               ...]
+
+        connector_metadata = self._get_connector_metadata()
+        assert isinstance(connector_metadata, DuSwiftConnectorMetadata)
+
+        if envs.VLLM_ENABLE_TBO or envs.VLLM_P2P_ASYNC:
+            for request in connector_metadata.requests:
+                request_id = request.request_id
+                ip, port = self.parse_request_id(request_id, True)
+                remote_address = ip + ":" + str(port + self._rank)
+                slot_mapping = request.slot_mapping
+                if request.slot_mapping_device is None:
+                    request.slot_mapping_device = \
+                        request.slot_mapping.pin_memory().to(device=kv_layer.device, non_blocking=True)
+                slot_mapping = request.slot_mapping_device
+                tbo_evt = torch.cuda.Event(enable_timing=False)
+                tbo_evt.record()
+                pp_rank = (self.parallel_config.rank //
+                        self.parallel_config.tensor_parallel_size) % \
+                        self.parallel_config.pipeline_parallel_size
+                if (self.pp_size == 1):
+                    self.du_swift_engine.p2p_async_send_tensor(request_id + "#" + layer_name,
+                                                    (kv_layer, slot_mapping), remote_address, tbo_evt)
+                elif (self.pp_size == 2):
+                    if (pp_rank == 0):
+                        self.du_swift_engine.p2p_async_send_tensor(request_id + "#" + layer_name,
+                                                        (kv_layer, slot_mapping), remote_address, tbo_evt)
+                        self.du_swift_engine.p2p_async_send_tensor(request_id + "#" + layer_name,
+                                                        (kv_layer, slot_mapping), ip + ":" + str(port + self._rank + 4), tbo_evt)
+                    else:
+                        self.du_swift_engine.p2p_async_send_tensor(request_id + "#" + layer_name,
+                                                        (kv_layer, slot_mapping), remote_address, tbo_evt)
+                        self.du_swift_engine.p2p_async_send_tensor(request_id + "#" + layer_name,
+                                                        (kv_layer, slot_mapping), ip + ":" + str(port + self._rank - 4), tbo_evt)
+                elif (self.pp_size == 8):
+                    for i in range(8):
+                        self.du_swift_engine.p2p_async_send_tensor(request_id + "#" + layer_name,
+                                                        (kv_layer, slot_mapping), ip + ":" + str(port + i), tbo_evt)
+                else:
+                    print("Error: only suppprt pp1 pp2 pp8!!!!!!")
+        else:
+            for request in connector_metadata.requests:
+                request_id = request.request_id
+                ip, port = self.parse_request_id(request_id, True)
+                p_ip, p_port = self.parse_request_id(request_id, False)
+                remote_address = ip + ":" + str(port + self._rank)
+                # pd_pair_id = p_ip + ":" + p_port + "_" + ip + ":" + port
+                kv_cache = extract_kv_from_layer(kv_layer, request.slot_mapping)
+
+                pp_rank = (self.parallel_config.rank // self.parallel_config.tensor_parallel_size
+                    ) % self.parallel_config.pipeline_parallel_size
+                if (self.multiple_machines_p and self.multiple_machines_d):
+                    ip_second = self.get_ip_value(ip)
+                    if (self.pp_size == 1):
+                        if self._rank < 8:
+                            self.du_swift_engine.send_tensor(request_id + "#" + layer_name,
+                                                        kv_cache, remote_address)
+                            self.du_swift_engine.send_tensor(request_id + "#" + layer_name,
+                                                        kv_cache, str(ip_second) + ":" + str(port + self._rank + 8))
+                    elif (self.pp_size == 2):
+                        if (pp_rank == 0):
+                            self.du_swift_engine.send_tensor(request_id + "#" + layer_name,
+                                                    kv_cache, remote_address)
+                        else:
+                            self.du_swift_engine.send_tensor(request_id + "#" + layer_name,
+                                                        kv_cache, str(ip_second) + ":" + str(port + self._rank))
+                    else:
+                        logger.error("Error: multiple machines only suppprt pp1tp16 and pp2tp8!!!!!!")
+                elif (self.multiple_machines_p and not self.multiple_machines_d):
+                    if (self.pp_size == 2):
+                        remote_address = ip + ":" + str(port + self._tp_rank)
+                        self.du_swift_engine.send_tensor(request_id + "#" + layer_name,
+                                                        kv_cache, remote_address)
+                    else:
+                        logger.error("Error: P multiple machines D machine only suppprt P:pp2tp8 D:tp8 !!!!!!")
+
+                elif (not self.multiple_machines_p and not self.multiple_machines_d):
+                    # remote_addr = RemoteAddr(pd_pair_id, remote_address, self._rank + self.num_card)
+                    self.du_swift_engine.send_tensor_new(request_id, layer_name, kv_cache,
+                                                is_mla)
+                    # if (self.pp_size == 1):
+                    #     self.du_swift_engine.send_tensor(request_id + "#" + layer_name,
+                    #                                 kv_cache, remote_address)
+                    # elif (self.pp_size == 2):
+                    #     if (pp_rank == 0):
+                    #         self.du_swift_engine.send_tensor(request_id + "#" + layer_name,
+                    #                                     kv_cache, remote_address)
+                    #         self.du_swift_engine.send_tensor(request_id + "#" + layer_name,
+                    #                                     kv_cache, ip + ":" + str(port + self._rank + 4))
+                    #     else:
+                    #         self.du_swift_engine.send_tensor(request_id + "#" + layer_name,
+                    #                                     kv_cache, remote_address)
+                    #         self.du_swift_engine.send_tensor(request_id + "#" + layer_name,
+                    #                                     kv_cache, ip + ":" + str(port + self._rank - 4))
+                    # elif (self.pp_size == 8):
+                    #     for i in range(8):
+                    #         self.du_swift_engine.send_tensor(request_id + "#" + layer_name,
+                    #                                     kv_cache, ip + ":" + str(port + i))
+                    # elif (self.enable_asymmetric_p2p):
+                    #     self.du_swift_engine.send_tensor(request_id + "#" + layer_name,
+                    #                                 kv_cache, remote_address)
+                    # else:
+                    #     logger.error("Error: P/D single machine only suppprt multiple tp:: (P: pp2tp4 D:tp8 P:pp8tp1 D:tp8) !!!!!!")
+                else:
+                        logger.error("Error: not support!!!!!!")
+    def wait_for_save(self):
+        pass
+        # if self.is_producer:
+        #     assert self.du_swift_engine is not None
+        #     self.du_swift_engine.wait_for_sent()
+
+    def get_finished(
+            self, finished_req_ids: set[str],
+            **kwargs) -> tuple[Optional[set[str]], Optional[set[str]]]:
+        """
+        Notifies worker-side connector ids of requests that have
+        finished generating tokens.
+
+        Returns:
+            ids of requests that have finished asynchronous transfer,
+            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).
+        """
+
+        assert self.du_swift_engine is not None
+
+        forward_context: ForwardContext = get_forward_context()
+        return self.du_swift_engine.get_finished(finished_req_ids,
+                                                 forward_context)
+
+    # ==============================
+    # Scheduler-side methods
+    # ==============================
+
+    def get_num_new_matched_tokens(
+        self,
+        request: "Request",
+        num_computed_tokens: int,
+    ) -> tuple[int, 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.
+        """
+        if self.is_producer:
+            return 0, False
+
+        num_external_tokens = (len(request.prompt_token_ids) - 1 -
+                               num_computed_tokens)
+
+        if num_external_tokens < 0:
+            num_external_tokens = 0
+
+        return num_external_tokens, False
+
+    def update_state_after_alloc(self, request: "Request",
+                                 blocks: "KVCacheBlocks",
+                                 num_external_tokens: int):
+        """
+        Update KVConnector state after block allocation.
+        """
+        if not self.is_producer and num_external_tokens > 0:
+            self._requests_need_load[request.request_id] = (
+                request, blocks.get_block_ids()[0])
+
+    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 = DuSwiftConnectorMetadata()
+
+        for new_req in scheduler_output.scheduled_new_reqs:
+            if self.is_producer:
+                num_scheduled_tokens = (
+                    scheduler_output.num_scheduled_tokens)[new_req.req_id]
+                num_tokens = num_scheduled_tokens + new_req.num_computed_tokens
+                # the request's prompt is chunked prefill
+                if num_tokens < len(new_req.prompt_token_ids):
+                    # 'CachedRequestData' has no attribute 'prompt_token_ids'
+                    self.chunked_prefill[new_req.req_id] = (
+                        new_req.block_ids[0], new_req.prompt_token_ids)
+                    continue
+                # the request's prompt is not chunked prefill
+                meta.add_request(request_id=new_req.req_id,
+                                 token_ids=new_req.prompt_token_ids,
+                                 block_ids=new_req.block_ids[0],
+                                 block_size=self._block_size)
+                continue
+            if new_req.req_id in self._requests_need_load:
+                meta.add_request(request_id=new_req.req_id,
+                                 token_ids=new_req.prompt_token_ids,
+                                 block_ids=new_req.block_ids[0],
+                                 block_size=self._block_size)
+                self._requests_need_load.pop(new_req.req_id)
+
+        cached_reqs = scheduler_output.scheduled_cached_reqs
+        for i, req_id in enumerate(cached_reqs.req_ids):
+            num_computed_tokens = cached_reqs.num_computed_tokens[i]
+            new_block_ids = cached_reqs.new_block_ids[i]
+            resumed_from_preemption = cached_reqs.resumed_from_preemption[i]
+
+            if self.is_producer:
+                num_scheduled_tokens = (
+                    scheduler_output.num_scheduled_tokens)[req_id]
+                num_tokens = (num_scheduled_tokens + num_computed_tokens)
+                # assert req_id in self.chunked_prefill
+                if  req_id not in self.chunked_prefill:
+                    continue
+                block_ids = new_block_ids[0]
+                if not resumed_from_preemption:
+                    block_ids = (self.chunked_prefill[req_id][0] + block_ids)
+                prompt_token_ids = self.chunked_prefill[req_id][1]
+                # the request's prompt is chunked prefill again
+                if num_tokens < len(prompt_token_ids):
+                    self.chunked_prefill[req_id] = (block_ids,
+                                                    prompt_token_ids)
+                    continue
+                # the request's prompt is all prefilled finally
+                meta.add_request(request_id=req_id,
+                                 token_ids=prompt_token_ids,
+                                 block_ids=block_ids,
+                                 block_size=self._block_size)
+                self.chunked_prefill.pop(req_id, None)
+                continue
+
+            # NOTE(rob): here we rely on the resumed requests being
+            # the first N requests in the list scheduled_cache_reqs.
+            if not resumed_from_preemption:
+                break
+            if req_id in self._requests_need_load:
+                request, _ = self._requests_need_load.pop(req_id)
+                total_tokens = num_computed_tokens + 1
+                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.
+                block_ids = new_block_ids[0]
+
+                meta.add_request(request_id=req_id,
+                                 token_ids=token_ids,
+                                 block_ids=block_ids,
+                                 block_size=self._block_size)
+
+        # Requests loaded asynchronously are not in the scheduler_output.
+        # for request_id in self._requests_need_load:
+        #     request, block_ids = self._requests_need_load[request_id]
+        #     meta.add_request(request_id=request.request_id,
+        #                      token_ids=request.prompt_token_ids,
+        #                      block_ids=block_ids,
+        #                      block_size=self._block_size)
+
+        self._requests_need_load.clear()
+        return meta
+
+    def request_finished(
+        self,
+        request: "Request",
+        block_ids: list[int],
+    ) -> tuple[bool, Optional[dict[str, Any]]]:
+        """
+        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.
+        """
+
+        self.chunked_prefill.pop(request.request_id, None)
+
+        return False, None
+
+    # ==============================
+    # Static methods
+    # ==============================
+
+    @staticmethod
+    def parse_request_id(request_id: str, is_prefill=True) -> tuple[str, int]:
+        # Regular expression to match the string hostname and integer port
+        if is_prefill:
+            pattern = r"___decode_addr_(.*):(\d+)"
+        else:
+            pattern = r"___prefill_addr_(.*):(\d+)___"
+
+        # Use re.search to find the pattern in the request_id
+        match = re.search(pattern, request_id)
+        if match:
+            # Extract the ranks
+            ip = match.group(1)
+            port = int(match.group(2))
+
+            return ip, port
+        raise ValueError(
+            f"Request id {request_id} does not contain hostname and port")
+
+    @staticmethod
+    def check_tensors_except_dim(tensor1, tensor2, dim):
+        shape1 = tensor1.size()
+        shape2 = tensor2.size()
+
+        if len(shape1) != len(shape2) or not all(
+                s1 == s2
+                for i, (s1, s2) in enumerate(zip(shape1, shape2)) if i != dim):
+            raise NotImplementedError(
+                "Currently, only symmetric TP is supported. Asymmetric TP, PP,"
+                "and others will be supported in future PRs.")

vllm/distributed/kv_transfer/kv_connector/v1/du/du_swift_engine.py

+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import logging
+import os
+import threading
+import time
+import typing
+from collections import deque
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Optional
+
+import msgpack
+import torch
+import zmq
+import regex
+
+from vllm.config import KVTransferConfig
+from vllm.distributed.device_communicators.pynccl_wrapper import (
+    NCCLLibrary, buffer_type, cudaStream_t, ncclComm_t, ncclDataTypeEnum)
+from vllm.distributed.kv_transfer.kv_connector.v1.du.tensor_memory_pool import (  # noqa: E501
+    TensorMemoryPool)
+from vllm.utils.torch_utils import current_stream
+from vllm.utils.network_utils import get_ip
+from vllm import envs
+from vllm.distributed.parallel_state import get_pp_group, get_tp_group
+
+from dataclasses import dataclass
+from vllm.model_executor.models.utils import extract_layer_index
+from vllm.distributed.utils import get_pp_indices
+from vllm.config import ModelConfig
+
+if TYPE_CHECKING:
+    from vllm.forward_context import ForwardContext
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_MEM_POOL_SIZE_GB = 32
+
+# @dataclass
+# class SendQueueItem:
+#     tensor_id: str
+#     remote_address: str
+#     tensor: torch.Tensor
+
+@contextmanager
+def set_du_swift_context(num_channels: str):
+    original_values: dict[str, Any] = {}
+    env_vars = [
+        'NCCL_MAX_NCHANNELS',
+        'NCCL_MIN_NCHANNELS',
+        'NCCL_CUMEM_ENABLE',
+        'NCCL_BUFFSIZE',
+        'NCCL_PROTO',  # LL,LL128,SIMPLE
+        'NCCL_ALGO',  # RING,TREE
+    ]
+
+    for var in env_vars:
+        original_values[var] = os.environ.get(var)
+
+    logger.info("set_du_swift_context, original_values: %s", original_values)
+
+    try:
+        os.environ['NCCL_MAX_NCHANNELS'] = num_channels
+        os.environ['NCCL_MIN_NCHANNELS'] = num_channels
+        os.environ['NCCL_CUMEM_ENABLE'] = '1'
+        yield
+    finally:
+        for var in env_vars:
+            if original_values[var] is not None:
+                os.environ[var] = original_values[var]
+            else:
+                os.environ.pop(var, None)
+
+
+@dataclass
+class RemoteAddr:
+    pd_pair_id: str = ""
+    zmq_address: str = ""
+    comm_rank: int = 0
+
+
+class DuSwiftEngine:
+
+    def __init__(self,
+                 local_rank: int,
+                 port_offset: int,
+                 config: KVTransferConfig,
+                 model_config: ModelConfig,
+                 dp_rank: int = 0,
+                 pp_rank: int = 0,
+                 tp_rank: int = 0,
+                 dp_size: int = 0,
+                 pp_size: int = 0,
+                 tp_size: int = 0,
+                 library_path: Optional[str] = None) -> None:
+        self.config = config
+        self.model_config = model_config
+        self.rank = port_offset
+        self.local_rank = local_rank
+        self.dp_rank = dp_rank
+        self.pp_rank = pp_rank
+        self.tp_rank = tp_rank
+        self.dp_size = dp_size
+        self.pp_size = pp_size
+        self.tp_size = tp_size
+        self.device = torch.device(f"cuda:{self.local_rank}")
+        self.nccl = NCCLLibrary(library_path)
+
+        self.total_num_hidden_layers = getattr(self.model_config.hf_text_config,
+                                              "num_hidden_layers", 0)
+        self.pp_rank = get_pp_group().rank_in_group
+        self.tp_rank = get_tp_group().rank_in_group
+        self.pp_size = get_pp_group().world_size 
+        self.tp_size = get_tp_group().world_size
+        if config.is_kv_producer:
+            self.remote_tp_size = self.config.get_from_extra_config(
+                "remote_tp_size", 1)
+            self.remote_pp_size = self.config.get_from_extra_config(
+                "remote_pp_size", 1)
+            self.enable_asymmetric_p2p = self.config.get_from_extra_config(
+                "enable_asymmetric_p2p", False)
+
+            if self.remote_tp_size % self.tp_size != 0:
+                logger.error(" the Prefill TP size must be less than or equal to the Decode TP size!!!!")
+            self.multp = int(self.remote_tp_size / self.tp_size)
+        self.multiple_machines = self.config.get_from_extra_config(
+                "enable_multiple_machines", False)
+        port = int(self.config.kv_port) + port_offset
+        if port == 0:
+            raise ValueError("Port cannot be 0")
+        self._hostname = get_ip()
+        self._port = port
+
+        # Each card corresponds to a ZMQ address.
+        self.zmq_address = f"{self._hostname}:{self._port}"
+
+        # The `http_port` must be consistent with the port of OpenAI.
+        self.http_address = (
+            f"{self._hostname}:"
+            f"{self.config.kv_connector_extra_config['http_port']}")
+
+        # If `proxy_ip` or `proxy_port` is `""`,
+        # then the ping thread will not be enabled.
+        proxy_ip = self.config.get_from_extra_config("proxy_ip", "")
+        proxy_port = self.config.get_from_extra_config("proxy_port", "")
+        if proxy_ip == "" or proxy_port == "":
+            self.proxy_address = ""
+        else:
+            self.proxy_address = proxy_ip + ":" + proxy_port
+
+        self.context = zmq.Context()
+        self.router_socket = self.context.socket(zmq.ROUTER)
+        self.router_socket.bind(f"tcp://{self.zmq_address}")
+
+        self.poller = zmq.Poller()
+        self.poller.register(self.router_socket, zmq.POLLIN)
+
+        self.send_store_cv = threading.Condition()
+        self.send_queue_cv = threading.Condition()
+        self.recv_store_cv = threading.Condition()
+
+        self.send_stream = torch.cuda.Stream()
+        self.recv_stream = torch.cuda.Stream()
+        
+        self.p2p_async_kv_tokens = envs.VLLM_P2P_BUF_TOKENS
+        self.p2p_async_buf = None
+        self.tensor_split_num: int = 0
+
+        mem_pool_size_gb = self.config.get_from_extra_config(
+            "mem_pool_size_gb", DEFAULT_MEM_POOL_SIZE_GB)
+        self.pool = TensorMemoryPool(max_block_size=int(mem_pool_size_gb) *
+                                     1024**3)  # GB
+
+        # The sending type includes tree mutually exclusive options:
+        # PUT, GET, PUT_ASYNC.
+        self.send_type = self.config.get_from_extra_config("send_type", "PUT")
+        if self.send_type == "GET":
+            # tensor_id: torch.Tensor
+            self.send_store: dict[str, torch.Tensor] = {}
+        else:
+            # PUT or PUT_ASYNC
+            # tensor_id: torch.Tensor
+            self.send_queue: deque[list[Any]] = deque()
+            self.send_request_id_to_tensor_ids: dict[str, set[str]] = {}
+            if self.send_type == "PUT_ASYNC":
+                self._send_thread = threading.Thread(target=self._send_async,
+                                                     daemon=True)
+                self._send_thread.start()
+
+        # tensor_id: torch.Tensor/(addr, dtype, shape)
+        self.recv_store: dict[str, Any] = {}
+        self.recv_request_id_to_tensor_ids: dict[str, set[str]] = {}
+        self.socks: dict[str, Any] = {}  # remote_address: client socket
+        self.comms: dict[str, Any] = {}  # remote_address: (ncclComm_t, rank)
+
+        self.buffer_size = 0
+        self.buffer_size_threshold = float(self.config.kv_buffer_size)
+
+        self.nccl_num_channels = self.config.get_from_extra_config(
+            "nccl_num_channels", "8")
+
+        self._listener_thread = threading.Thread(
+            target=self._listen_for_requests, daemon=True)
+        self._listener_thread.start()
+
+        self._ping_thread = None
+        if self.multiple_machines:
+            if port_offset == 0 and self.proxy_address != "":
+                self._ping_thread = threading.Thread(target=self._ping,
+                                                    daemon=True)
+                self._ping_thread.start()
+        else:
+            if self.proxy_address != "":
+                self._ping_thread = threading.Thread(target=self._ping_new,
+                                                    daemon=True)
+                self._ping_thread.start()
+        logger.info(
+            "💯DuSwiftEngine init, rank:%d, local_rank:%d, http_address:%s, "
+            "zmq_address:%s, proxy_address:%s, send_type:%s, buffer_size_"
+            "threshold:%.2f, nccl_num_channels:%s", self.rank, self.local_rank,
+            self.http_address, self.zmq_address, self.proxy_address,
+            self.send_type, self.buffer_size_threshold, self.nccl_num_channels)
+
+    def _create_connect_new(self, remote_address: typing.Optional[str] = None):
+        assert remote_address is not None
+        
+        if remote_address not in self.socks:
+            sock = self.context.socket(zmq.DEALER)
+            sock.setsockopt(zmq.SNDHWM, 10000)
+            sock.setsockopt(zmq.RCVHWM, 5000) 
+            sock.setsockopt(zmq.LINGER, 0)  
+            sock.setsockopt(zmq.TCP_KEEPALIVE, 1)
+            sock.setsockopt_string(zmq.IDENTITY, f"P-{self.zmq_address}")
+            sock.connect(f"tcp://{remote_address}")
+            self.socks[remote_address] = sock
+
+        return self.socks[remote_address]
+
+    def _create_connect(self, remote_address: typing.Optional[str] = None):
+        assert remote_address is not None
+        if remote_address not in self.socks:
+            sock = self.context.socket(zmq.DEALER)
+            sock.setsockopt_string(zmq.IDENTITY, self.zmq_address)
+            sock.connect(f"tcp://{remote_address}")
+            self.socks[remote_address] = sock
+            if remote_address in self.comms:
+                logger.info("👋comm exists, remote_address:%s, comms:%s",
+                            remote_address, self.comms)
+                return sock, self.comms[remote_address]
+
+            unique_id = self.nccl.ncclGetUniqueId()
+            data = {"cmd": "NEW", "unique_id": bytes(unique_id.internal)}
+            sock.send(msgpack.dumps(data))
+
+            with torch.cuda.device(self.device):
+                rank = 0
+                with set_du_swift_context(self.nccl_num_channels):
+                    comm: ncclComm_t = self.nccl.ncclCommInitRank(
+                        2, unique_id, rank)
+                self.comms[remote_address] = (comm, rank)
+                logger.info("🤝ncclCommInitRank Success, %s👉%s, MyRank: %s",
+                            self.zmq_address, remote_address, rank)
+
+        return self.socks[remote_address], self.comms[remote_address]
+
+    def get_send_queue_items(self, request_id: str, layer_name: str,
+                             tensor: torch.Tensor,
+                             is_mla: bool) -> list[any]:
+        tensor_id = self.get_tensor_id(request_id, layer_name)
+        remote_ip, remote_port = self.parse_request_id(request_id, True)
+        
+        p_ip, p_port = self.parse_request_id(request_id, False)
+        pd_pair_id = p_ip + ":" + str(p_port) + "_" + remote_ip + ":" + str(remote_port)
+
+
+        if not self.enable_asymmetric_p2p:
+                remote_address = remote_ip + ":" + str(remote_port + self.rank)
+                remote_addr = RemoteAddr(pd_pair_id, remote_address, self.rank + self.pp_size * self.tp_size)
+                # logger.info(f"""+++++xiabo tensor_id:{tensor_id} request_id:{request_id} remote_address:{remote_address}""")
+                return [(tensor_id, remote_addr, tensor)]
+        
+        if not is_mla:
+            logger.error(" DuSwift only support mla model symmetric PP/TP!!!!")
+        
+        remote_pp_rank = self.compute_remote_pp_rank(layer_name)
+        items: list[Any] = []
+
+        for d_tp_rank in range(self.remote_tp_size):
+            for mul_tp in range(self.multp):
+                if self.tp_rank + mul_tp * self.tp_size == d_tp_rank:
+                    remote_port_offset = remote_pp_rank * self.remote_tp_size + d_tp_rank
+                    remote_address = remote_ip + ":" + str(remote_port + remote_port_offset)
+                    remote_addr = RemoteAddr(pd_pair_id, remote_address, remote_port_offset + self.pp_size * self.tp_size)
+                    logger.debug(
+                        "Wait to send::%s, tensor_shape:%s, "
+                        "(pp=%d, tp=%d) -> remote_address=%s(pp=%d, tp=%d) comm_rank (%d -> %d)", tensor_id,
+                        tensor.shape, self.pp_rank, self.tp_rank, remote_address,
+                        remote_pp_rank, self.rank * mul_tp + self.rank, self.rank, remote_port_offset + self.pp_size * self.tp_size)
+                    
+                    items.append([tensor_id, remote_addr, tensor])
+        return items
+        
+    def send_tensor_new(
+        self,
+        request_id: str,
+        layer_name: str,
+        tensor: torch.Tensor,
+        is_mla: bool = False,
+    ) -> bool:
+        tensor_id = self.get_tensor_id(request_id, layer_name)
+
+        if self.send_type == "PUT":
+            return all(
+                self._send_sync_new(item) for item in self.get_send_queue_items(
+                    request_id, layer_name, tensor, is_mla))
+
+        if self.send_type == "PUT_ASYNC":
+            with self.send_queue_cv:
+                for item in self.get_send_queue_items(request_id, layer_name,
+                                                      tensor, is_mla):
+                    self.send_queue.append(item)
+                self.send_queue_cv.notify()
+            return True
+        if self.send_type == "GET":
+            logger.error(" DuSwift new not support GET model, please set VLLM_P2PNCCL_NEW=0 use defalut model!!!!")
+
+    def send_tensor(
+        self,
+        tensor_id: str,
+        tensor: torch.Tensor,
+        remote_address: typing.Optional[str] = None,
+        tbo_evt = None,
+    ) -> bool:
+        if remote_address is None:
+            with self.recv_store_cv:
+                self.recv_store[tensor_id] = tensor
+                self.recv_store_cv.notify()
+            return True
+        else:
+            if self.send_type == "PUT":
+                return self._send_sync(tensor_id, tensor, remote_address)
+            elif self.send_type == "PUT_ASYNC":
+                with self.send_queue_cv:
+                    self.send_queue.append([tensor_id, remote_address, tensor])
+                    self.send_queue_cv.notify()
+            else:  # GET
+                with self.send_store_cv:
+                    tensor_size = tensor.element_size() * tensor.numel()
+                    while (self.buffer_size + tensor_size
+                           > self.buffer_size_threshold):
+                        oldest_tenser_id = next(iter(self.send_store))
+                        oldest_tenser = self.send_store.pop(oldest_tenser_id)
+                        oldest_tenser_size = oldest_tenser.element_size(
+                        ) * oldest_tenser.numel()
+                        self.buffer_size -= oldest_tenser_size
+                        logger.info(
+                            "⛔[GET]Send to %s, tensor_id:%s, tensor_size:%d,"
+                            " buffer_size:%d, oldest_tenser_size:%d, rank:%d",
+                            remote_address, tensor_id, tensor_size,
+                            self.buffer_size, oldest_tenser_size, self.rank)
+
+                    self.send_store[tensor_id] = tensor
+                    self.buffer_size += tensor_size
+                    logger.debug(
+                        "🔵[GET]Send to %s, tensor_id:%s, tensor_size:%d, "
+                        "shape:%s, rank:%d, buffer_size:%d(%.2f%%)",
+                        remote_address, tensor_id, tensor_size, tensor.shape,
+                        self.rank, self.buffer_size,
+                        self.buffer_size / self.buffer_size_threshold * 100)
+
+        return True
+    
+    def p2p_async_send_tensor(
+        self,
+        tensor_id: str,
+        tensor: torch.Tensor,
+        remote_address: typing.Optional[str] = None,
+        tbo_evt = None,
+    ) -> bool:
+        if remote_address is None:
+            with self.recv_store_cv:
+                self.recv_store[tensor_id] = tensor
+                self.recv_store_cv.notify()
+            return True
+        else:
+            if self.send_type == "PUT":
+                return self._send_sync(tensor_id, tensor, remote_address)
+            elif self.send_type == "PUT_ASYNC":
+                with self.send_queue_cv:
+                    kv_layer, slot_mapping = tensor  #  tesor (kv_layer, slot_mapping)
+                    self.send_queue.append([tensor_id, remote_address, kv_layer, slot_mapping, tbo_evt])
+                    self.send_queue_cv.notify()
+            else:  # GET
+                with self.send_store_cv:
+                    tensor_size = tensor.element_size() * tensor.numel()
+                    while (self.buffer_size + tensor_size
+                           > self.buffer_size_threshold):
+                        oldest_tenser_id = next(iter(self.send_store))
+                        oldest_tenser = self.send_store.pop(oldest_tenser_id)
+                        oldest_tenser_size = oldest_tenser.element_size(
+                        ) * oldest_tenser.numel()
+                        self.buffer_size -= oldest_tenser_size
+                        logger.info(
+                            "⛔[GET]Send to %s, tensor_id:%s, tensor_size:%d,"
+                            " buffer_size:%d, oldest_tenser_size:%d, rank:%d",
+                            remote_address, tensor_id, tensor_size,
+                            self.buffer_size, oldest_tenser_size, self.rank)
+
+                    self.send_store[tensor_id] = tensor
+                    self.buffer_size += tensor_size
+                    logger.debug(
+                        "🔵[GET]Send to %s, tensor_id:%s, tensor_size:%d, "
+                        "shape:%s, rank:%d, buffer_size:%d(%.2f%%)",
+                        remote_address, tensor_id, tensor_size, tensor.shape,
+                        self.rank, self.buffer_size,
+                        self.buffer_size / self.buffer_size_threshold * 100)
+
+        return True
+
+    def recv_tensor(
+        self,
+        tensor_id: str,
+        remote_address: typing.Optional[str] = None,
+    ) -> torch.Tensor:
+        if self.send_type == "PUT" or self.send_type == "PUT_ASYNC":
+            start_time = time.time()
+            with self.recv_store_cv:
+                while tensor_id not in self.recv_store:
+                    self.recv_store_cv.wait()
+                tensor = self.recv_store[tensor_id]
+
+            if tensor is not None:
+                if isinstance(tensor, tuple):
+                    addr, dtype, shape = tensor
+                    tensor = self.pool.load_tensor(addr, dtype, shape,
+                                                   self.device)
+                else:
+                    self.buffer_size -= (tensor.element_size() *
+                                         tensor.numel())
+            else:
+                duration = time.time() - start_time
+                logger.warning(
+                    "🔴[PUT]Recv From %s, tensor_id:%s, duration:%.3fms, "
+                    "rank:%d", remote_address, tensor_id, duration * 1000,
+                    self.rank)
+            return tensor
+
+        # GET
+        if remote_address is None:
+            return None
+
+        if remote_address not in self.socks:
+            self._create_connect(remote_address)
+
+        sock = self.socks[remote_address]
+        comm, rank = self.comms[remote_address]
+
+        data = {"cmd": "GET", "tensor_id": tensor_id}
+        sock.send(msgpack.dumps(data))
+
+        message = sock.recv()
+        data = msgpack.loads(message)
+        if data["ret"] != 0:
+            logger.warning("🔴[GET]Recv From %s, tensor_id: %s, ret: %d",
+                           remote_address, tensor_id, data["ret"])
+            return None
+
+        tensor = torch.empty(data["shape"],
+                             dtype=getattr(torch, data["dtype"]),
+                             device=self.device)
+
+        self._recv(comm, tensor, rank ^ 1, self.recv_stream)
+
+        return tensor
+
+    def _listen_for_requests(self):
+        while True:
+            socks = dict(self.poller.poll())
+            if self.router_socket in socks:
+                remote_address, message = self.router_socket.recv_multipart()
+                data = msgpack.loads(message)
+                if data["cmd"] == "NEW":
+                    unique_id = self.nccl.unique_id_from_bytes(
+                        bytes(data["unique_id"]))
+                    with torch.cuda.device(self.device):
+                        rank = 1
+                        with set_du_swift_context(self.nccl_num_channels):
+                            comm: ncclComm_t = self.nccl.ncclCommInitRank(
+                                2, unique_id, rank)
+                        self.comms[remote_address.decode()] = (comm, rank)
+                        logger.info(
+                            "🤝ncclCommInitRank Success, %s👈%s, MyRank:%s",
+                            self.zmq_address, remote_address.decode(), rank)
+                elif data["cmd"] == "PUT":
+                    tensor_id = data["tensor_id"]
+                    if "tensor_split_num" in data:
+                        self.tensor_split_num = data["tensor_split_num"]
+                    try:
+                        with torch.cuda.stream(self.recv_stream):
+                            tensor = torch.empty(data["shape"],
+                                                 dtype=getattr(
+                                                     torch, data["dtype"]),
+                                                 device=self.device)
+                        self.router_socket.send_multipart(
+                            [remote_address, b"0"])
+                        comm, rank = self.comms[remote_address.decode()]
+                        self._recv(comm, tensor, rank ^ 1, self.recv_stream)
+                        tensor_size = tensor.element_size() * tensor.numel()
+                        if (self.buffer_size + tensor_size
+                                > self.buffer_size_threshold):
+                            # Store Tensor in memory pool
+                            addr = self.pool.store_tensor(tensor)
+                            tensor = (addr, tensor.dtype, tensor.shape)
+                        else:
+                            self.buffer_size += tensor_size
+
+                    except torch.cuda.OutOfMemoryError:
+                        self.router_socket.send_multipart(
+                            [remote_address, b"1"])
+                        tensor = None
+                        logger.warning(
+                            "🔴[PUT]Recv Tensor, Out Of Memory, %s👈%s, "
+                            "data:%s", self.zmq_address,
+                            remote_address.decode(), data)
+
+                    with self.recv_store_cv:
+                        self.recv_store[tensor_id] = tensor
+                        self._have_received_tensor_id(tensor_id)
+                        self.recv_store_cv.notify()
+                elif data["cmd"] == "PUT_NEW":
+                    tensor_id = data["tensor_id"]
+                    if "tensor_split_num" in data:
+                        self.tensor_split_num = data["tensor_split_num"]
+                    try:
+                        with torch.cuda.stream(self.recv_stream):
+                            tensor = torch.empty(data["shape"],
+                                                 dtype=getattr(
+                                                     torch, data["dtype"]),
+                                                 device=self.device)
+                        self.router_socket.send_multipart(
+                            [remote_address, b"0"])
+                        # comm, rank = self.comms[remote_address.decode()]
+                        # self._recv(comm, tensor, rank ^ 1, self.recv_stream)
+                        comm, rank = self.comms[data["pd_pair_id"]]
+                        self._recv(comm, tensor, int(data["comm_rank"]), self.recv_stream)
+                        tensor_size = tensor.element_size() * tensor.numel()
+                        if (self.buffer_size + tensor_size
+                                > self.buffer_size_threshold):
+                            # Store Tensor in memory pool
+                            addr = self.pool.store_tensor(tensor)
+                            tensor = (addr, tensor.dtype, tensor.shape)
+                        else:
+                            self.buffer_size += tensor_size
+
+                    except torch.cuda.OutOfMemoryError:
+                        self.router_socket.send_multipart(
+                            [remote_address, b"1"])
+                        tensor = None
+                        logger.warning(
+                            "🔴[PUT]Recv Tensor, Out Of Memory, %s👈%s, "
+                            "data:%s", self.zmq_address,
+                            remote_address.decode(), data)
+                    with self.recv_store_cv:
+                        self.recv_store[tensor_id] = tensor
+                        self._have_received_tensor_id(tensor_id)
+                        self.recv_store_cv.notify()
+                elif data["cmd"] == "comm_init":
+                    unique_id = self.nccl.unique_id_from_bytes(
+                        bytes(data["unique_id"]))
+                    with torch.cuda.device(self.device):
+                        rank = int(data["rank"])
+                        world_size = int(data["world_size"])
+                        with set_du_swift_context(self.nccl_num_channels):
+                            comm: ncclComm_t = self.nccl.ncclCommInitRank(
+                                    world_size, unique_id, rank)
+                        self.comms[data["pd_pair_id"]] = (comm, rank)
+                        logger.info(
+                            "🤝ncclCommInitRank Success, %s👈%s, MyRank:%s",
+                            self.zmq_address, data["pd_pair_id"], rank)
+                elif data["cmd"] == "GET":
+                    tensor_id = data["tensor_id"]
+                    with self.send_store_cv:
+                        tensor = self.send_store.pop(tensor_id, None)
+                        if tensor is not None:
+                            data = {
+                                "ret": 0,
+                                "shape": tensor.shape,
+                                "dtype":
+                                str(tensor.dtype).replace("torch.", "")
+                            }
+                            # LRU
+                            self.send_store[tensor_id] = tensor
+                            self._have_sent_tensor_id(tensor_id)
+                        else:
+                            data = {"ret": 1}
+
+                    self.router_socket.send_multipart(
+                        [remote_address, msgpack.dumps(data)])
+
+                    if data["ret"] == 0:
+                        comm, rank = self.comms[remote_address.decode()]
+                        self._send(comm, tensor.to(self.device), rank ^ 1,
+                                   self.send_stream)
+                else:
+                    logger.warning(
+                        "🚧Unexpected, Received message from %s, data:%s",
+                        remote_address, data)
+
+    def _have_sent_tensor_id(self, tensor_id: str):
+        request_id = tensor_id.split('#')[0]
+        if request_id not in self.send_request_id_to_tensor_ids:
+            self.send_request_id_to_tensor_ids[request_id] = set()
+        self.send_request_id_to_tensor_ids[request_id].add(tensor_id)
+
+    def _have_received_tensor_id(self, tensor_id: str):
+        request_id = tensor_id.split('#')[0]
+        if request_id not in self.recv_request_id_to_tensor_ids:
+            self.recv_request_id_to_tensor_ids[request_id] = set()
+        self.recv_request_id_to_tensor_ids[request_id].add(tensor_id)
+
+    def _send_async(self):
+        while True:
+            with self.send_queue_cv:
+                while not self.send_queue:
+                    self.send_queue_cv.wait()
+                if envs.VLLM_ENABLE_TBO or envs.VLLM_P2P_ASYNC:
+                    tensor_id, remote_address, kv_layer, slot_mapping, tbo_evt = self.send_queue.popleft()
+                else:
+                    tensor_id, remote_address, tensor = self.send_queue.popleft()
+                if not self.send_queue:
+                    self.send_queue_cv.notify()
+            if (envs.VLLM_ENABLE_TBO or envs.VLLM_P2P_ASYNC) and tbo_evt is not None:
+                self.send_stream.wait_event(tbo_evt)
+                self._send_kv_p2p_sync(tensor_id, kv_layer, slot_mapping, remote_address)
+            else:
+                if self.multiple_machines:
+                    self._send_sync(tensor_id, tensor, remote_address)
+                else:
+                    # logger.info(f"""=============xiabo tensor_id:{tensor_id} remote_address:{remote_address}""")
+                    self._send_sync_new(tensor_id, tensor, remote_address)
+
+    def wait_for_sent(self):
+        if self.send_type == "PUT_ASYNC":
+            start_time = time.time()
+            with self.send_queue_cv:
+                while self.send_queue:
+                    self.send_queue_cv.wait()
+            duration = time.time() - start_time
+            logger.debug(
+                "🚧[PUT_ASYNC]It took %.3fms to wait for the send_queue"
+                " to be empty, rank:%d", duration * 1000, self.rank)
+
+    def _send_kv_p2p_sync(self, tensor_id: str, kv_layer: torch.Tensor,
+                           slot_mapping: torch.Tensor, remote_address: str) -> bool:
+        if remote_address not in self.socks:
+            self._create_connect(remote_address)
+        sock = self.socks[remote_address]
+        comm, rank = self.comms[remote_address]
+
+        is_mla = (kv_layer.ndim == 3)
+        hidden_dim = kv_layer.shape[-1]
+        
+        if self.p2p_async_buf is None:
+            if is_mla:
+                self.p2p_async_buf = torch.empty((self.p2p_async_kv_tokens, hidden_dim),
+                               dtype=kv_layer.dtype, device=kv_layer.device)
+            else:
+                self.p2p_async_buf = torch.empty((2, self.p2p_async_kv_tokens, hidden_dim),
+                               dtype=kv_layer.dtype, device=kv_layer.device)
+
+        pack_num = (slot_mapping.shape[0] - 1) // self.p2p_async_kv_tokens + 1
+        self.tensor_split_num = pack_num
+
+        with torch.cuda.stream(self.send_stream):
+            for pack_idx in range(pack_num):
+                start = pack_idx * self.p2p_async_kv_tokens
+                end = min((pack_idx + 1) * self.p2p_async_kv_tokens, slot_mapping.shape[0])
+                sub_index = slot_mapping[start:end]
+
+                if is_mla:
+                    num_pages, page_size = kv_layer.shape[0], kv_layer.shape[1]
+                    data = kv_layer.reshape(num_pages * page_size, -1)
+                    torch.index_select(data, dim=0, index=sub_index, out=self.p2p_async_buf[:end-start])
+                    tx_shape = (end - start, hidden_dim)
+                else:
+                    num_pages, page_size = kv_layer.shape[1], kv_layer.shape[2]
+                    data = kv_layer.reshape(2, num_pages * page_size, -1)
+                    torch.index_select(data, dim=1, index=sub_index, out=self.p2p_async_buf[:, :end-start])
+                    tx_shape = (2, end - start, hidden_dim)
+
+                if is_mla:
+                    send_tensor = self.p2p_async_buf[:end-start]
+                else:
+                    send_tensor = self.p2p_async_buf[:, :end-start]
+                    
+                header = {
+                    "cmd": "PUT",
+                    "tensor_id": tensor_id + "#" + str(pack_idx), # 拼 pack_idx
+                    "pack_idx": pack_idx,
+                    "tensor_split_num": pack_num,
+                    "shape": tx_shape,
+                    "dtype": str(kv_layer.dtype).replace("torch.", "")
+                }
+                sock.send(msgpack.dumps(header))
+                response = sock.recv()
+                if response != b"0":
+                    logger.error(
+                        "🔴Send Tensor Failed | %s 👉 %s | Rank:%s | shape:%s | size:%.4f GB | response:%s",
+                        self.zmq_address, remote_address, rank,
+                        tuple(send_tensor.shape), send_tensor.element_size() * send_tensor.numel() / 1024**3, 
+                        response.decode()
+                    )
+                    return False
+                
+                self._send(comm, send_tensor, rank ^ 1, self.send_stream)
+
+        if self.send_type == "PUT_ASYNC":
+            self._have_sent_tensor_id(tensor_id)
+
+        return True
+
+    def _send_sync_new(
+        self,
+        tensor_id: str,
+        tensor: torch.Tensor,
+        remote_address: typing.Optional[RemoteAddr] = None,
+    ) -> bool:
+        if remote_address is None:
+            return False
+        
+        if remote_address.zmq_address not in self.socks:
+            # logger.info(f"""=============xiabo remote_address.zmq_address:{remote_address.zmq_address}""")
+            self._create_connect_new(remote_address.zmq_address)
+
+        sock = self.socks[remote_address.zmq_address]
+        comm, rank = self.comms[remote_address.pd_pair_id]
+        data = {
+            "cmd": "PUT_NEW",
+            "tensor_id": tensor_id,
+            "shape": tensor.shape,
+            "dtype": str(tensor.dtype).replace("torch.", ""),
+            "pd_pair_id": remote_address.pd_pair_id,
+            "comm_rank": rank
+        }
+        logger.info(f"""_send_sync_new:{data}""")
+        sock.send(msgpack.dumps(data))
+
+        response = sock.recv()
+        if response != b"0":
+            logger.error(
+                "🔴Send Tensor, Peer Out Of Memory/Threshold, %s 👉 %s, "
+                "MyRank:%s, data:%s, tensor:%s, size:%fGB, response:%s",
+                self.zmq_address, remote_address.zmq_address, rank, data, tensor.shape,
+                tensor.element_size() * tensor.numel() / 1024**3,
+                response.decode())
+            return False
+
+        self._send(comm, tensor.to(self.device), remote_address.comm_rank, self.send_stream)
+        if self.send_type == "PUT_ASYNC":
+            self._have_sent_tensor_id(tensor_id)
+
+        return True
+
+    def _send_sync(
+        self,
+        tensor_id: str,
+        tensor: torch.Tensor,
+        remote_address: typing.Optional[str] = None,
+    ) -> bool:
+        if remote_address is None:
+            return False
+        if remote_address not in self.socks:
+            self._create_connect(remote_address)
+
+        sock = self.socks[remote_address]
+        comm, rank = self.comms[remote_address]
+        data = {
+            "cmd": "PUT",
+            "tensor_id": tensor_id,
+            "shape": tensor.shape,
+            "dtype": str(tensor.dtype).replace("torch.", "")
+        }
+        sock.send(msgpack.dumps(data))
+
+        response = sock.recv()
+        if response != b"0":
+            logger.error(
+                "🔴Send Tensor, Peer Out Of Memory/Threshold, %s 👉 %s, "
+                "MyRank:%s, data:%s, tensor:%s, size:%fGB, response:%s",
+                self.zmq_address, remote_address, rank, data, tensor.shape,
+                tensor.element_size() * tensor.numel() / 1024**3,
+                response.decode())
+            return False
+
+        self._send(comm, tensor.to(self.device), rank ^ 1, self.send_stream)
+
+        if self.send_type == "PUT_ASYNC":
+            self._have_sent_tensor_id(tensor_id)
+
+        return True
+
+    def get_finished(
+        self, finished_req_ids: set[str], forward_context: "ForwardContext"
+    ) -> tuple[Optional[set[str]], Optional[set[str]]]:
+        """
+        Notifies worker-side connector ids of requests that have
+        finished generating tokens.
+
+        Returns:
+            ids of requests that have finished asynchronous transfer,
+            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).
+        """
+
+        # Clear the buffer upon request completion.
+        for request_id in finished_req_ids:
+            for layer_name in forward_context.no_compile_layers:
+                tensor_id = request_id + "#" + layer_name
+                if tensor_id in self.recv_store:
+                    with self.recv_store_cv:
+                        tensor = self.recv_store.pop(tensor_id, None)
+                        self.send_request_id_to_tensor_ids.pop(
+                            request_id, None)
+                        self.recv_request_id_to_tensor_ids.pop(
+                            request_id, None)
+                    addr = 0
+                    if isinstance(tensor, tuple):
+                        addr, _, _ = tensor
+                        self.pool.free(addr)
+
+        # TODO:Retrieve requests that have already sent the KV cache.
+        finished_sending: set[str] = set()
+
+        # TODO:Retrieve requests that have already received the KV cache.
+        finished_recving: set[str] = set()
+
+        return finished_sending or None, finished_recving or None
+
+    def _ping(self):
+        sock = self.context.socket(zmq.DEALER)
+        sock.setsockopt_string(zmq.IDENTITY, self.zmq_address)
+        logger.debug("ping start, zmq_address:%s", self.zmq_address)
+        sock.connect(f"tcp://{self.proxy_address}")
+        data = {
+            "type": "P" if self.config.is_kv_producer else "D",
+            "http_address": self.http_address,
+            "zmq_address": self.zmq_address
+        }
+        while True:
+            sock.send(msgpack.dumps(data))
+            time.sleep(3)
+
+    def _ping_new(self):
+        sock = self.context.socket(zmq.DEALER)
+        sock.setsockopt_string(zmq.IDENTITY, self.zmq_address)
+        logger.debug("ping start, zmq_address:%s", self.zmq_address)
+        sock.connect(f"tcp://{self.proxy_address}")
+
+        if self.rank == 0:
+            data = {
+                "type": "P_init" if self.config.is_kv_producer else "D_init",
+                "http_address": self.http_address,
+                "zmq_address": self.zmq_address,
+                "dp_size" : self.dp_size,
+                "pp_size" : self.pp_size,
+                "tp_size" : self.tp_size
+            }
+            # logger.info(f"""_ping data:{data}""")
+            sock.send(msgpack.dumps(data))
+        data = {
+            "type": "P" if self.config.is_kv_producer else "D",
+            "http_address": self.http_address,
+            "dp_rank" : self.dp_rank,
+            "pp_rank" : self.pp_rank,
+            "tp_rank" : self.tp_rank,
+            "zmq_address": self.zmq_address
+        }
+        # while True:
+        # logger.info(f"""_ping data:{data}""")
+        sock.send(msgpack.dumps(data))
+            # time.sleep(3)
+
+    def _send(self, comm, tensor: torch.Tensor, dst: int, stream=None):
+        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()
+
+        with torch.cuda.stream(stream):
+            self.nccl.ncclSend(buffer_type(tensor.data_ptr()), tensor.numel(),
+                               ncclDataTypeEnum.from_torch(tensor.dtype), dst,
+                               comm, cudaStream_t(stream.cuda_stream))
+        stream.synchronize()
+
+    def _recv(self, comm, tensor: torch.Tensor, src: int, stream=None):
+        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()
+
+        with torch.cuda.stream(stream):
+            self.nccl.ncclRecv(buffer_type(tensor.data_ptr()), tensor.numel(),
+                               ncclDataTypeEnum.from_torch(tensor.dtype), src,
+                               comm, cudaStream_t(stream.cuda_stream))
+        stream.synchronize()
+
+    def close(self) -> None:
+        self._listener_thread.join()
+        if self.send_type == "PUT_ASYNC":
+            self._send_thread.join()
+        if self._ping_thread is not None:
+            self._ping_thread.join()
+
+    def get_pp_indices_d(self, num_hidden_layers: int, pp_rank: int,
+                   pp_size: int) -> tuple[int, int]:
+        partition_list_str = envs.VLLM_PP_LAYER_PARTITION_D
+        if partition_list_str is not None:
+            try:
+                partitions = [
+                    int(layer) for layer in partition_list_str.split(",")
+                ]
+            except ValueError as err:
+                raise ValueError("Invalid partition string: {}".format(
+                    partition_list_str)) from err
+            if len(partitions) != pp_size:
+                raise ValueError(f"{len(partitions)=} does not match {pp_size=}.")
+            if sum(partitions) != num_hidden_layers:
+                raise ValueError(
+                    f"{sum(partitions)=} does not match {num_hidden_layers=}.")
+        else:
+            layers_per_partition = num_hidden_layers // pp_size
+            partitions = [layers_per_partition for _ in range(pp_size)]
+
+            if remaining_layers := num_hidden_layers % pp_size:
+                for i in range(2, remaining_layers + 2):
+                    partitions[-i] += 1
+                logger.info(
+                    "Hidden layers were unevenly partitioned: [%s]. "
+                    "This can be manually overridden using the "
+                    "VLLM_PP_LAYER_PARTITION_D environment variable",
+                    ",".join(str(p) for p in partitions))
+
+        start_layer = sum(partitions[:pp_rank])
+        end_layer = start_layer + partitions[pp_rank]
+
+        return (start_layer, end_layer)
+
+    def compute_remote_pp_rank(self, layer_name: str) -> int:
+        current_layer_idx = extract_layer_index(layer_name)
+        for d_pp_rank in range(self.remote_pp_size):
+            start, end = self.get_pp_indices_d(self.total_num_hidden_layers, d_pp_rank, self.remote_pp_size)
+            # logger.info(f"""compute_remote_pp_rank : current_layer_idx:{current_layer_idx} start:{start} end:{end}""")
+            if (current_layer_idx == self.total_num_hidden_layers):
+                return self.remote_pp_size - 1
+            if start <= current_layer_idx < end:
+                return d_pp_rank
+        return -1
+
+
+    @staticmethod
+    def get_tensor_id(request_id: str, layer_name: str) -> str:
+        return request_id + "#" + layer_name
+
+    @staticmethod
+    def parse_request_id(request_id: str, is_prefill=True) -> tuple[str, int]:
+        # Regular expression to match the string hostname and integer port
+        if is_prefill:
+            pattern = r"___decode_addr_(.*):(\d+)"
+        else:
+            pattern = r"___prefill_addr_(.*):(\d+)___"
+
+        # Use re.search to find the pattern in the request_id
+        match = regex.search(pattern, request_id)
+        if match:
+            # Extract the ranks
+            ip = match.group(1)
+            port = int(match.group(2))
+
+            return ip, port
+        raise ValueError(
+            f"Request id {request_id} does not contain hostname and port")

vllm/distributed/kv_transfer/kv_connector/v1/du/tensor_memory_pool.py

+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import atexit
+import ctypes
+import math
+from dataclasses import dataclass
+
+import torch
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class MemoryBlock:
+    size: int
+    addr: int
+
+
+"""A memory pool for managing pinned host memory allocations for tensors.
+
+This class implements a buddy allocation system to efficiently manage pinned
+host memory for tensor storage. It supports allocation, deallocation, and
+tensor storage/retrieval operations.
+
+Key Features:
+- Uses power-of-two block sizes for efficient buddy allocation
+- Supports splitting and merging of memory blocks
+- Provides methods to store CUDA tensors in pinned host memory
+- Allows loading tensors from pinned memory back to device
+- Automatically cleans up memory on destruction
+
+Attributes:
+    max_block_size (int): Maximum block size (rounded to nearest power of two)
+    min_block_size (int): Minimum block size (rounded to nearest power of two)
+    free_lists (dict): Dictionary of free memory blocks by size
+    allocated_blocks (dict): Dictionary of currently allocated blocks
+    base_tensor (torch.Tensor): Base pinned memory tensor
+    base_address (int): Base memory address of the pinned memory region
+
+Example:
+    >>> pool = TensorMemoryPool(max_block_size=1024*1024)
+    >>> tensor = torch.randn(100, device='cuda')
+    >>> addr = pool.store_tensor(tensor)
+    >>> loaded_tensor = pool.load_tensor(addr, tensor.dtype,
+    ...                                  tensor.shape, 'cuda')
+    >>> pool.free(addr)
+"""
+
+
+class TensorMemoryPool:
+    """Initializes the memory pool with given size constraints.
+
+    Args:
+        max_block_size (int): Maximum size of memory blocks to manage
+        min_block_size (int, optional): Minimum size of memory blocks
+            to manage. Defaults to 512.
+
+    Raises:
+        ValueError: If block sizes are invalid or max_block_size is less
+            than min_block_size
+    """
+
+    def __init__(self, max_block_size: int, min_block_size: int = 128):
+        if max_block_size <= 0 or min_block_size <= 0:
+            raise ValueError("Block sizes must be positive")
+        if max_block_size < min_block_size:
+            raise ValueError(
+                "Max block size must be greater than min block size")
+
+        self.max_block_size = self._round_to_power_of_two(max_block_size)
+        self.min_block_size = self._round_to_power_of_two(min_block_size)
+
+        self.free_lists: dict[int, dict[int, MemoryBlock]] = {}
+        self.allocated_blocks: dict[int, MemoryBlock] = {}
+
+        self._initialize_free_lists()
+        self._allocate_pinned_memory()
+
+        atexit.register(self.cleanup)
+
+    def _round_to_power_of_two(self, size: int) -> int:
+        return 1 << (size - 1).bit_length()
+
+    def _initialize_free_lists(self):
+        size = self.max_block_size
+        while size >= self.min_block_size:
+            self.free_lists[size] = {}
+            size //= 2
+
+    def _allocate_pinned_memory(self):
+        self.base_tensor = torch.empty(self.max_block_size // 4,
+                                       dtype=torch.float32,
+                                       pin_memory=True)
+        self.base_address = self.base_tensor.data_ptr()
+        initial_block = MemoryBlock(size=self.max_block_size,
+                                    addr=self.base_address)
+        self.free_lists[self.max_block_size][
+            initial_block.addr] = initial_block
+        logger.debug("TensorMemoryPool, base_address:", self.base_address,
+                     self.base_address % self.max_block_size)
+
+    def allocate(self, size: int) -> int:
+        """Allocates a memory block of at least the requested size.
+
+        Args:
+            size (int): Minimum size of memory to allocate
+
+        Returns:
+            int: Address of the allocated memory block
+
+        Raises:
+            ValueError: If size is invalid or insufficient memory is available
+        """
+        if size <= 0:
+            raise ValueError("Allocation size must be positive")
+
+        required_size = self._round_to_power_of_two(
+            max(size, self.min_block_size))
+        if required_size > self.max_block_size:
+            raise ValueError("Requested size exceeds maximum block size")
+
+        current_size = required_size
+        while current_size <= self.max_block_size:
+            if self.free_lists[current_size]:
+                _, block = self.free_lists[current_size].popitem()
+                self._split_block(block, required_size)
+                self.allocated_blocks[block.addr] = block
+                return block.addr
+            current_size *= 2
+
+        raise ValueError("Insufficient memory")
+
+    def _split_block(self, block: MemoryBlock, required_size: int):
+        while (block.size > required_size
+               and block.size // 2 >= self.min_block_size):
+            buddy_size = block.size // 2
+            buddy_addr = block.addr + buddy_size
+
+            buddy = MemoryBlock(size=buddy_size, addr=buddy_addr)
+            block.size = buddy_size
+
+            self.free_lists[buddy_size][buddy.addr] = buddy
+
+    def free(self, addr: int):
+        """Frees an allocated memory block.
+
+        Args:
+            addr (int): Address of the block to free
+
+        Raises:
+            ValueError: If address is invalid or not allocated
+        """
+        if addr not in self.allocated_blocks:
+            raise ValueError("Invalid address to free")
+
+        block = self.allocated_blocks.pop(addr)
+        self._merge_buddies(block)
+
+    def _merge_buddies(self, block: MemoryBlock):
+        MAX_MERGE_DEPTH = 30
+        depth = 0
+
+        while depth < MAX_MERGE_DEPTH:
+            buddy_offset = block.size if (block.addr - self.base_address) % (
+                2 * block.size) == 0 else -block.size
+            buddy_addr = block.addr + buddy_offset
+            buddy = self.free_lists[block.size].get(buddy_addr)
+            if buddy:
+                del self.free_lists[buddy.size][buddy.addr]
+                merged_addr = min(block.addr, buddy.addr)
+                merged_size = block.size * 2
+                block = MemoryBlock(size=merged_size, addr=merged_addr)
+                depth += 1
+            else:
+                break
+        self.free_lists[block.size][block.addr] = block
+
+    def store_tensor(self, tensor: torch.Tensor) -> int:
+        """Stores a CUDA tensor in pinned host memory.
+
+        Args:
+            tensor (torch.Tensor): CUDA tensor to store
+
+        Returns:
+            int: Address where the tensor is stored
+
+        Raises:
+            ValueError: If tensor is not on CUDA or allocation fails
+        """
+        if not tensor.is_cuda:
+            raise ValueError("Only CUDA tensors can be stored")
+
+        size = tensor.element_size() * tensor.numel()
+        addr = self.allocate(size)
+        block = self.allocated_blocks[addr]
+
+        if block.size < size:
+            self.free(addr)
+            raise ValueError(
+                f"Allocated block size {block.size} is smaller than "
+                f"required size {size}")
+
+        try:
+            buffer = (ctypes.c_byte * block.size).from_address(block.addr)
+            cpu_tensor = torch.frombuffer(buffer,
+                                          dtype=tensor.dtype,
+                                          count=tensor.numel()).reshape(
+                                              tensor.shape)
+        except ValueError as err:
+            self.free(addr)
+            raise ValueError(f"Failed to create tensor view: {err}") from err
+
+        cpu_tensor.copy_(tensor)
+
+        return addr
+
+    def load_tensor(self, addr: int, dtype: torch.dtype,
+                    shape: tuple[int, ...], device) -> torch.Tensor:
+        """Loads a tensor from pinned host memory to the specified device.
+
+        Args:
+            addr (int): Address where tensor is stored
+            dtype (torch.dtype): Data type of the tensor
+            shape (tuple[int, ...]): Shape of the tensor
+            device: Target device for the loaded tensor
+
+        Returns:
+            torch.Tensor: The loaded tensor on the specified device
+
+        Raises:
+            ValueError: If address is invalid or sizes don't match
+        """
+        if addr not in self.allocated_blocks:
+            raise ValueError("Invalid address to load")
+
+        block = self.allocated_blocks[addr]
+        num_elements = math.prod(shape)
+        dtype_size = torch.tensor([], dtype=dtype).element_size()
+        required_size = num_elements * dtype_size
+
+        if required_size > block.size:
+            raise ValueError("Requested tensor size exceeds block size")
+
+        buffer = (ctypes.c_byte * block.size).from_address(block.addr)
+        cpu_tensor = torch.frombuffer(buffer, dtype=dtype,
+                                      count=num_elements).reshape(shape)
+
+        cuda_tensor = torch.empty(shape, dtype=dtype, device=device)
+
+        cuda_tensor.copy_(cpu_tensor)
+
+        return cuda_tensor
+
+    def cleanup(self):
+        """Cleans up all memory resources and resets the pool state."""
+        self.free_lists.clear()
+        self.allocated_blocks.clear()
+        if hasattr(self, 'base_tensor'):
+            del self.base_tensor
+
+    def __del__(self):
+        self.cleanup()

vllm/envs.py

@@ -281,6 +281,8 @@ if TYPE_CHECKING:
     VLLM_USE_LIGHTOP_MOE_ALIGN: bool = False
     VLLM_USE_MERGE_ATTN_STATES_OPT: bool = False
     USE_FUSED_RMS_QUANT: bool = False
+    VLLM_P2P_ASYNC: bool = False
+    VLLM_P2P_BUF_TOKENS: int = 30000
     USE_FUSED_SILU_MUL_QUANT: bool = False
     VLLM_USE_PD_SPLIT: bool = False
     VLLM_USE_PP_SYNC: bool = False
@@ -1805,7 +1807,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # vllm will use rmsquant fused op
     "USE_FUSED_RMS_QUANT":
         lambda: bool(int(os.getenv("USE_FUSED_RMS_QUANT", "0"))),
-        
+    # vllm pd separation will be used async
+    "VLLM_P2P_ASYNC":
+    lambda: bool(int(os.getenv("VLLM_P2P_ASYNC", "0"))),
+    # pd separation p2p async buf tokens
+    "VLLM_P2P_BUF_TOKENS":
+    lambda: int(os.getenv("VLLM_P2P_BUF_TOKENS", "30000")),
     # vllm will use silu_mul_quant fused op
     "USE_FUSED_SILU_MUL_QUANT":
         lambda: (os.getenv("USE_FUSED_SILU_MUL_QUANT", "False").lower() in