[Perf] Overlap result D2H copy with the next forward step (#29075)

python/sglang/srt/eplb/expert_distribution.py

@@ -52,8 +52,10 @@ _OutputMode = Literal["file", "object"]
 class ExpertDistributionMetrics:
     eplb_balancedness: torch.Tensor
 
-    def copy_to_cpu(self):
-        self.eplb_balancedness = self.eplb_balancedness.to("cpu", non_blocking=True)
+    def map_device_tensors(self, fn):
+        # Device-tensor fields only; caller injects the copy+safety primitive
+        # (see GenerationBatchResult.copy_to_cpu).
+        self.eplb_balancedness = fn(self.eplb_balancedness)
 
 
 class ExpertDistributionRecorder(ABC):