[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):
+    def map_device_tensors(self, fn):
-        self.eplb_balancedness = self.eplb_balancedness.to("cpu", non_blocking=True)
+        # 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):

python/sglang/srt/managers/scheduler.py

@@ -3253,10 +3253,15 @@ class Scheduler(
                                 else batch_result.next_token_ids
                             )
                             self.future_map.stash(future_indices, stash_payload)
-                            batch_result.copy_to_cpu(
+                            # Result D2H on copy_stream overlaps the next forward
-                                return_logprob=batch.return_logprob,
+                            # instead of serializing on forward_stream; it's a leaf
-                                return_hidden_states=batch.return_hidden_states,
+                            # gated by copy_done, so nothing on forward_stream waits.
-                            )
+                            self.copy_stream.wait_stream(self.forward_stream)
+                            with self.copy_stream_ctx:
+                                batch_result.copy_to_cpu(
+                                    return_logprob=batch.return_logprob,
+                                    return_hidden_states=batch.return_hidden_states,
+                                )
                         else:
                             batch_result.future_indices = future_indices
 

python/sglang/srt/managers/utils.py

@@ -22,6 +22,19 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+def _async_d2h(t: torch.Tensor) -> torch.Tensor:
+    """Async D2H copy for overlap scheduling. On CUDA the dest is pinned (a D2H
+    to pageable host memory blocks the caller until done) and record_stream keeps
+    the source alive until the copy stream drains, so the caching allocator can't
+    recycle it early. Non-CUDA falls back to a plain copy."""
+    if not t.is_cuda:
+        return t.to("cpu", non_blocking=True)
+    cpu_t = torch.empty(t.shape, dtype=t.dtype, pin_memory=True)
+    cpu_t.copy_(t, non_blocking=True)
+    t.record_stream(torch.cuda.current_stream(t.device))
+    return cpu_t
+
+
 @dataclasses.dataclass
 class GenerationBatchResult:
     logits_output: Optional[LogitsProcessorOutput] = None
@@ -72,6 +85,7 @@ class GenerationBatchResult:
     fpm_start_event: Optional[torch.cuda.Event] = None
     fpm_end_event: Optional[torch.cuda.Event] = None
 
+    @torch.profiler.record_function("copy_result_to_cpu")
     def copy_to_cpu(self, return_logprob: bool, return_hidden_states: bool = True):
         """Copy tensors to CPU in overlap scheduling.
         Only the tensors which are needed for processing results are copied,
@@ -79,45 +93,47 @@ class GenerationBatchResult:
         """
         if return_logprob:
             if self.logits_output.next_token_logprobs is not None:
-                self.logits_output.next_token_logprobs = (
+                self.logits_output.next_token_logprobs = _async_d2h(
-                    self.logits_output.next_token_logprobs.to("cpu", non_blocking=True)
+                    self.logits_output.next_token_logprobs
                 )
             if self.logits_output.input_token_logprobs is not None:
-                self.logits_output.input_token_logprobs = (
+                self.logits_output.input_token_logprobs = _async_d2h(
-                    self.logits_output.input_token_logprobs.to("cpu", non_blocking=True)
+                    self.logits_output.input_token_logprobs
                 )
             if self.logits_output.next_token_top_logprobs_val is not None:
                 self.logits_output.next_token_top_logprobs_val = [
-                    v.to("cpu", non_blocking=True) if torch.is_tensor(v) else v
+                    _async_d2h(v) if torch.is_tensor(v) else v
                     for v in self.logits_output.next_token_top_logprobs_val
                 ]
             if self.logits_output.next_token_top_logprobs_idx is not None:
                 self.logits_output.next_token_top_logprobs_idx = [
-                    x.to("cpu", non_blocking=True) if torch.is_tensor(x) else x
+                    _async_d2h(x) if torch.is_tensor(x) else x
                     for x in self.logits_output.next_token_top_logprobs_idx
                 ]
             if self.logits_output.next_token_token_ids_logprobs_val is not None:
                 self.logits_output.next_token_token_ids_logprobs_val = [
-                    v.to("cpu", non_blocking=True) if torch.is_tensor(v) else v
+                    _async_d2h(v) if torch.is_tensor(v) else v
                     for v in self.logits_output.next_token_token_ids_logprobs_val
                 ]
         if return_hidden_states and self.logits_output.hidden_states is not None:
-            self.logits_output.hidden_states = self.logits_output.hidden_states.to(
+            self.logits_output.hidden_states = _async_d2h(
-                "cpu", non_blocking=True
+                self.logits_output.hidden_states
             )
-        self.next_token_ids = self.next_token_ids.to("cpu", non_blocking=True)
+        self.next_token_ids = _async_d2h(self.next_token_ids)
 
         if self.accept_lens is not None:
-            self.accept_lens = self.accept_lens.to("cpu", non_blocking=True)
+            self.accept_lens = _async_d2h(self.accept_lens)
 
-        if self.routed_experts_output is not None:
+        # Sub-objects only declare their device fields; the single copy+safety
-            self.routed_experts_output.copy_to_cpu()
+        # primitive (_async_d2h: pinned D2H + record_stream) is injected here so
-
+        # all device->host copying and lifetime safety lives in one place.
-        if self.indexer_topk_output is not None:
+        for holder in (
-            self.indexer_topk_output.copy_to_cpu()
+            self.routed_experts_output,
-
+            self.indexer_topk_output,
-        if (x := self.expert_distribution_metrics) is not None:
+            self.expert_distribution_metrics,
-            x.copy_to_cpu()
+        ):
+            if holder is not None:
+                holder.map_device_tensors(_async_d2h)
 
         self.copy_done.record()
 
@@ -243,30 +259,27 @@ class EmbeddingBatchResult:
     def can_run_cuda_graph(self) -> bool:
         return False
 
+    @torch.profiler.record_function("copy_embedding_to_cpu")
     def copy_to_cpu(self):
         """Copy embeddings and pooled hidden states to CPU for overlap scheduling."""
         if isinstance(self.embeddings, torch.Tensor):
             self.copy_done = torch.get_device_module(self.embeddings.device).Event()
-            self.embeddings = self.embeddings.to("cpu", non_blocking=True)
+            self.embeddings = _async_d2h(self.embeddings)
         else:
             assert isinstance(self.embeddings, list)
             if len(self.embeddings) == 0:
                 return
 
             self.copy_done = torch.get_device_module(self.embeddings[0].device).Event()
-            self.embeddings = [
+            self.embeddings = [_async_d2h(emb) for emb in self.embeddings]
-                emb.to("cpu", non_blocking=True) for emb in self.embeddings
-            ]
 
         if self.pooled_hidden_states is not None:
             if isinstance(self.pooled_hidden_states, list):
                 self.pooled_hidden_states = [
-                    t.to("cpu", non_blocking=True) for t in self.pooled_hidden_states
+                    _async_d2h(t) for t in self.pooled_hidden_states
                 ]
             else:
-                self.pooled_hidden_states = self.pooled_hidden_states.to(
+                self.pooled_hidden_states = _async_d2h(self.pooled_hidden_states)
-                    "cpu", non_blocking=True
-                )
 
         self.copy_done.record()
 

python/sglang/srt/state_capturer/base.py

@@ -79,17 +79,19 @@ class BaseHostCache:
 @dataclasses.dataclass
 class TopkCaptureOutput:
     """Holds GPU tensors captured during forward for overlap scheduling.
-    Call copy_to_cpu() inside forward stream (before copy_done.record()),
+    map_device_tensors() D2H-copies them before copy_done.record() (may run on
-    then finalize() after copy_done.synchronize().
+    the dedicated result-copy stream); finalize() runs after copy_done.synchronize().
     """
 
     out_cache_loc: torch.Tensor
     topk: torch.Tensor
     host_cache: BaseHostCache
 
-    def copy_to_cpu(self):
+    def map_device_tensors(self, fn):
-        self.out_cache_loc = self.out_cache_loc.to("cpu", non_blocking=True)
+        # Device-tensor fields only; caller injects the copy+safety primitive
-        self.topk = self.topk.to("cpu", non_blocking=True)
+        # (see GenerationBatchResult.copy_to_cpu).
+        self.out_cache_loc = fn(self.out_cache_loc)
+        self.topk = fn(self.topk)
 
     def finalize(self):
         self.host_cache.buffer[self.out_cache_loc] = self.topk