Prefer layer reshape cache event

Author: dante159753

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

@@ -224,7 +224,7 @@ The default metrics configuration contains the following UCM metrics.
 | `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_reshape_cache_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. |
 

examples/metrics/metrics_configs.yaml

@@ -65,7 +65,7 @@ counter:
     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"
+  - name: "dump_event_reshape_cache_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"

ucm/integration/vllm/ucm_connector.py

@@ -836,17 +836,14 @@ def _get_reshape_cache_event(
         if attn_metadata is None:
             return None, "none"
 
-        if self.use_mla:
-            if not layer_name:
-                return None, "mla_missing_layer"
+        if layer_name:
             try:
                 layer_metadata = attn_metadata[layer_name]
                 event = getattr(layer_metadata, "reshape_cache_event", None)
                 if event is not None:
-                    return event, "mla_layer"
+                    return event, "layer"
             except (KeyError, TypeError, AttributeError):
                 pass
-            return None, "mla_missing"
 
         event = getattr(attn_metadata, "reshape_cache_event", None)
         if event is not None:
@@ -876,9 +873,9 @@ def _get_dump_event_handle(
                     ucmmetrics.update_stats(
                         "dump_event_reshape_cache_direct_used_total", 1.0
                     )
-                elif event_source == "mla_layer":
+                elif event_source == "layer":
                     ucmmetrics.update_stats(
-                        "dump_event_reshape_cache_mla_layer_used_total", 1.0
+                        "dump_event_reshape_cache_layer_used_total", 1.0
                     )
                 return event_handle