Merge branch 'v0.9.2-dev-du_connector' into 'v0.9.2-dev'

V0.9.2 dev du connector

See merge request dcutoolkit/deeplearing/vllm!409

vllm/distributed/kv_transfer/kv_connector/factory.py

@@ -117,6 +117,11 @@ KVConnectorFactory.register_connector(
     "vllm.distributed.kv_transfer.kv_connector.v1.p2p.p2p_nccl_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/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/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_engine.py

@@ -13,7 +13,6 @@ 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 (
@@ -21,13 +20,6 @@ from vllm.distributed.device_communicators.pynccl_wrapper import (
 from vllm.distributed.kv_transfer.kv_connector.v1.p2p.tensor_memory_pool import (  # noqa: E501
     TensorMemoryPool)
 from vllm.utils import current_stream, 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
@@ -36,11 +28,6 @@ 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_p2p_nccl_context(num_channels: str):
@@ -72,37 +59,17 @@ def set_p2p_nccl_context(num_channels: str):
                 os.environ.pop(var, None)
 
 
-@dataclass
-class RemoteAddr:
-    pd_pair_id: str = ""
-    zmq_address: str = ""
-    comm_rank: int = 0
-
-
 class P2pNcclEngine:
 
     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,
+                 hostname: str = "",
+                 port_offset: 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)
 
@@ -128,7 +95,7 @@ class P2pNcclEngine:
         port = int(self.config.kv_port) + port_offset
         if port == 0:
             raise ValueError("Port cannot be 0")
-        self._hostname = get_ip()
+        self._hostname = hostname
         self._port = port
 
         # Each card corresponds to a ZMQ address.
@@ -161,10 +128,6 @@ class P2pNcclEngine:
 
         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)
@@ -204,16 +167,11 @@ class P2pNcclEngine:
         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()
+        if port_offset == 0 and self.proxy_address != "":
+            self._ping_thread = threading.Thread(target=self._ping,
+                                                 daemon=True)
+            self._ping_thread.start()
+
         logger.info(
             "💯P2pNcclEngine init, rank:%d, local_rank:%d, http_address:%s, "
             "zmq_address:%s, proxy_address:%s, send_type:%s, buffer_size_"
@@ -221,21 +179,6 @@ class P2pNcclEngine:
             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:
@@ -263,73 +206,11 @@ class P2pNcclEngine:
 
         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(" P2PNCCL 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(" P2PNCCL 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:
@@ -369,53 +250,6 @@ class P2pNcclEngine:
                         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,
@@ -493,8 +327,6 @@ class P2pNcclEngine:
                             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"],
@@ -511,6 +343,10 @@ class P2pNcclEngine:
                             # Store Tensor in memory pool
                             addr = self.pool.store_tensor(tensor)
                             tensor = (addr, tensor.dtype, tensor.shape)
+                            logger.warning(
+                                "🔴[PUT]Recv Tensor, Out Of Threshold, "
+                                "%s👈%s, data:%s, addr:%d", self.zmq_address,
+                                remote_address.decode(), data, addr)
                         else:
                             self.buffer_size += tensor_size
 
@@ -527,56 +363,7 @@ class P2pNcclEngine:
                         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_p2p_nccl_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:
@@ -623,21 +410,10 @@ class P2pNcclEngine:
             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()
+                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)
+            self._send_sync(tensor_id, tensor, remote_address)
 
     def wait_for_sent(self):
         if self.send_type == "PUT_ASYNC":
@@ -851,36 +627,6 @@ class P2pNcclEngine:
             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}, "

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

@@ -63,7 +63,7 @@ class TensorMemoryPool:
             than min_block_size
     """
 
-    def __init__(self, max_block_size: int, min_block_size: int = 128):
+    def __init__(self, max_block_size: int, min_block_size: int = 512):
         if max_block_size <= 0 or min_block_size <= 0:
             raise ValueError("Block sizes must be positive")
         if max_block_size < min_block_size: