nixl pd support qwen3.5 (#1340)

Author: hiworldwzj

lightllm/common/kv_cache_mem_manager/deepseek2_mem_manager.py

@@ -56,7 +56,10 @@ def write_mem_to_page_kv_move_buffer(
         dp_index: int,
         mem_managers: List["MemoryManager"],
         dp_world_size: int,
+        page_kind: str = "kv",
+        req_idx: int = None,
     ):
+        assert page_kind == "kv", f"{type(self).__name__} does not support page_kind={page_kind}"
         cur_page = self.kv_move_buffer[page_index]
         pin_mem_indexes = self._buffer_mem_indexes_tensors[page_index][0 : len(mem_indexes)]
         pin_mem_indexes.numpy()[:] = mem_indexes
@@ -77,7 +80,10 @@ def read_page_kv_move_buffer_to_mem(
         dp_index: int,
         mem_managers: List["MemoryManager"],
         dp_world_size: int,
+        page_kind: str = "kv",
+        req_idx: int = None,
     ):
+        assert page_kind == "kv", f"{type(self).__name__} does not support page_kind={page_kind}"
         cur_page = self.kv_move_buffer[page_index]
         pin_mem_indexes = self._buffer_mem_indexes_tensors[page_index][0 : len(mem_indexes)]
         pin_mem_indexes.numpy()[:] = mem_indexes

lightllm/common/kv_cache_mem_manager/mem_manager.py

@@ -102,9 +102,6 @@ def alloc_kv_move_buffer(self, max_req_total_len):
         return
 
     def alloc_paged_kv_move_buffer(self, page_num, page_size) -> torch.Tensor:
-        if isinstance(self, MemoryManager) and type(self) is not 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"
@@ -121,7 +118,10 @@ def write_mem_to_page_kv_move_buffer(
         dp_index: int,
         mem_managers: List["MemoryManager"],
         dp_world_size: int,
+        page_kind: str = "kv",
+        req_idx: int = None,
     ):
+        assert page_kind == "kv", f"{type(self).__name__} does not support page_kind={page_kind}"
         cur_page = self.kv_move_buffer[page_index]
         pin_mem_indexes = self._buffer_mem_indexes_tensors[page_index][0 : len(mem_indexes)]
         pin_mem_indexes.numpy()[:] = mem_indexes
@@ -150,7 +150,10 @@ def read_page_kv_move_buffer_to_mem(
         dp_index: int,
         mem_managers: List["MemoryManager"],
         dp_world_size: int,
+        page_kind: str = "kv",
+        req_idx: int = None,
     ):
+        assert page_kind == "kv", f"{type(self).__name__} does not support page_kind={page_kind}"
         cur_page = self.kv_move_buffer[page_index]
         pin_mem_indexes = self._buffer_mem_indexes_tensors[page_index][0 : len(mem_indexes)]
         pin_mem_indexes.numpy()[:] = mem_indexes