revert layerwise

Author: qyh111

ucm/integration/vllm/ucm_connector.py

@@ -722,6 +722,8 @@ def _track_async_dump_requests(
         self,
         requests_dispatch_meta: dict[str, RequestDispatchMeta],
     ) -> None:
+        if self.use_layerwise:
+            return
         self._async_dump_req_ids.update(
             request_id
             for request_id, dispatch_meta in requests_dispatch_meta.items()
@@ -858,6 +860,8 @@ def _flush_pending_dump_tasks(self, request_ids: Optional[set[str]] = None) -> N
         self._pending_dump_tasks = remaining_tasks
 
     def handle_preemptions(self, kv_connector_metadata: KVConnectorMetadata):
+        if self.use_layerwise:
+            return
         preempted_req_ids = getattr(kv_connector_metadata, "preempted_req_ids", None)
         if preempted_req_ids:
             self._flush_pending_dump_tasks(preempted_req_ids)
@@ -948,6 +952,8 @@ def request_finished(
         request: "Request",
         block_ids: list[int],
     ) -> tuple[bool, dict[str, Any] | None]:
+        if self.use_layerwise:
+            return False, None
         if request.request_id in self._async_dump_req_ids:
             self._async_dump_req_ids.discard(request.request_id)
             return True, None
@@ -966,6 +972,8 @@ def get_finished(
         self,
         finished_req_ids: set[str],
     ) -> tuple[Optional[set[str]], Optional[set[str]]]:
+        if self.use_layerwise:
+            return None, None
         async_finished_req_ids = finished_req_ids & self._async_dump_req_ids
 
         if async_finished_req_ids:
@@ -1160,46 +1168,58 @@ def save_kv_layer(
             total_vllm_block_ids.extend(vllm_block_ids)
 
         if dump_request_ids:
-            self._async_dump_req_ids.update(dump_request_ids)
             if self.dump_total_ptrs is None:
                 self.dump_total_ptrs = self.kv_cache_layout.extract_block_addrs(
                     total_vllm_block_ids, layer_first=True
                 )
             shard_indexs = [layer_id] * len(total_ucm_block_ids)
-            event_handle = 0
             try:
                 layer_ptrs = np.ascontiguousarray(self.dump_total_ptrs[local_layer_id])
                 event_handle = self._get_dump_event_handle()
                 task = self.store.dump_data(
                     total_ucm_block_ids, shard_indexs, layer_ptrs, event_handle
                 )
-                pending_dump_task = PendingDumpTask(
-                    task=task,
-                    request_ids=set(dump_request_ids),
-                    event_handle=event_handle,
-                )
-                self._pending_dump_tasks.append(pending_dump_task)
+                self.dump_tasks[layer_name] = task
             except Exception as e:
                 logger.error(f"submit dump task failed. {type(e).__name__}: {e}")
-                if self.enable_event_sync and event_handle and self.device is not None:
-                    self.device.destroy_event_handle(event_handle)
         if self.is_save:
             submit_end = time.perf_counter()
             ucmmetrics.update_stats(
                 {"layerwise_save_submit_ms": (submit_end - submit_start) * 1000}
             )
 
     def wait_for_save(self) -> None:
-        if self._connector_metadata:
-            metadata = self._get_connector_metadata()
-            self._async_dump_req_ids.update(
-                request_id
-                for request_id, request in metadata.request_meta.items()
-                if len(request.dump_block_ids[0]) > 0
-            )
+        if not self.is_save:
+            total_end = time.perf_counter()
+            if self._layerwise_batch_start is not None:
+                batch_total_ms = (total_end - self._layerwise_batch_start) * 1000
+                ucmmetrics.update_stats({"layerwise_batch_total_ms": batch_total_ms})
+                self._layerwise_batch_start = None
+            return
+
+        total_start = time.perf_counter()
+        try:
+            for layer_name in self.kv_caches:
+                if layer_name not in self.dump_tasks:
+                    continue
+                self.store.wait(self.dump_tasks[layer_name])
+        except Exception as e:
+            logger.error(f"wait for dump kv cache failed. {type(e).__name__}: {e}")
+
+        total_end = time.perf_counter()
+        stats = {"layerwise_save_tail_total_ms": (total_end - total_start) * 1000}
+        if self._layerwise_batch_start is not None:
+            stats["layerwise_batch_total_ms"] = (
+                total_end - self._layerwise_batch_start
+            ) * 1000
+            self._layerwise_batch_start = None
+        ucmmetrics.update_stats(stats)
+
         self.dump_tasks.clear()
         self.is_save = False
         self.dump_total_ptrs = None
+        if self.enable_event_sync:
+            self.device.destroy_event_handles()
 
 
 class UCMCPConnector(UCMLayerWiseConnector):