[perf]prepare offset in advance (#188)

* prepare offset in advance

* fix ci problem

* fix mla offset error

* scheduler do not need dataoffset

Author: qyh111

test/test_uc_connector.py

@@ -81,13 +81,19 @@ def setUp(self):
         self.block_size = 8
         self.num_layers = 48
         self.total_blocks_num = 40
+        self.total_tp_size = 2
         self.kv_caches = {}
+        self.k_data_offsets = {}
         for i in range(self.num_layers):
             layer_name = f"model.layers.{i}.self_attn.attn"
             kv_tensor = torch.rand(
                 (2, self.total_blocks_num, self.block_size, 4, 8), dtype=torch.bfloat16
             )
             self.kv_caches[layer_name] = kv_tensor
+        for layer_id in range(self.num_layers):
+            self.k_data_offsets[layer_id] = {}
+            for i in range(self.total_tp_size):
+                self.k_data_offsets[layer_id][i] = 0
 
     def init_uc(
         self, mock_connector, metadata=Mock(), use_layerwise=True
@@ -102,14 +108,16 @@ def init_uc(
             ucconnector.connector = mock_connector
             ucconnector.request_block_infos: dict[str, RequestBlockInfo] = {}
             ucconnector.dump_tasks: dict[str, dict[str, List[Task]]] = {}
-            ucconnector.total_tp_size = 2
+            ucconnector.total_tp_size = self.total_tp_size
             ucconnector._connector_metadata = metadata
             ucconnector.layerwise_load_tasks: dict[
                 str, dict[str, tuple[Task, Task]]
             ] = {}
             ucconnector._need_load_reqs: dict[str, Union[list[int], list[Task]]] = {}
             ucconnector._load_failed_reqs: set[str] = set()
             ucconnector._load_req_to_blocks: dict[str, set[int]] = {}
+            ucconnector.k_data_offsets = self.k_data_offsets
+            ucconnector.min_block_size = 0
         return ucconnector
 
     def test_get_num_new_matched_tokens_hit_all_on_storage(self):

ucm/integration/vllm/uc_connector.py

@@ -109,6 +109,8 @@ def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
         self._need_load_reqs: dict[str, Union[list[int], list[Task]]] = {}
         self._load_failed_reqs: set[str] = set()
         self._load_req_to_blocks: dict[str, set[int]] = {}
+        if role == KVConnectorRole.WORKER:
+            self._initialize_dataoffset(vllm_config)
         if (
             self._vllm_config.kv_transfer_config is not None
             and "ucm_connector_name"
@@ -156,37 +158,34 @@ def _init_kv_caches_from_forward_context(self, forward_context: "ForwardContext"
                     forward_context.virtual_engine
                 ]
 
-    def DataOffset(self, kv_layer, rank, layer_id, is_v):
-        # Non-MLA scene: one layer shape is (2, num_blocks, block_size, num_kv_heads, head_size)
-        # MLA scene: one layer shape is (num_blocks, block_size, head_size)
-        # Element size
-        elem_size = kv_layer[0].element_size()
-        logger.debug(
-            f"total_tp_size = {self.total_tp_size},\n" f"element size = {elem_size}."
+    def _initialize_dataoffset(self, vllm_config: "VllmConfig"):
+        num_kv_heads = vllm_config.model_config.get_num_kv_heads(
+            vllm_config.parallel_config
+        )
+        head_size = vllm_config.model_config.get_head_size()
+        self.min_block_size = (
+            self.block_size * num_kv_heads * head_size * self.element_size
         )
-        # One block size
-        k_min_data_block_size = (
-            kv_layer[0][0].numel() if not self.is_mla else kv_layer[0].numel()
-        ) * elem_size
-        v_min_data_block_size = (
-            kv_layer[1][0].numel() if not self.is_mla else 0
-        ) * elem_size
-        # When tp > 1 layer_size = (k_min_data_block_size + v_min_data_block_size) * tp_size
         layer_size = (
-            k_min_data_block_size + v_min_data_block_size
-        ) * self.total_tp_size
-        if is_v:
-            # Offset of v = Offset of k + k_min_data_block_size
-            return int(
-                self.DataOffset(kv_layer, rank, layer_id, False) + k_min_data_block_size
-            )
-        if self.is_mla:
-            return int(layer_size * layer_id)
-        else:
-            # Offset of k = layer_size * layer_id + layer_size / tp_size * current rank
-            return int(
-                layer_size * layer_id + layer_size / self.total_tp_size * self.rank
-            )
+            self.min_block_size * 2 * self.total_tp_size
+            if not self.is_mla
+            else self.min_block_size
+        )
+        # layer_id -> rank -> k_offset
+        self.k_data_offsets: dict[int, dict[int, int]] = {}
+
+        for layer_id in range(self.num_layers):
+            self.k_data_offsets[layer_id] = {}
+            for rank in range(self.total_tp_size):
+                if self.is_mla:
+                    self.k_data_offsets[layer_id][0] = layer_size * layer_id
+                    break
+                else:
+                    offset = (
+                        layer_size * layer_id
+                        + (layer_size // self.total_tp_size) * rank
+                    )
+                    self.k_data_offsets[layer_id][rank] = offset
 
     def get_tensor_and_offset_layerwise(
         self, vllm_block_ids: List[int], kv_layer: torch.Tensor, layer_name: str
@@ -198,14 +197,17 @@ def get_tensor_and_offset_layerwise(
         layer_id = self._extract_layer_index(layer_name)
 
         for blk_id in vllm_block_ids:
-            k_data_offset = self.DataOffset(kv_layer, self.rank, layer_id, False)
             if self.is_mla:
+                k_data_offset = self.k_data_offsets[layer_id][0]
                 k_tensors.append(kv_layer[blk_id])
             else:
+                k_data_offset = self.k_data_offsets[layer_id][self.rank]
                 k_tensors.append(kv_layer[0][blk_id])
             k_offsets.append(k_data_offset)
             if not self.is_mla:
-                v_data_offset = self.DataOffset(kv_layer, self.rank, layer_id, True)
+                v_data_offset = (
+                    self.k_data_offsets[layer_id][self.rank] + self.min_block_size
+                )
                 v_tensors.append(kv_layer[1][blk_id])
                 v_offsets.append(v_data_offset)
         return k_tensors + v_tensors, k_offsets + v_offsets