pd mode use p2p triton kernel to manage kv trans && refactor deepseekv2 code (#691)

Co-authored-by: wangzaijun <[email protected]>

Author: hiworldwzj

lightllm/common/deepseek2_mem_manager.py

@@ -1,10 +1,11 @@
 import torch
 import os
-
+import torch.distributed as dist
 from lightllm.server.pd_io_struct import KVMoveTask
 from .mem_manager import MemoryManager
 from typing import List
 from lightllm.utils.log_utils import init_logger
+from lightllm.common.kv_trans_kernel.kv_trans import kv_trans
 
 logger = init_logger(__name__)
 
@@ -33,6 +34,7 @@ def alloc_kv_move_buffer(self, max_req_total_len):
         self.kv_move_buffer = torch.empty(
             (1, max_req_total_len + 8, self.head_num, self.head_dim), dtype=self.dtype, device="cuda"
         )
+        self.kv_move_buf_indexes = torch.arange(0, max_req_total_len + 8, dtype=torch.int64, device="cuda")
         return
 
     def send_to_decode_node(
@@ -41,8 +43,6 @@ def send_to_decode_node(
         assert dp_size == 1
 
         # 先将数据发送到指定的一张卡上的buffer，再发送。
-        import torch.distributed as dist
-
         move_token_indexes = []
         for task in move_tasks:
             if task.move_kv_len != 0:
@@ -69,8 +69,6 @@ def receive_from_prefill_node(
         assert dp_size == 1
 
         # 先将数据接受到指定的一张卡上的buffer，再复制到其他的卡上。
-        import torch.distributed as dist
-
         move_token_indexes = []
         for task in move_tasks:
             if task.move_kv_len != 0:
@@ -97,6 +95,58 @@ def _write_kv_move_data(self, token_indexes: torch.Tensor, buffer_tensor: torch.
         self.kv_buffer[layer_index : layer_index + 1, token_indexes, :, :] = buffer_tensor
         return
 
+    def send_to_decode_node_p2p(self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size: int):
+        """
+        使用 p2p triton kernel 进行数据复制和传输的实现方式。
+        """
+        assert dp_size == 1
+
+        move_token_indexes = []
+        for task in move_tasks:
+            if task.move_kv_len != 0:
+                move_token_indexes.extend(task.prefill_token_indexes[-task.move_kv_len :])
+
+        move_token_indexes = torch.tensor(move_token_indexes, dtype=torch.int64, device="cuda")
+        for layer_index in range(self.layer_num):
+            move_buffer = self._get_kv_move_data_p2p(move_token_indexes, layer_index, self.kv_move_buffer)
+            dist.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):
+        move_token_num = len(token_indexes)
+        move_size = self.kv_buffer.numel() // self.layer_num // self.size * move_token_num
+        move_buffer = kv_move_buffer.view(-1)[0:move_size].view(move_token_num, self.head_num, self.head_dim)
+        kv_trans(
+            self.kv_buffer[layer_index, :, :, :], token_indexes, move_buffer, self.kv_move_buf_indexes[0:move_token_num]
+        )
+        return move_buffer
+
+    def receive_from_prefill_node_p2p(
+        self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size: int
+    ):
+        assert dp_size == 1
+
+        move_token_indexes = []
+        for task in move_tasks:
+            if task.move_kv_len != 0:
+                move_token_indexes.extend(task.decode_token_indexes[-task.move_kv_len :])
+
+        move_token_indexes = torch.tensor(move_token_indexes, dtype=torch.int64, device="cuda")
+
+        token_num = len(move_token_indexes)
+        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)
+            for i, mem in enumerate(mem_managers):
+                mem._write_kv_move_data_p2p(move_token_indexes, recive_buffer, layer_index)
+        return
+
+    def _write_kv_move_data_p2p(self, token_indexes: torch.Tensor, buffer_tensor: torch.Tensor, layer_index):
+        move_token_num = len(token_indexes)
+        kv_trans(buffer_tensor, self.kv_move_buf_indexes[0:move_token_num], self.kv_buffer[layer_index], token_indexes)
+        return
+
     @torch.no_grad()
     def free_all(self):
         self.can_use_mem_size = len(self.mem_state) - self.holding_size

lightllm/common/kv_trans_kernel/__init__.py

@@ -0,0 +1,78 @@
+import torch
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _kv_trans_kernel(
+    input_ptr,
+    input_stride_0,
+    input_stride_1,
+    input_stride_2,
+    input_token_idx_ptr,
+    output_ptr,
+    output_stride_0,
+    output_stride_1,
+    output_stride_2,
+    output_token_idx_ptr,
+    token_num: int,
+    head_num: int,
+    head_dim: int,
+    grid_count: int,
+    BLOCK_SIZE: tl.constexpr,
+    NUM_STAGES: tl.constexpr,
+):
+    input_stride_0 = tl.cast(input_stride_0, dtype=tl.int64)
+    input_stride_1 = tl.cast(input_stride_1, dtype=tl.int64)
+    output_stride_0 = tl.cast(output_stride_0, dtype=tl.int64)
+    output_stride_1 = tl.cast(output_stride_1, dtype=tl.int64)
+
+    head_num_dim = head_num * head_dim
+    tid = tl.program_id(0)
+
+    offs = tl.arange(0, BLOCK_SIZE)
+    while tid < token_num:
+        input_token_idx = tl.load(input_token_idx_ptr + tid)
+        output_token_idx = tl.load(output_token_idx_ptr + tid)
+        for block_idx in tl.range(0, tl.cdiv(head_num_dim, BLOCK_SIZE), 1, num_stages=NUM_STAGES):
+            cur_offs = block_idx * BLOCK_SIZE + offs
+            in_datas = tl.load(input_ptr + input_stride_0 * input_token_idx + cur_offs, mask=cur_offs < head_num_dim)
+            tl.store(output_ptr + output_stride_0 * output_token_idx + cur_offs, in_datas, mask=cur_offs < head_num_dim)
+
+        tid += grid_count
+
+    return
+
+
+def kv_trans(input: torch.Tensor, input_idx: torch.Tensor, output: torch.Tensor, output_idx: torch.Tensor):
+    assert input.is_contiguous()
+    assert output.is_contiguous()
+    assert len(input.shape) == 3
+    assert len(output.shape) == 3
+    assert len(input_idx) == len(output_idx)
+
+    _, head_num, head_dim = input.shape
+    token_num = len(input_idx)
+    # 用较少的资源来做数据传输，防止占用过多的 sm 计算单元
+    grid_count = 20
+    BLOCK_SIZE = 256
+    NUM_STAGES = 3
+    grid = (grid_count,)
+
+    _kv_trans_kernel[grid](
+        input,
+        *input.stride(),
+        input_idx,
+        output,
+        *output.stride(),
+        output_idx,
+        token_num=token_num,
+        head_num=head_num,
+        head_dim=head_dim,
+        grid_count=grid_count,
+        BLOCK_SIZE=BLOCK_SIZE,
+        NUM_STAGES=NUM_STAGES,
+        num_warps=1,
+    )
+    return

lightllm/common/kv_trans_kernel/kv_trans.py

@@ -7,6 +7,7 @@
 from lightllm.utils.log_utils import init_logger
 from lightllm.server.router.dynamic_prompt.shared_arr import SharedInt
 from lightllm.utils.profile_max_tokens import get_available_gpu_memory, get_total_gpu_memory
+from lightllm.common.kv_trans_kernel.kv_trans import kv_trans
 
 logger = init_logger(__name__)
 
@@ -78,6 +79,7 @@ def alloc_kv_move_buffer(self, max_req_total_len):
         self.kv_move_buffer = torch.empty(
             (1, max_req_total_len + 8, 2 * self.head_num, self.head_dim), dtype=self.dtype, device="cuda"
         )
+        self.kv_move_buf_indexes = torch.arange(0, max_req_total_len + 8, dtype=torch.int64, device="cuda")
         return
 
     def send_to_decode_node(self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size: int):
@@ -159,6 +161,65 @@ def _write_kv_move_data(self, token_indexes: torch.Tensor, buffer_tensor: torch.
         self.kv_buffer[layer_index : layer_index + 1, token_indexes, :, :] = buffer_tensor
         return
 
+    def send_to_decode_node_p2p(self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size: int):
+        """
+        使用 p2p triton kernel 进行数据复制和传输的实现方式。
+        """
+        assert dp_size == 1
+
+        # 先将数据发送到指定的一张卡上的buffer，再发送。
+        import torch.distributed as dist
+
+        move_token_indexes = []
+        for task in move_tasks:
+            if task.move_kv_len != 0:
+                move_token_indexes.extend(task.prefill_token_indexes[-task.move_kv_len :])
+
+        move_token_indexes = torch.tensor(move_token_indexes, dtype=torch.int64, device="cuda")
+        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)
+        return
+
+    def _get_kv_move_data_p2p(self, token_indexes: torch.Tensor, layer_index: int, kv_move_buffer: torch.Tensor):
+        move_token_num = len(token_indexes)
+        move_size = self.kv_buffer.numel() // self.layer_num // self.size * move_token_num
+        move_buffer = kv_move_buffer.view(-1)[0:move_size].view(move_token_num, 2 * self.head_num, self.head_dim)
+        kv_trans(
+            self.kv_buffer[layer_index, :, :, :], token_indexes, move_buffer, self.kv_move_buf_indexes[0:move_token_num]
+        )
+        return move_buffer
+
+    def receive_from_prefill_node_p2p(
+        self, move_tasks: List[KVMoveTask], mem_managers: List["MemoryManager"], dp_size: int
+    ):
+        assert dp_size == 1
+
+        # 先将数据接受到指定的一张卡上的buffer，再复制到其他的卡上。
+        import torch.distributed as dist
+
+        move_token_indexes = []
+        for task in move_tasks:
+            if task.move_kv_len != 0:
+                move_token_indexes.extend(task.decode_token_indexes[-task.move_kv_len :])
+
+        move_token_indexes = torch.tensor(move_token_indexes, dtype=torch.int64, device="cuda")
+
+        token_num = len(move_token_indexes)
+        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, 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)
+                mem._write_kv_move_data_p2p(move_token_indexes, recive_buffer, layer_index)
+        return
+
+    def _write_kv_move_data_p2p(self, token_indexes: torch.Tensor, buffer_tensor: torch.Tensor, layer_index):
+        move_token_num = len(token_indexes)
+        kv_trans(buffer_tensor, self.kv_move_buf_indexes[0:move_token_num], self.kv_buffer[layer_index], token_indexes)
+        return
+
     def _free_buffers(self):
         self.kv_buffer = None