layer into triton op

Author: WANDY666

lightllm/common/kv_trans_kernel/kv_trans_v2.py

@@ -199,13 +199,16 @@ def _kv_trans_for_dp_kernel(
     input_stride_0,
     input_stride_1,
     input_stride_2,
+    input_stride_3,
     input_token_idx_ptr,
     input_token_dp_index_ptr,
     output_ptr,
     output_stride_0,
     output_stride_1,
     output_stride_2,
+    output_stride_3,
     output_token_idx_ptr,
+    layer_num: tl.constexpr,
     token_num: int,
     head_num: int,
     head_dim: int,
@@ -229,11 +232,20 @@ def _kv_trans_for_dp_kernel(
         mem_index = RANK_IN_DP + dp_index * CARD_NUM_PER_D
         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
-            input_ptr = tl.load(input_mems_ptr + mem_index).to(tl.pointer_type(output_ptr.dtype.element_ty))
-            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)
+
+        input_ptr = tl.load(input_mems_ptr + mem_index).to(tl.pointer_type(output_ptr.dtype.element_ty))
+        for layer_idx in tl.range(0, layer_num, 1):
+            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 * layer_idx + input_stride_1 * input_token_idx + cur_offs,
+                    mask=cur_offs < head_num_dim,
+                )
+                tl.store(
+                    output_ptr + output_stride_0 * layer_idx + output_stride_1 * output_token_idx + cur_offs,
+                    in_datas,
+                    mask=cur_offs < head_num_dim,
+                )
 
         tid += grid_count
 
@@ -250,19 +262,19 @@ def kv_trans_for_dp(
     rank_in_dp: int,
 ):
     """
-    input_mems 是一个 torch.uint64 的tensor, 其内部存储了当前使用的对应的mem_manager对象中kv cache的首指针。
+    input_mems 是一个 torch.uint64 的tensor, shape为(layer_num, mem_num)，其内部存储了当前使用的对应的mem_manager对象中kv cache的首指针。
     """
     assert input_mems.is_contiguous()
     assert output.is_contiguous()
     assert len(input_mems.shape) == 1
-    assert len(output.shape) == 3
+    assert len(output.shape) == 4
     assert len(input_idx) == len(output_idx)
     assert len(output_idx) == len(input_dp_idx)
     assert len(input_mems) % dp_size_in_node == 0
 
     card_num_per_d = len(input_mems) // dp_size_in_node
 
-    _, head_num, head_dim = output.shape
+    layer_num, _, head_num, head_dim = output.shape
     token_num = len(output_idx)
     # 用较少的资源来做数据传输，防止占用过多的 sm 计算单元
     grid_count = 20
@@ -278,6 +290,7 @@ def kv_trans_for_dp(
         output,
         *output.stride(),
         output_idx,
+        layer_num=layer_num,
         token_num=token_num,
         head_num=head_num,
         head_dim=head_dim,

lightllm/common/mem_manager.py

@@ -414,25 +414,23 @@ def copy_kv_from_other_dp_ranks(
         dp_size_in_node: int,
         rank_in_dp: int,
     ):
-        if not hasattr(self, "mem_ptrs_dict"):
-            self.mem_ptrs_dict = {}
-            for layer_index in range(self.layer_num):
-                mems_ptr = []
-                for i in range(0, len(mem_managers)):
-                    mems_ptr.append(mem_managers[i].kv_buffer[layer_index, :, :, :].data_ptr())
-                mems_ptr = torch.tensor(mems_ptr, dtype=torch.uint64, device="cuda")
-                self.mem_ptrs_dict[layer_index] = mems_ptr
-
-        for layer_index in range(self.layer_num):
-            kv_trans_for_dp(
-                input_mems=self.mem_ptrs_dict[layer_index],
-                input_idx=move_token_indexes,
-                input_dp_idx=token_dp_indexes,
-                output=self.kv_buffer[layer_index],
-                output_idx=mem_indexes,
-                dp_size_in_node=dp_size_in_node,
-                rank_in_dp=rank_in_dp,
-            )
+        if not hasattr(self, "mem_ptrs_tensor"):
+            # 构建一个2D tensor，shape为(layer_num, mem_num)
+            mems_ptr_list = []
+            for i in range(0, len(mem_managers)):
+                mems_ptr_list.append(mem_managers[i].kv_buffer.data_ptr())
+            self.mem_ptrs_tensor = torch.tensor(mems_ptr_list, dtype=torch.uint64, device="cuda")
+
+        # 一次性传输所有层
+        kv_trans_for_dp(
+            input_mems=self.mem_ptrs_tensor,
+            input_idx=move_token_indexes,
+            input_dp_idx=token_dp_indexes,
+            output=self.kv_buffer,
+            output_idx=mem_indexes,
+            dp_size_in_node=dp_size_in_node,
+            rank_in_dp=rank_in_dp,
+        )
 
     def create_shm(self):
         obj_bytes = ForkingPickler.dumps(self)
@@ -449,7 +447,9 @@ def from_shm(rank_in_node):
             f"{get_unique_server_name()}_mem_manager_{rank_in_node}", LIGHTLLM_MEM_MANAGER_SHM_SIZE
         )
         bytes_len = int.from_bytes(shm.buf[0:4], "little")
-        return ForkingPickler.loads(shm.buf[4 : 4 + bytes_len])
+        obj_bytes = shm.buf[4 : 4 + bytes_len].tobytes()
+        shm.close()
+        return ForkingPickler.loads(obj_bytes)
 
 
 class ReadOnlyStaticsMemoryManager:

lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py

@@ -29,8 +29,8 @@
 from lightllm.common.mem_manager import MemoryManager
 import torch.multiprocessing as mp
 
-min_trans_token_num = os.getenv("LIGHTLLM_MIN_TRANS_TOKEN_NUM", 512)
-dp_kv_transfer_req_num = os.getenv("LIGHTLLM_DP_KV_TRANSFER_REQ_NUM", 16)
+min_trans_token_num = int(os.getenv("LIGHTLLM_MIN_TRANS_TOKEN_NUM", "512"))
+dp_kv_transfer_req_num = int(os.getenv("LIGHTLLM_DP_KV_TRANSFER_REQ_NUM", "16"))
 
 
 class DPChunkedPrefillBackend(ModeBackend):
@@ -167,7 +167,6 @@ def _fetch_dp_prompt_cache(
                 if sampling_param.disable_prompt_cache:
                     continue
                 shm_req.link_prompt_ids_shm_array()
-                shm_req.link_kv_indexes_shm_array()
 
                 kv_len, value_tensor = self._match_radix_cache(shm_req)
                 with g_infer_context.shm_req_manager.get_req_lock_by_index(shm_req.index_in_shm_mem):
@@ -210,7 +209,7 @@ def _fetch_dp_prompt_cache(
     def _transfer_dp_kv_cache(self, my_match: List[Tuple], other_match: List[Tuple]):
         other_shm_reqs = []
         for match, index in other_match:
-            shm_req, kv_len, value_tensor = match
+            shm_req, kv_len, value_tensor, _ = match
             trans_len = kv_len - shm_req.dp_origin_kv_len
             if shm_req.dp_max_kv_rank == self.dp_rank_in_node:
                 self.shared_kv_indexes.arr[index, 0:trans_len] = value_tensor[shm_req.dp_origin_kv_len : kv_len]
@@ -227,7 +226,7 @@ def _transfer_dp_kv_cache(self, my_match: List[Tuple], other_match: List[Tuple])
         trans_info = []
         alloc_size = 0
         for match, index in my_match:
-            shm_req, kv_len, value_tensor = match
+            shm_req, kv_len, value_tensor, _ = match
             trans_len = shm_req.dp_max_kv_len - kv_len
             if trans_len > 0 and shm_req.dp_max_kv_rank != self.dp_rank_in_node:
                 move_token_indexes.extend(self.shared_kv_indexes.arr[index, 0:trans_len])

lightllm/utils/log_utils.py

@@ -7,7 +7,7 @@
 import time
 from typing import Optional
 
-_FORMAT = "%(levelname)s %(asctime)s [%(filename)s:%(lineno)d] %(message)s"
+_FORMAT = "%(levelname)s %(asctime)s,%(msecs)03d [%(filename)s:%(lineno)d] %(message)s"
 _DATE_FORMAT = "%m-%d %H:%M:%S"
 
 _LOG_LEVEL = os.environ.get("LIGHTLLM_LOG_LEVEL", "debug")

unit_tests/common/kv_trans_kernel/test_kv_trans_v2.py

@@ -1,7 +1,7 @@
 import pytest
 import torch
 import random
-from lightllm.common.kv_trans_kernel.kv_trans_v2 import kv_trans_v2_for_p_node, kv_trans_v2_for_d_node
+from lightllm.common.kv_trans_kernel.kv_trans_v2 import kv_trans_v2_for_p_node, kv_trans_v2_for_d_node, kv_trans_for_dp
 
 
 @pytest.mark.parametrize(
@@ -73,5 +73,51 @@ def test_kv_trans_v2_for_d_node(token_num):
     return
 
 
+@pytest.mark.parametrize(
+    "token_num",
+    [token_num for token_num in range(5, 10)],
+)
+def test_kv_trans_for_dp(token_num):
+    card_num = 8
+    dp_size_in_node = 4
+    layer_num = 3
+    head_num = 2
+    head_dim = 512
+    kv_buffer_token_num = 512
+    rank_in_dp = 1
+
+    card_num_per_d = card_num // dp_size_in_node
+
+    # 创建多层的 mem，每个 mem 包含所有层的数据
+    mems = []
+    for _ in range(card_num):
+        mems.append(
+            torch.randn((layer_num, kv_buffer_token_num, head_num, head_dim), dtype=torch.float16, device="cuda")
+        )
+
+    input_mems = torch.tensor([e.data_ptr() for e in mems], dtype=torch.uint64, device="cuda")
+    input_idx = [random.randint(0, kv_buffer_token_num - 1) for _ in range(token_num)]
+    input_idx = torch.tensor(input_idx, dtype=torch.int32, device="cuda")
+    input_dp_idx = [random.randint(0, dp_size_in_node - 1) for _ in range(token_num)]
+    input_dp_idx = torch.tensor(input_dp_idx, dtype=torch.int32, device="cuda")
+
+    true_output = torch.zeros((layer_num, kv_buffer_token_num, head_num, head_dim), dtype=torch.float16, device="cuda")
+    test_output = torch.zeros((layer_num, kv_buffer_token_num, head_num, head_dim), dtype=torch.float16, device="cuda")
+    output_idx = torch.arange(0, token_num, 1, dtype=torch.int32, device="cuda")
+
+    kv_trans_for_dp(input_mems, input_idx, input_dp_idx, test_output, output_idx, dp_size_in_node, rank_in_dp)
+
+    # 验证结果
+    for dest_token_index, src_token_index, dp_index in zip(
+        list(range(token_num)), input_idx.cpu().numpy(), input_dp_idx.cpu().numpy()
+    ):
+        mem_index = rank_in_dp + dp_index * card_num_per_d
+        # 所有 layer 都从同一个 mem 的对应层读取
+        true_output[:, dest_token_index, :, :] = mems[mem_index][:, src_token_index, :, :]
+
+    assert torch.equal(true_output, test_output), "kv_trans_for_dp output mismatch"
+    return
+
+
 if __name__ == "__main__":
     pytest.main()