fix

Author: wangzaijun

lightllm/common/deepseek2_mem_manager.py

@@ -42,13 +42,15 @@ def alloc_paged_kv_move_buffer(self, page_num, page_size) -> torch.Tensor:
             (page_num, page_size, self.layer_num, self.head_num, self.head_dim), dtype=self.dtype, device="cuda"
         )
         return self.kv_move_buffer
-    
-    def write_mem_to_page_kv_move_buffer(self,
-                                        mem_indexes: List[int], 
-                                        page_index: int,
-                                        dp_index: int,
-                                        mem_managers: List["MemoryManager"],
-                                        dp_world_size:int):
+
+    def write_mem_to_page_kv_move_buffer(
+        self,
+        mem_indexes: List[int],
+        page_index: int,
+        dp_index: int,
+        mem_managers: List["MemoryManager"],
+        dp_world_size: int,
+    ):
         cur_page = self.kv_move_buffer[page_index]
         dp_mems = mem_managers[(dp_index * dp_world_size) : ((dp_index + 1) * dp_world_size)]
         mla_page_io(
@@ -58,13 +60,15 @@ def write_mem_to_page_kv_move_buffer(self,
             mode="write",
         )
         return
-    
-    def read_page_kv_move_buffer_to_mem(self,
-                                        mem_indexes: List[int], 
-                                        page_index: int,
-                                        dp_index: int,
-                                        mem_managers: List["MemoryManager"],
-                                        dp_world_size:int):
+
+    def read_page_kv_move_buffer_to_mem(
+        self,
+        mem_indexes: List[int],
+        page_index: int,
+        dp_index: int,
+        mem_managers: List["MemoryManager"],
+        dp_world_size: int,
+    ):
         cur_page = self.kv_move_buffer[page_index]
         dp_mems = mem_managers[(dp_index * dp_world_size) : ((dp_index + 1) * dp_world_size)]
         mem_indexes = torch.tensor(mem_indexes, dtype=torch.int64, device="cuda")
@@ -76,7 +80,6 @@ def read_page_kv_move_buffer_to_mem(self,
                 mode="read",
             )
 
-
     def send_to_decode_node(
         self,
         move_tasks: List[KVMoveTask],

lightllm/common/kv_trans_kernel/nixl_kv_trans.py

@@ -6,6 +6,7 @@
 
 logger = init_logger(__name__)
 
+
 @triton.jit
 def _page_io(
     mem_index_ptr,
@@ -43,7 +44,7 @@ def _page_io(
     v_stride_layer_num = tl.cast(v_stride_layer_num, dtype=tl.int64)
     k_stride_size = tl.cast(k_stride_size, dtype=tl.int64)
     v_stride_size = tl.cast(v_stride_size, dtype=tl.int64)
-    
+
     tid = tl.program_id(0)
     kv_head_id = tl.program_id(1)
     page_head_id = page_head_start + kv_head_id
@@ -57,18 +58,86 @@ def _page_io(
 
     for layer_index in tl.range(layer_num, num_stages=3):
         if IS_WRITE:
-            k_tensor = tl.load(k_ptr + layer_index * k_stride_layer_num + mem_index * k_stride_size + kv_head_id * k_stride_head + off_dim * k_stride_dim, mask=mask)
-            v_tensor = tl.load(v_ptr + layer_index * v_stride_layer_num + mem_index * v_stride_size + kv_head_id * v_stride_head + off_dim * v_stride_dim, mask=mask)
-            tl.store(k_page_ptr + tid * k_page_stride_size + layer_index * k_page_stride_layer_num + page_head_id * k_page_stride_head + off_dim * k_page_stride_dim, k_tensor, mask=mask)
-            tl.store(v_page_ptr + tid * v_page_stride_size + layer_index * v_page_stride_layer_num + page_head_id * v_page_stride_head + off_dim * v_page_stride_dim, v_tensor, mask=mask)
+            k_tensor = tl.load(
+                k_ptr
+                + layer_index * k_stride_layer_num
+                + mem_index * k_stride_size
+                + kv_head_id * k_stride_head
+                + off_dim * k_stride_dim,
+                mask=mask,
+            )
+            v_tensor = tl.load(
+                v_ptr
+                + layer_index * v_stride_layer_num
+                + mem_index * v_stride_size
+                + kv_head_id * v_stride_head
+                + off_dim * v_stride_dim,
+                mask=mask,
+            )
+            tl.store(
+                k_page_ptr
+                + tid * k_page_stride_size
+                + layer_index * k_page_stride_layer_num
+                + page_head_id * k_page_stride_head
+                + off_dim * k_page_stride_dim,
+                k_tensor,
+                mask=mask,
+            )
+            tl.store(
+                v_page_ptr
+                + tid * v_page_stride_size
+                + layer_index * v_page_stride_layer_num
+                + page_head_id * v_page_stride_head
+                + off_dim * v_page_stride_dim,
+                v_tensor,
+                mask=mask,
+            )
         else:
-            k_page_tensor = tl.load(k_page_ptr + tid * k_page_stride_size + layer_index * k_page_stride_layer_num + page_head_id * k_page_stride_head + off_dim * k_page_stride_dim, mask=mask)
-            v_page_tensor = tl.load(v_page_ptr + tid * v_page_stride_size + layer_index * v_page_stride_layer_num + page_head_id * v_page_stride_head + off_dim * v_page_stride_dim, mask=mask)
-            tl.store(k_ptr + layer_index * k_stride_layer_num + mem_index * k_stride_size + kv_head_id * k_stride_head + off_dim * k_stride_dim, k_page_tensor, mask=mask)
-            tl.store(v_ptr + layer_index * v_stride_layer_num + mem_index * v_stride_size + kv_head_id * v_stride_head + off_dim * v_stride_dim, v_page_tensor, mask=mask)
+            k_page_tensor = tl.load(
+                k_page_ptr
+                + tid * k_page_stride_size
+                + layer_index * k_page_stride_layer_num
+                + page_head_id * k_page_stride_head
+                + off_dim * k_page_stride_dim,
+                mask=mask,
+            )
+            v_page_tensor = tl.load(
+                v_page_ptr
+                + tid * v_page_stride_size
+                + layer_index * v_page_stride_layer_num
+                + page_head_id * v_page_stride_head
+                + off_dim * v_page_stride_dim,
+                mask=mask,
+            )
+            tl.store(
+                k_ptr
+                + layer_index * k_stride_layer_num
+                + mem_index * k_stride_size
+                + kv_head_id * k_stride_head
+                + off_dim * k_stride_dim,
+                k_page_tensor,
+                mask=mask,
+            )
+            tl.store(
+                v_ptr
+                + layer_index * v_stride_layer_num
+                + mem_index * v_stride_size
+                + kv_head_id * v_stride_head
+                + off_dim * v_stride_dim,
+                v_page_tensor,
+                mask=mask,
+            )
     return
 
-def page_io(mem_indexes:torch.Tensor, page_tensor: torch.Tensor, kv_buffer: torch.Tensor, tp_index:int, tp_world_size:int, mode:str):
+
+def page_io(
+    mem_indexes: torch.Tensor,
+    page_tensor: torch.Tensor,
+    kv_buffer: torch.Tensor,
+    tp_index: int,
+    tp_world_size: int,
+    mode: str,
+):
     assert mode in ["read", "write"]
     assert mem_indexes.is_contiguous()
     assert page_tensor.is_contiguous()
@@ -86,9 +155,10 @@ def page_io(mem_indexes:torch.Tensor, page_tensor: torch.Tensor, kv_buffer: torc
     v_page_tensor = page_tensor[:, :, -page_v_head_num:, :]
 
     k_head_num, v_head_num = kv_head_num // 2, kv_head_num // 2
+    assert k_head_num == v_head_num
     k_buffer = kv_buffer[:, :, 0:k_head_num, :]
     v_buffer = kv_buffer[:, :, k_head_num:, :]
-    
+
     tp_index = tp_index // repeat_count
     tp_world_size = tp_world_size // repeat_count
 
@@ -127,14 +197,13 @@ def page_io(mem_indexes:torch.Tensor, page_tensor: torch.Tensor, kv_buffer: torc
         layer_num=layer_num,
         head_dim=head_dim,
         HEAD_DIM_BLOCK=triton.next_power_of_2(head_dim),
-        IS_WRITE=mode=="write",
+        IS_WRITE=mode == "write",
         NEED_MASK=triton.next_power_of_2(head_dim) != head_dim,
         num_warps=1,
     )
     return
 
 
-
 @triton.jit
 def _mla_page_io(
     mem_index_ptr,
@@ -157,7 +226,7 @@ def _mla_page_io(
     page_stride_size = tl.cast(page_stride_size, dtype=tl.int64)
     kv_stride_layer_num = tl.cast(kv_stride_layer_num, dtype=tl.int64)
     kv_stride_size = tl.cast(kv_stride_size, dtype=tl.int64)
-    
+
     tid = tl.program_id(0)
 
     mem_index = tl.load(mem_index_ptr + tid)
@@ -169,14 +238,45 @@ def _mla_page_io(
 
     for layer_index in tl.range(layer_num, num_stages=3):
         if IS_WRITE:
-            kv_tensor = tl.load(kv_ptr + layer_index * kv_stride_layer_num + mem_index * kv_stride_size + 0 * kv_stride_head + off_dim * kv_stride_dim, mask=mask)
-            tl.store(page_ptr + tid * page_stride_size + layer_index * page_stride_layer_num + 0 * page_stride_head + off_dim * page_stride_dim, kv_tensor, mask=mask)
+            kv_tensor = tl.load(
+                kv_ptr
+                + layer_index * kv_stride_layer_num
+                + mem_index * kv_stride_size
+                + 0 * kv_stride_head
+                + off_dim * kv_stride_dim,
+                mask=mask,
+            )
+            tl.store(
+                page_ptr
+                + tid * page_stride_size
+                + layer_index * page_stride_layer_num
+                + 0 * page_stride_head
+                + off_dim * page_stride_dim,
+                kv_tensor,
+                mask=mask,
+            )
         else:
-            page_tensor = tl.load(page_ptr + tid * page_stride_size + layer_index * page_stride_layer_num + 0 * page_stride_head + off_dim * page_stride_dim, mask=mask)
-            tl.store(kv_ptr + layer_index * kv_stride_layer_num + mem_index * kv_stride_size + 0 * kv_stride_head + off_dim * kv_stride_dim, page_tensor, mask=mask)
+            page_tensor = tl.load(
+                page_ptr
+                + tid * page_stride_size
+                + layer_index * page_stride_layer_num
+                + 0 * page_stride_head
+                + off_dim * page_stride_dim,
+                mask=mask,
+            )
+            tl.store(
+                kv_ptr
+                + layer_index * kv_stride_layer_num
+                + mem_index * kv_stride_size
+                + 0 * kv_stride_head
+                + off_dim * kv_stride_dim,
+                page_tensor,
+                mask=mask,
+            )
     return
 
-def mla_page_io(mem_indexes:torch.Tensor, page_tensor: torch.Tensor, kv_buffer: torch.Tensor, mode:str):
+
+def mla_page_io(mem_indexes: torch.Tensor, page_tensor: torch.Tensor, kv_buffer: torch.Tensor, mode: str):
     assert mode in ["read", "write"]
     assert mem_indexes.is_contiguous()
     assert page_tensor.is_contiguous()
@@ -189,7 +289,6 @@ def mla_page_io(mem_indexes:torch.Tensor, page_tensor: torch.Tensor, kv_buffer:
     assert page_head_dim == head_dim
     assert page_head_num == kv_head_num == 1
 
-
     token_num = len(mem_indexes)
     grid = (token_num,)
 
@@ -208,7 +307,7 @@ def mla_page_io(mem_indexes:torch.Tensor, page_tensor: torch.Tensor, kv_buffer:
         layer_num=layer_num,
         head_dim=head_dim,
         HEAD_DIM_BLOCK=triton.next_power_of_2(head_dim),
-        IS_WRITE=mode=="write",
+        IS_WRITE=mode == "write",
         NEED_MASK=triton.next_power_of_2(head_dim) != head_dim,
         num_warps=1,
     )

lightllm/common/mem_manager.py

@@ -105,55 +105,62 @@ def alloc_kv_move_buffer(self, max_req_total_len):
     def alloc_paged_kv_move_buffer(self, page_num, page_size) -> torch.Tensor:
         if isinstance(self, MemoryManager) and type(self) != MemoryManager:
             raise NotImplementedError("subclass need reimpl this method")
-        
+
         num_kv_head = get_num_key_value_heads(get_env_start_args().model_dir)
         self.kv_move_buffer = torch.empty(
             (page_num, page_size, self.layer_num, 2 * num_kv_head, self.head_dim), dtype=self.dtype, device="cuda"
         )
         return self.kv_move_buffer
-    
-    def write_mem_to_page_kv_move_buffer(self,
-                                        mem_indexes: List[int], 
-                                        page_index: int,
-                                        dp_index: int,
-                                        mem_managers: List["MemoryManager"],
-                                        dp_world_size:int):
+
+    def write_mem_to_page_kv_move_buffer(
+        self,
+        mem_indexes: List[int],
+        page_index: int,
+        dp_index: int,
+        mem_managers: List["MemoryManager"],
+        dp_world_size: int,
+    ):
         cur_page = self.kv_move_buffer[page_index]
         repeat_count = dp_world_size * self.kv_buffer.shape[2] // self.kv_move_buffer.shape[3]
         dp_mems = mem_managers[(dp_index * dp_world_size) : ((dp_index + 1) * dp_world_size)]
         for tp_index in range(dp_world_size):
             if tp_index % repeat_count == 0:
-                page_io(torch.tensor(mem_indexes, dtype=torch.int64, device="cuda"),
-                              page_tensor=cur_page,
-                              kv_buffer=dp_mems[tp_index].kv_buffer,
-                              tp_index=tp_index,
-                              tp_world_size=dp_world_size,
-                              mode="write")
+                page_io(
+                    torch.tensor(mem_indexes, dtype=torch.int64, device="cuda"),
+                    page_tensor=cur_page,
+                    kv_buffer=dp_mems[tp_index].kv_buffer,
+                    tp_index=tp_index,
+                    tp_world_size=dp_world_size,
+                    mode="write",
+                )
         # keep for debug
         # logger.info(f"src token tensor {self.kv_buffer[:, mem_indexes[0], 0, 0]}")
         # logger.info(f"src page token tensor {cur_page[0, :, 0, 0]}")
         return
-    
-    def read_page_kv_move_buffer_to_mem(self,
-                                        mem_indexes: List[int], 
-                                        page_index: int,
-                                        dp_index: int,
-                                        mem_managers: List["MemoryManager"],
-                                        dp_world_size:int):
+
+    def read_page_kv_move_buffer_to_mem(
+        self,
+        mem_indexes: List[int],
+        page_index: int,
+        dp_index: int,
+        mem_managers: List["MemoryManager"],
+        dp_world_size: int,
+    ):
         cur_page = self.kv_move_buffer[page_index]
         dp_mems = mem_managers[(dp_index * dp_world_size) : ((dp_index + 1) * dp_world_size)]
         for tp_index in range(dp_world_size):
-            page_io(torch.tensor(mem_indexes, dtype=torch.int64, device="cuda"),
-                            page_tensor=cur_page,
-                            kv_buffer=dp_mems[tp_index].kv_buffer,
-                            tp_index=tp_index,
-                            tp_world_size=dp_world_size,
-                            mode="read")
+            page_io(
+                torch.tensor(mem_indexes, dtype=torch.int64, device="cuda"),
+                page_tensor=cur_page,
+                kv_buffer=dp_mems[tp_index].kv_buffer,
+                tp_index=tp_index,
+                tp_world_size=dp_world_size,
+                mode="read",
+            )
         # keep for debug
         # logger.info(f"dst token tensor {self.kv_buffer[:, mem_indexes[0], 0, 0]}")
         # logger.info(f"dst page token tensor {cur_page[0, :, 0, 0]}")
 
-
     def send_to_decode_node(
         self,
         move_tasks: List[KVMoveTask],

lightllm/server/core/objs/req.py

@@ -254,7 +254,7 @@ def get_all_prompt_metadata(self):
         metadata["prompt_token_ids"] = [int(e) for e in cur_ids]
         self._cache_prompt_metadata = metadata
         return metadata
-    
+
     def is_infer_decode(self) -> bool:
         """
         judge the req is in decode stage

lightllm/server/core/objs/sampling_params.py

@@ -222,7 +222,8 @@ def initialize(self, inputs: Tuple[int, float]):
 
     def to_tuple(self):
         return (self.item0, self.item1)
-    
+
+
 class NodeUUId(ctypes.Structure):
     _pack_ = 4
     _fields_ = [
@@ -236,7 +237,7 @@ def initialize(self, node_id: int):
         return
 
     def get(self) -> int:
-        return ((self.node_id_high << 64) | self.node_id_low)
+        return (self.node_id_high << 64) | self.node_id_low
 
 
 class DecodeNode(ctypes.Structure):
@@ -308,7 +309,7 @@ class SamplingParams(ctypes.Structure):
         ("group_request_id", ctypes.c_int64),  # p d mode used params
         ("suggested_dp_index", ctypes.c_int),  # suggest dp index, deepseekv2 dp mode, use to suggest used dp_index
         ("move_kv_to_decode_node", DecodeNode),  # move kv to deocde node, only used in pd mode
-         # in pd split mode, use to keep the id of pd master
+        # in pd split mode, use to keep the id of pd master
         ("pd_master_node_id", NodeUUId),
         # nixl params object, only used in nixl pd mode, used to build nixl connection in p and d
         ("nixl_params", NIXLParamObj),