Add opt-in reshape cache event sync

Author: dante159753

docs/source/user-guide/metrics/metrics.md

@@ -223,6 +223,10 @@ The default metrics configuration contains the following UCM metrics.
 | `ucm:posix_h2s_bytes_total` | Total bytes transferred from host buffer to Posix storage. |
 | `ucm:load_bytes_total` | Total bytes loaded through the UCM connector. |
 | `ucm:save_bytes_total` | Total bytes saved through the UCM connector. |
+| `ucm:dump_event_reshape_cache_direct_used_total` | Dump submissions synchronized with `attn_metadata.reshape_cache_event`. |
+| `ucm:dump_event_reshape_cache_mla_layer_used_total` | Dump submissions synchronized with `attn_metadata[layer_name].reshape_cache_event`. |
+| `ucm:dump_event_current_stream_used_total` | Dump submissions synchronized with a UCM-recorded current stream event fallback. |
+| `ucm:dump_event_sync_fallback_used_total` | Dump submissions that fell back to device synchronization because no event handle was available. |
 
 ### Gauges
 

docs/source/user-guide/prefix-cache/pipeline_store.md

@@ -216,6 +216,7 @@ ucm_connectors:
       use_gdr: false
 enable_event_sync: true
 use_layerwise: true
+enable_reshape_cache_event_sync: false
 enable_record_traces: false
 use_lite: false
 persist_token_threshold: 0
@@ -320,4 +321,4 @@ This log indicates that the **Posix Store** has received a **load or dump task**
 ```text
 [UC][D] Posix task({task_id},{operation},{subtask_number},{size}) finished, cost {time}ms. [PID,TID]
 ```
-This log indicates that a load or dump task in the **Posix Store** has completed, along with its execution time in **in ms**.
+This log indicates that a load or dump task in the **Posix Store** has completed, along with its execution time in **in ms**.

examples/metrics/metrics_configs.yaml

@@ -63,6 +63,14 @@ counter:
     documentation: "Total bytes loaded through the UCM connector (summed across all start_load_kv calls)"
   - name: "save_bytes_total"
     documentation: "Total bytes saved through the UCM connector (summed across all wait_for_save calls)"
+  - name: "dump_event_reshape_cache_direct_used_total"
+    documentation: "Number of dump submissions synchronized with attn_metadata.reshape_cache_event"
+  - name: "dump_event_reshape_cache_mla_layer_used_total"
+    documentation: "Number of dump submissions synchronized with attn_metadata[layer_name].reshape_cache_event"
+  - name: "dump_event_current_stream_used_total"
+    documentation: "Number of dump submissions synchronized with a UCM-recorded current stream event fallback"
+  - name: "dump_event_sync_fallback_used_total"
+    documentation: "Number of dump submissions that fell back to device synchronization because no event handle was available"
 
 # Gauge metrics configuration
 gauge:

examples/ucm_config_example.yaml

@@ -26,6 +26,9 @@ ucm_connectors:
 
 # When you use UcmNfsStore, you should set enable_event_sync to false.
 enable_event_sync: true
+# Use event from vllm ascend to do D2H immediatly after kv cache dump.
+# Enable for better dump performance.
+enable_reshape_cache_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/ucm_connector.py

@@ -364,6 +364,9 @@ def __init__(
         self.launch_config = ucm_config.get_config()
         self.connector_configs = self.launch_config.get("ucm_connectors", [])
         self.enable_event_sync = self.launch_config.get("enable_event_sync", True)
+        self.enable_reshape_cache_event_sync = self.launch_config.get(
+            "enable_reshape_cache_event_sync", False
+        )
         self.enable_record_traces = self.launch_config.get(
             "enable_record_traces", False
         )
@@ -857,36 +860,37 @@ def _get_dump_event_handle(
     ) -> 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}."
+        if self.enable_reshape_cache_event_sync:
+            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
             )
-            return event_handle
+            if event_handle != 0:
+                if event_source == "direct":
+                    ucmmetrics.update_stats(
+                        "dump_event_reshape_cache_direct_used_total", 1.0
+                    )
+                elif event_source == "mla_layer":
+                    ucmmetrics.update_stats(
+                        "dump_event_reshape_cache_mla_layer_used_total", 1.0
+                    )
+                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}."
+            ucmmetrics.update_stats(
+                "dump_event_sync_fallback_used_total", 1.0
             )
         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}."
+            ucmmetrics.update_stats(
+                "dump_event_current_stream_used_total", 1.0
             )
         return event_handle