fix code style

Author: qyh111

ucm/integration/vllm/ucm_connector.py

@@ -21,7 +21,7 @@
 from ucm.shared.metrics import ucmmonitor
 from ucm.shared.metrics.observability import UCMStatsLogger
 from ucm.store.factory import UcmConnectorFactory
-from ucm.store.ucmstore_v1 import UcmKVStoreBaseV1, Task
+from ucm.store.ucmstore_v1 import Task, UcmKVStoreBaseV1
 from ucm.utils import Config
 
 if TYPE_CHECKING:
@@ -394,20 +394,25 @@ def _extract_layer_index(layer_name: str) -> Optional[int]:
                 return int(chunk)
         return None
 
-    def _get_tensors(self, vllm_block_id: int) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
+    def _get_tensors(
+        self, vllm_block_id: int
+    ) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
         """
         GQA/MHA: one layer shape is (2, num_blocks, block_size, num_kv_heads, head_size)
         MLA: one layer shape is (num_blocks, block_size, head_size)
         """
         k_tensors, v_tensors = [], []
         for _, kv_layer in self.kv_caches.items():
-            k_tensors.append(kv_layer[vllm_block_id] if self.is_mla else kv_layer[0][vllm_block_id])
+            k_tensors.append(
+                kv_layer[vllm_block_id] if self.is_mla else kv_layer[0][vllm_block_id]
+            )
             if not self.is_mla:
                 v_tensors.append(kv_layer[1][vllm_block_id])
         return k_tensors, v_tensors
 
-
-    def _generate_task(self, vllm_block_ids: List[int], ucm_block_ids: List[bytes]) -> Tuple[List[bytes], List[int], List[List[torch.Tensor]]]:
+    def _generate_task(
+        self, vllm_block_ids: List[int], ucm_block_ids: List[bytes]
+    ) -> Tuple[List[bytes], List[int], List[List[torch.Tensor]]]:
         """
         GQA/MHA: one layer shape is (2, num_blocks, block_size, num_kv_heads, head_size)
         MLA: one layer shape is (num_blocks, block_size, head_size)
@@ -416,7 +421,7 @@ def _generate_task(self, vllm_block_ids: List[int], ucm_block_ids: List[bytes])
         for i, vllm_block_id in enumerate(vllm_block_ids):
             k_tensors, v_tensors = self._get_tensors(vllm_block_id)
             block_ids.append(ucm_block_ids[i])
-            tensors.append(k_tensors+v_tensors)
+            tensors.append(k_tensors + v_tensors)
             shard_indexs.append(0)
 
         return block_ids, shard_indexs, tensors

ucm/store/factory.py

@@ -60,4 +60,4 @@ def create_connector(cls, connector_name: str, config: dict) -> UcmKVStoreBaseV1
 
 UcmConnectorFactory.register_connector(
     "UcmPcStore", "ucm.store.pcstore.pcstore_connector", "UcmPcStore"
-)
+)

ucm/store/pcstore/cc/domain/trans/trans_task.h

@@ -61,7 +61,7 @@ class TransTask {
     void Append(const std::string& block, const std::vector<uintptr_t>& addresses)
     {
         grouped_[block] = addresses;
-        number_ += addresses.size(); 
+        number_ += addresses.size();
     }
     auto Str() const noexcept { return fmt::format("{},{},{}", id, brief_, number_); }
     size_t GroupNumber() const { return grouped_.size(); }

ucm/store/pcstore/pcstore_connector.py

@@ -69,14 +69,20 @@ def prefetch(self, block_ids: List[bytes]) -> None:
         pass
 
     def load(
-        self, block_ids: List[bytes], shard_index: List[int], dst_tensor: List[List[torch.Tensor]]
+        self,
+        block_ids: List[bytes],
+        shard_index: List[int],
+        dst_tensor: List[List[torch.Tensor]],
     ) -> Task:
         dst_tensor_ptrs = [[t.data_ptr() for t in tensors] for tensors in dst_tensor]
         task_id = self.store.LoadToDevice(block_ids, dst_tensor_ptrs)
         return PcTask(task_id=task_id)
 
     def dump(
-        self, block_ids: List[bytes], shard_index: List[int], src_tensor: List[List[torch.Tensor]]
+        self,
+        block_ids: List[bytes],
+        shard_index: List[int],
+        src_tensor: List[List[torch.Tensor]],
     ) -> Task:
         src_tensor_ptrs = [[t.data_ptr() for t in tensors] for tensors in src_tensor]
         task_id = self.store.DumpFromDevice(block_ids, src_tensor_ptrs)

ucm/store/test/e2e/pcstore_embed.py

@@ -62,14 +62,18 @@ def make_buffers(
     return hashes, tensors
 
 
-def embed(store: UcmKVStoreBaseV1, hashes: List[bytes], tensors: List[List[torch.Tensor]]):
+def embed(
+    store: UcmKVStoreBaseV1, hashes: List[bytes], tensors: List[List[torch.Tensor]]
+):
     shard_index = [0] * len(hashes)
     task = store.dump(hashes, shard_index, tensors)
     assert task.task_id > 0
     store.wait(task)
 
 
-def fetch(store: UcmKVStoreBaseV1, hashes: List[bytes], tensors: List[List[torch.Tensor]]):
+def fetch(
+    store: UcmKVStoreBaseV1, hashes: List[bytes], tensors: List[List[torch.Tensor]]
+):
     founds = store.lookup(hashes)
     for found in founds:
         assert found

ucm/store/test/e2e/pcstore_fetch.py

@@ -70,7 +70,9 @@ def make_buffers(device_id, batch_size, block_dim, block_len, block_layer):
     return tensors
 
 
-def fetch(store: UcmKVStoreBaseV1, hashes: List[bytes], tensors: List[List[torch.Tensor]]):
+def fetch(
+    store: UcmKVStoreBaseV1, hashes: List[bytes], tensors: List[List[torch.Tensor]]
+):
     founds = store.lookup(hashes)
     for found in founds:
         assert found