Support EPLB balancedness prometheus metric without GPU->CPU synchronize (#15401)

python/sglang/srt/managers/utils.py

@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, List, Optional
 
 import torch
 
+from sglang.srt.eplb.expert_distribution import ExpertDistributionMetrics
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.managers.overlap_utils import FutureIndices
 from sglang.srt.managers.schedule_batch import Req
@@ -44,6 +45,9 @@ class GenerationBatchResult:
     # relay path: forward stream -> next step forward
     next_draft_input: Optional[EagleDraftInput] = None
 
+    # metrics
+    expert_distribution_metrics: Optional[ExpertDistributionMetrics] = None
+
     def copy_to_cpu(self, return_logprob: bool):
         """Copy tensors to CPU in overlap scheduling.
         Only the tensors which are needed for processing results are copied,
@@ -67,6 +71,9 @@ class GenerationBatchResult:
         if self.accept_lens is not None:
             self.accept_lens = self.accept_lens.to("cpu", non_blocking=True)
 
+        if (x := self.expert_distribution_metrics) is not None:
+            x.copy_to_cpu()
+
         self.copy_done.record()
 
     @classmethod