single kv transfer process for pd.

Author: Weichao Luo

lightllm/common/deepseek2_mem_manager.py

@@ -7,6 +7,7 @@
 from lightllm.utils.log_utils import init_logger
 from lightllm.common.kv_trans_kernel.kv_trans import kv_trans
 from lightllm.common.kv_trans_kernel.kv_trans_v2 import kv_trans_v2_for_d_node, kv_trans_v2_for_p_node
+from lightllm.distributed.pynccl import PyNcclCommunicator
 
 logger = init_logger(__name__)
 
@@ -41,7 +42,8 @@ def alloc_kv_move_buffer(self, max_req_total_len):
         return
 
     def send_to_decode_node(
-        self, move_tasks: List[KVMoveTask], mem_managers: List["Deepseek2MemoryManager"], dp_size_in_node: int
+        self, move_tasks: List[KVMoveTask], mem_managers: List["Deepseek2MemoryManager"], dp_size_in_node: int,
+        nccl_comm: PyNcclCommunicator
     ):
         assert dp_size_in_node == 1
 
@@ -55,7 +57,7 @@ def send_to_decode_node(
         cur_mem = mem_managers[cur_device_index]
         for layer_index in range(cur_mem.layer_num):
             move_buffer = cur_mem._get_kv_move_data(move_token_indexes, layer_index)
-            dist.send(move_buffer, dst=1)
+            nccl_comm.send(move_buffer, dst=1)
         return
 
     def _get_kv_move_data(self, token_indexes: List[int], layer_index: int):
@@ -67,7 +69,8 @@ def _get_kv_move_data(self, token_indexes: List[int], layer_index: int):
         return move_buffer
 
     def receive_from_prefill_node(
-        self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size_in_node: int
+        self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size_in_node: int,
+        nccl_comm: PyNcclCommunicator
     ):
         assert dp_size_in_node == 1
 
@@ -82,7 +85,7 @@ def receive_from_prefill_node(
         move_size = self.kv_buffer.numel() // self.layer_num // self.size * token_num
         recive_buffer = self.kv_move_buffer.view(-1)[0:move_size].view(1, token_num, self.head_num, self.head_dim)
         for layer_index in range(self.layer_num):
-            dist.recv(recive_buffer, src=0)
+            nccl_comm.recv(recive_buffer, src=0)
             for i, mem in enumerate(mem_managers):
                 if i == cur_device_index:
                     mem._write_kv_move_data(move_token_indexes, recive_buffer, layer_index)
@@ -99,7 +102,8 @@ def _write_kv_move_data(self, token_indexes: torch.Tensor, buffer_tensor: torch.
         return
 
     def send_to_decode_node_p2p(
-        self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size_in_node: int
+        self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size_in_node: int,
+        nccl_comm: PyNcclCommunicator
     ):
         """
         使用 p2p triton kernel 进行数据复制和传输的实现方式。
@@ -126,7 +130,7 @@ def send_to_decode_node_p2p(
             move_buffer = self._get_kv_move_data_p2p(
                 move_token_indexes, token_dp_indexes, layer_index, self.kv_move_buffer, dp_size_in_node
             )
-            dist.send(move_buffer, dst=1)
+            nccl_comm.send(move_buffer, dst=1)
         return
 
     def _get_kv_move_data_p2p(
@@ -151,7 +155,8 @@ def _get_kv_move_data_p2p(
         return move_buffer
 
     def receive_from_prefill_node_p2p(
-        self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size_in_node: int
+        self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size_in_node: int,
+        nccl_comm: PyNcclCommunicator
     ):
         if not hasattr(self, "mem_ptrs_dict"):
             self.mem_ptrs_dict = {}
@@ -176,7 +181,7 @@ def receive_from_prefill_node_p2p(
         move_size = self.kv_buffer.numel() // self.layer_num // self.size * token_num
         recive_buffer = self.kv_move_buffer.view(-1)[0:move_size].view(token_num, self.head_num, self.head_dim)
         for layer_index in range(self.layer_num):
-            dist.recv(recive_buffer, src=0)
+            nccl_comm.recv(recive_buffer, src=0)
             self._write_kv_move_data_p2p(
                 move_token_indexes, token_dp_indexes, recive_buffer, layer_index, dp_size_in_node
             )

lightllm/common/mem_manager.py

@@ -10,6 +10,7 @@
 from lightllm.common.kv_trans_kernel.kv_trans import kv_trans
 from lightllm.utils.dist_utils import get_current_rank_in_node
 from lightllm.utils.envs_utils import get_unique_server_name, get_env_start_args
+from lightllm.distributed.pynccl import PyNcclCommunicator
 
 
 logger = init_logger(__name__)
@@ -86,7 +87,8 @@ def alloc_kv_move_buffer(self, max_req_total_len):
         return
 
     def send_to_decode_node(
-        self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size_in_node: int
+        self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size_in_node: int,
+        nccl_comm: PyNcclCommunicator
     ):
         assert dp_size_in_node == 1
 
@@ -103,14 +105,14 @@ def send_to_decode_node(
             for layer_index in range(mem.layer_num):
                 move_buffer = mem._get_kv_move_data(move_token_indexes, layer_index)
                 if i == cur_device_index:
-                    dist.send(move_buffer, dst=1)
+                    nccl_comm.send(move_buffer, dst=1)
                 else:
                     move_size = move_buffer.numel()
                     new_move_buffer = cur_mem.kv_move_buffer.view(-1)[0:move_size].view(move_buffer.shape)
                     from torch.cuda import comm
 
                     comm.broadcast(move_buffer, out=[new_move_buffer])
-                    dist.send(new_move_buffer, dst=1)
+                    nccl_comm.send(new_move_buffer, dst=1)
         return
 
     def _get_kv_move_data(self, token_indexes: List[int], layer_index: int):
@@ -122,7 +124,8 @@ def _get_kv_move_data(self, token_indexes: List[int], layer_index: int):
         return move_buffer
 
     def receive_from_prefill_node(
-        self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size_in_node: int
+        self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size_in_node: int,
+        nccl_comm: PyNcclCommunicator,
     ):
         assert dp_size_in_node == 1
 
@@ -139,7 +142,7 @@ def receive_from_prefill_node(
         recive_buffer = self.kv_move_buffer.view(-1)[0:move_size].view(1, token_num, 2 * self.head_num, self.head_dim)
         for i, mem in enumerate(mem_managers):
             for layer_index in range(mem.layer_num):
-                dist.recv(recive_buffer, src=0)
+                nccl_comm.recv(recive_buffer, src=0)
                 if i == cur_device_index:
                     mem._write_kv_move_data(move_token_indexes, recive_buffer, layer_index)
                 else:
@@ -155,7 +158,8 @@ def _write_kv_move_data(self, token_indexes: torch.Tensor, buffer_tensor: torch.
         return
 
     def send_to_decode_node_p2p(
-        self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size_in_node: int
+        self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size_in_node: int,
+        nccl_comm: PyNcclCommunicator
     ):
         """
         使用 p2p triton kernel 进行数据复制和传输的实现方式。
@@ -173,7 +177,7 @@ def send_to_decode_node_p2p(
         for i, mem in enumerate(mem_managers):
             for layer_index in range(mem.layer_num):
                 move_buffer = mem._get_kv_move_data_p2p(move_token_indexes, layer_index, self.kv_move_buffer)
-                dist.send(move_buffer, dst=1)
+                nccl_comm.send(move_buffer, dst=1)
         return
 
     def _get_kv_move_data_p2p(self, token_indexes: torch.Tensor, layer_index: int, kv_move_buffer: torch.Tensor):
@@ -186,7 +190,8 @@ def _get_kv_move_data_p2p(self, token_indexes: torch.Tensor, layer_index: int, k
         return move_buffer
 
     def receive_from_prefill_node_p2p(
-        self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size_in_node: int
+        self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size_in_node: int,
+        nccl_comm: PyNcclCommunicator,
     ):
         assert dp_size_in_node == 1
 
@@ -204,7 +209,7 @@ def receive_from_prefill_node_p2p(
         recive_buffer = self.kv_move_buffer.view(-1)[0:move_size].view(token_num, 2 * self.head_num, self.head_dim)
         for i, mem in enumerate(mem_managers):
             for layer_index in range(mem.layer_num):
-                dist.recv(recive_buffer, src=0)
+                nccl_comm.recv(recive_buffer, src=0)
                 mem._write_kv_move_data_p2p(move_token_indexes, recive_buffer, layer_index)
         return