Use reshape cache event for dump sync

Author: dante159753

ucm/integration/vllm/device.py

@@ -13,7 +13,7 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from itertools import accumulate
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
 from vllm.platforms import current_platform
@@ -26,12 +26,16 @@
 class Device(ABC):
     def __init__(self):
         self.events = {}
+        self.borrowed_events = []
 
     @abstractmethod
     def get_event_handle(self) -> int:
         """Return event handle for stream sync. 0 means no event (use synchronize instead)."""
         pass
 
+    def get_event_handle_from_event(self, event: Any) -> int:
+        return 0
+
     @abstractmethod
     def synchronize(self):
         pass
@@ -114,11 +118,23 @@ def get_event_handle(self) -> int:
             logger.error(f"get cuda event handle failed. {e}")
             return 0
 
+    def get_event_handle_from_event(self, event: Any) -> int:
+        try:
+            handle = getattr(event, "cuda_event", None)
+            if handle is None or int(handle) == 0:
+                return 0
+            self.borrowed_events.append(event)
+            return int(handle)
+        except Exception as e:
+            logger.error(f"get cuda event handle from existing event failed. {e}")
+            return 0
+
     def synchronize(self):
         torch.cuda.current_stream().synchronize()
 
     def destroy_event_handles(self):
         self.events.clear()
+        self.borrowed_events.clear()
 
     def destroy_event_handle(self, handle: int):
         self.events.pop(handle, None)
@@ -248,6 +264,22 @@ def get_event_handle(self) -> int:
             logger.error(f"get npu event handle failed. {e}")
             return 0
 
+    def get_event_handle_from_event(self, event: Any) -> int:
+        try:
+            handle = getattr(event, "npu_event", None)
+            if handle is None:
+                parameter = getattr(event, "_as_parameter_", None)
+                if callable(parameter):
+                    parameter = parameter()
+                handle = getattr(parameter, "value", None)
+            if handle is None or int(handle) == 0:
+                return 0
+            self.borrowed_events.append(event)
+            return int(handle)
+        except Exception as e:
+            logger.error(f"get npu event handle from existing event failed. {e}")
+            return 0
+
     def synchronize(self):
         torch.npu.current_stream().synchronize()
 
@@ -260,6 +292,7 @@ def destroy_event_handles(self):
             except Exception as e:
                 logger.error(f"destroy npu event failed. {e}")
         self.events.clear()
+        self.borrowed_events.clear()
 
     def destroy_event_handle(self, handle: int):
         import acl

ucm/integration/vllm/ucm_connector.py

@@ -323,6 +323,7 @@ def __init__(
         self.tp_rank = self._vllm_config.parallel_config.rank
         self.block_size = self._vllm_config.cache_config.block_size
         self.is_mla = self._vllm_config.model_config.is_deepseek_mla
+        self.use_mla = getattr(self._vllm_config.model_config, "use_mla", self.is_mla)
         self.num_layers = self._vllm_config.model_config.get_num_layers(
             self._vllm_config.parallel_config
         )
@@ -826,14 +827,67 @@ 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:
+    def _get_reshape_cache_event(
+        self, layer_name: Optional[str], attn_metadata: Optional["AttentionMetadata"]
+    ) -> tuple[Optional[Any], str]:
+        if attn_metadata is None:
+            return None, "none"
+
+        if self.use_mla:
+            if not layer_name:
+                return None, "mla_missing_layer"
+            try:
+                layer_metadata = attn_metadata[layer_name]
+                event = getattr(layer_metadata, "reshape_cache_event", None)
+                if event is not None:
+                    return event, "mla_layer"
+            except (KeyError, TypeError, AttributeError):
+                pass
+            return None, "mla_missing"
+
+        event = getattr(attn_metadata, "reshape_cache_event", None)
+        if event is not None:
+            return event, "direct"
+        return None, "direct_missing"
+
+    def _get_dump_event_handle(
+        self,
+        layer_name: Optional[str] = None,
+        attn_metadata: Optional["AttentionMetadata"] = None,
+    ) -> int:
         if not self.enable_event_sync:
             self.device.synchronize()
+            logger.debug("dump event sync disabled; synchronized current stream.")
             return 0
 
+        reshape_cache_event, event_source = self._get_reshape_cache_event(
+            layer_name, attn_metadata
+        )
+        event_handle = (
+            self.device.get_event_handle_from_event(reshape_cache_event)
+            if reshape_cache_event is not None
+            else 0
+        )
+        if event_handle != 0:
+            logger.debug(
+                f"dump event handle from reshape_cache_event source={event_source}, "
+                f"layer_name={layer_name}, handle={event_handle}."
+            )
+            return event_handle
+
         event_handle = self.device.get_event_handle()
         if event_handle == 0:
             self.device.synchronize()
+            logger.debug(
+                f"dump event fallback synchronized current stream, "
+                f"reshape_event_source={event_source}, layer_name={layer_name}."
+            )
+        else:
+            logger.debug(
+                f"dump event handle from current stream fallback, "
+                f"reshape_event_source={event_source}, layer_name={layer_name}, "
+                f"handle={event_handle}."
+            )
         return event_handle
 
     def save_kv_layer(
@@ -1140,7 +1194,7 @@ def save_kv_layer(
             shard_indexs = [layer_id] * len(total_ucm_block_ids)
             try:
                 layer_ptrs = np.ascontiguousarray(self.dump_total_ptrs[local_layer_id])
-                event_handle = self._get_dump_event_handle()
+                event_handle = self._get_dump_event_handle(layer_name, attn_metadata)
                 task = self.store.dump_data(
                     total_ucm_block_ids, shard_indexs, layer_ptrs, event_handle
                 )