[opt]Add event to sync (#768)

## Purpose
This PR adds a config switch `enable_event_sync` to control whether dump
tasks wait on a prerequisite compute event.
## Modifications 
Add `enable_event_sync`plumbing in `ucm_connector.py`.
Pass event handle to `dump_data` only when sync is enabled; otherwise
pass 0.
Add `CopyStream::WaitEvent(void* event)` to apply wait on all internal
copy streams.
Update `DumpQueue` to issue prerequisite wait once via `CopyStream`,
ensuring all streams are covered.
Make `Trans::Stream::WaitEvent` pure virtual and implement it in
SimuStream (no-op).
## Test
<img width="1309" height="759" alt="image"
src="https://github.com/user-attachments/assets/f9fd0590-154f-4399-b64a-3592fde511df"
/>

<img width="1345" height="825" alt="image"
src="https://github.com/user-attachments/assets/ae79846a-05e5-4c85-9d63-776027ebca61"
/>
<img width="933" height="356" alt="image"
src="https://github.com/user-attachments/assets/65f356b6-367e-4454-81e7-3f2386a46c6c"
/>

Author: qyh111

examples/ucm_config_example.yaml

@@ -14,6 +14,7 @@ ucm_connectors:
       storage_backends: "/mnt/test"
       io_direct: false
 
+enable_event_sync: false
 # Enable UCM metrics so they can be monitored online via Grafana and Prometheus.
 # metrics_config_path: "/workspace/unified-cache-management/examples/metrics/metrics_configs.yaml"
 

ucm/integration/vllm/device.py

@@ -0,0 +1,113 @@
+# -*- coding: utf-8 -*-
+"""
+Event-based sync between Python compute stream and C++ cache stream.
+
+When dump_data is called, the cache's C++ stream does D2H from device memory.
+We must ensure the Python compute stream has finished writing KVCache before
+the cache reads. Event sync: record event on compute stream, pass to C++,
+cache stream waits for event before D2H. This avoids blocking the CPU.
+"""
+from abc import ABC, abstractmethod
+from typing import Optional
+
+import torch
+from vllm.platforms import current_platform
+
+from ucm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class Device(ABC):
+    def __init__(self):
+        self.events = []
+
+    @abstractmethod
+    def get_event_handle(self) -> int:
+        """Return event handle for stream sync. 0 means no event (use synchronize instead)."""
+        pass
+
+    @abstractmethod
+    def synchronize(self):
+        pass
+
+    @abstractmethod
+    def destroy_event_handles(self):
+        pass
+
+
+class CudaDevice(Device):
+    def __init__(self):
+        super().__init__()
+
+    def get_event_handle(self) -> int:
+        try:
+            cuda_event = torch.cuda.Event(enable_timing=False)
+            stream = torch.cuda.current_stream()
+            cuda_event.record(stream)
+            handle = int(cuda_event.cuda_event)
+            if handle is None or handle == 0:
+                return 0
+            self.events.append(cuda_event)
+            return handle
+        except Exception as e:
+            logger.error(f"get cuda event handle failed. {e}")
+            return 0
+
+    def synchronize(self):
+        torch.cuda.current_stream().synchronize()
+
+    def destroy_event_handles(self):
+        self.events.clear()
+
+
+class NpuDevice(Device):
+    def __init__(self):
+        super().__init__()
+
+    def get_event_handle(self) -> int:
+        import acl
+        import torch_npu
+
+        try:
+            stream = torch_npu.npu.current_stream().npu_stream
+            event, ret = acl.rt.create_event()
+            if ret != 0:
+                logger.error(f"acl create_event failed: {ret}")
+                return 0
+            self.events.append(event)
+            ret = acl.rt.record_event(event, stream)
+            if ret != 0:
+                logger.error(f"acl record_event failed: {ret}")
+                return 0
+            handle = int(event)
+            if not handle:
+                return 0
+            return handle
+        except Exception as e:
+            logger.error(f"get npu event handle failed. {e}")
+            return 0
+
+    def synchronize(self):
+        torch.npu.current_stream().synchronize()
+
+    def destroy_event_handles(self):
+        import acl
+
+        for event in self.events:
+            try:
+                acl.rt.destroy_event(event)
+            except Exception as e:
+                logger.error(f"destroy npu event failed. {e}")
+                continue
+        self.events.clear()
+
+
+def create_device() -> Optional[Device]:
+    if current_platform.is_cuda_alike():
+        return CudaDevice()
+
+    if current_platform.device_type == "npu":
+        return NpuDevice()
+
+    return None

ucm/integration/vllm/ucm_connector.py

@@ -21,6 +21,7 @@
 from vllm.platforms import current_platform
 from vllm.v1.core.sched.output import SchedulerOutput
 
+from ucm.integration.vllm.device import create_device
 from ucm.logger import init_logger
 from ucm.observability import PrometheusStatsLogger
 from ucm.shared.metrics import ucmmetrics
@@ -207,6 +208,7 @@ def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
         self.launch_config = ucm_config.get_config()
         logger.info(f"self.launch_config: {self.launch_config}")
         self.connector_configs = self.launch_config.get("ucm_connectors", [])
+        self.enable_event_sync = self.launch_config.get("enable_event_sync", False)
         assert len(self.connector_configs) > 0, "no storage connector name in config."
 
         self.chunk_size = self.block_size
@@ -236,12 +238,6 @@ def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
                 f"metrics_config_path: {self.metrics_config}, set worker_id: {worker_id}"
             )
 
-        self.synchronize = lambda: (
-            torch.cuda.current_stream().synchronize()
-            if current_platform.is_cuda_alike()
-            else torch.npu.current_stream().synchronize()
-        )
-
         # invlalid block ids due to load errors
         self._invalid_block_ids: set[int] = set()
 
@@ -319,6 +315,9 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         }
 
         self.store = self._create_store(self.kv_cache_layout)
+        self.device = create_device()
+        if self.device is None:
+            raise RuntimeError(f"Unsupported device platform for UCMDirectConnector.")
 
     def get_num_new_matched_tokens(
         self,
@@ -547,6 +546,16 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
     def wait_for_layer_load(self, layer_name: str) -> None:
         pass
 
+    def _get_dump_event_handle(self) -> int:
+        if not self.enable_event_sync:
+            self.device.synchronize()
+            return 0
+
+        event_handle = self.device.get_event_handle()
+        if event_handle == 0:
+            self.device.synchronize()
+        return event_handle
+
     def save_kv_layer(
         self,
         layer_name: str,
@@ -588,10 +597,10 @@ def wait_for_save(self) -> None:
             total_ptrs = total_ptrs.reshape(total_ptrs.shape[0], -1)
             shard_indexs = [0] * len(total_ucm_block_ids)
             try:
-                self.synchronize()
+                event_handle = self._get_dump_event_handle()
                 save_start_time = time.perf_counter() * 1000
                 task = self.store.dump_data(
-                    total_ucm_block_ids, shard_indexs, total_ptrs
+                    total_ucm_block_ids, shard_indexs, total_ptrs, event_handle
                 )
                 dump_tasks.append(task)
             except RuntimeError as e:
@@ -726,9 +735,9 @@ def save_kv_layer(
             shard_indexs = [layer_id] * len(total_ucm_block_ids)
             try:
                 layer_ptrs = np.ascontiguousarray(total_ptrs[:, layer_id, :])
-                self.synchronize()
+                event_handle = self._get_dump_event_handle()
                 task = self.store.dump_data(
-                    total_ucm_block_ids, shard_indexs, layer_ptrs
+                    total_ucm_block_ids, shard_indexs, layer_ptrs, event_handle
                 )
                 self.dump_tasks[layer_name] = task
             except RuntimeError as e:
@@ -745,6 +754,8 @@ def wait_for_save(self) -> None:
             logger.error(f"wait for dump kv cache failed. {e}")
         self.dump_tasks.clear()
         self.is_save = False
+        if self.enable_event_sync:
+            self.device.destroy_event_handles()
 
 
 class UCMCPConnector(UCMLayerWiseConnector):

ucm/shared/trans/ascend/ascend_stream.cc

@@ -175,4 +175,12 @@ Status AscendStream::Synchronized()
     return Status{ret, std::to_string(ret)};
 }
 
-} // namespace UC::Trans
+Status AscendStream::WaitEvent(void* event)
+{
+    if (event == nullptr) { return Status::OK(); }
+    auto ret = aclrtStreamWaitEvent(stream_, static_cast<aclrtEvent>(event));
+    if (ret != ACL_SUCCESS) [[unlikely]] { return Status{ret, std::to_string(ret)}; }
+    return Status::OK();
+}
+
+}  // namespace UC::Trans

ucm/shared/trans/ascend/ascend_stream.h

@@ -57,8 +57,9 @@ class AscendStream : public Stream {
 
     Status AppendCallback(std::function<void(bool)> cb) override;
     Status Synchronized() override;
+    Status WaitEvent(void* event) override;
 };
 
-} // namespace UC::Trans
+}  // namespace UC::Trans
 
 #endif

ucm/shared/trans/cuda/cuda_stream.cc

@@ -155,4 +155,12 @@ Status Trans::CudaStream::Synchronized()
     return Status::OK();
 }
 
-} // namespace UC::Trans
+Status Trans::CudaStream::WaitEvent(void* event)
+{
+    if (event == nullptr) { return Status::OK(); }
+    auto ret = cudaStreamWaitEvent(stream_, static_cast<cudaEvent_t>(event), 0);
+    if (ret != cudaSuccess) [[unlikely]] { return Status{ret, cudaGetErrorString(ret)}; }
+    return Status::OK();
+}
+
+}  // namespace UC::Trans

ucm/shared/trans/cuda/cuda_stream.h

@@ -52,8 +52,9 @@ class CudaStream : public Stream {
 
     Status AppendCallback(std::function<void(bool)> cb) override;
     Status Synchronized() override;
+    Status WaitEvent(void* event) override;
 };
 
-} // namespace UC::Trans
+}  // namespace UC::Trans
 
 #endif

ucm/shared/trans/simu/simu_stream.cc

@@ -157,6 +157,12 @@ Status SimuStream::AppendCallback(std::function<void(bool)> cb)
     return Status::OK();
 }
 
+Status SimuStream::WaitEvent(void* event)
+{
+    (void)event;
+    return Status::OK();
+}
+
 Status SimuStream::Synchronized()
 {
     std::mutex mutex;

ucm/shared/trans/simu/simu_stream.h

@@ -62,6 +62,7 @@ class SimuStream : public Stream {
     Status HostToDeviceAsync(void* host, void* device[], size_t size, size_t number) override;
 
     Status AppendCallback(std::function<void(bool)> cb) override;
+    Status WaitEvent(void* event) override;
     Status Synchronized() override;
 };
 

ucm/shared/trans/stream.h

@@ -50,8 +50,9 @@ class Stream {
 
     virtual Status AppendCallback(std::function<void(bool)> cb) = 0;
     virtual Status Synchronized() = 0;
+    virtual Status WaitEvent(void* event) = 0;
 };
 
-} // namespace UC::Trans
+}  // namespace UC::Trans
 
 #endif