Add metrics for having prefill and decode in different ranks (#15752)

python/sglang/srt/utils/device_timer.py

@@ -1,23 +1,23 @@
 from collections import deque
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import Callable, Deque, Optional
+from typing import Callable, Deque, Dict, Optional
 
 import torch
 
 
 class DeviceTimer:
-    def __init__(self, reporter: Callable[[str, float], None]):
+    def __init__(self, reporter: Callable):
         self._intervals: Deque[_TimingInterval] = deque()
         self._reporter = reporter
 
     @contextmanager
-    def wrap(self, category: str):
+    def wrap(self, metadata: Dict):
         self._intervals.append(_TimingInterval.create())
         try:
             yield
         finally:
-            self._intervals[-1].end(category=category)
+            self._intervals[-1].end(metadata=metadata)
             self._report()
 
     def _report(self):
@@ -27,14 +27,14 @@ class DeviceTimer:
                 break
 
             self._intervals.popleft()
-            self._reporter(interval.category, interval.elapsed_time() / 1000.0)
+            self._reporter(t=interval.elapsed_time() / 1000.0, **interval.metadata)
 
 
 @dataclass
 class _TimingInterval:
     start_event: torch.cuda.Event
     end_event: Optional[torch.cuda.Event] = None
-    category: Optional[str] = None
+    metadata: Optional[Dict] = None
 
     @staticmethod
     def create():
@@ -42,13 +42,13 @@ class _TimingInterval:
         start_event.record()
         return _TimingInterval(start_event=start_event)
 
-    def end(self, category: str):
+    def end(self, metadata: Dict):
         end_event = torch.cuda.Event(enable_timing=True)
         end_event.record()
 
         assert self.end_event is None
         self.end_event = end_event
-        self.category = category
+        self.metadata = metadata
 
     def elapsed_time(self) -> float:
         return self.start_event.elapsed_time(self.end_event)