Reduce spec-v2 overlap copy overhead

Spec-v2 decode was paying unnecessary scheduler-side copy and replay costs: hidden states were copied back to CPU even when not requested, overlap result D2H ran on the forward stream, and draft graph replay issued several small copies separately.

Port the focused upstream optimizations without taking the larger spec-v2 relay refactor: gate hidden-state D2H on return_hidden_states, run result copies on copy_stream with pinned async D2H, use stride arange for draft select_index, and group small draft graph replay copies while keeping the large hidden-state DMA copy separate.

Constraint: Current branch carries local GLM/NSA/PD/CP changes, so upstream spec-v2 large refactors are not safe to cherry-pick wholesale.

Rejected: Cherry-pick upstream spec-v2 relay/dataclass refactor | too broad for this performance fix and conflicts with current branch structure

Rejected: Copy hidden states unconditionally | default chat output does not consume them after draft extend

Confidence: medium

Scope-risk: moderate

Directive: Keep result-copy lifetime tied to copy_done; do not move copy_to_cpu back onto forward_stream without measuring decode throughput.

Tested: local py_compile for changed production files and new test

Tested: local git diff --check

Tested: remote g0034 cjy-glm5-new as ubuntu on /sgl-workspace/sglang-tai: pytest -q -p no:cacheprovider test_generation_batch_result_copy.py test_eagle_worker_v2_cp_hidden.py test_spec_utils.py

Not-tested: full two-node decode throughput A/B after restart

python/sglang/srt/managers/scheduler.py

@@ -3115,7 +3115,12 @@ class Scheduler(
                     batch_result.copy_done = self.device_module.Event()
                     if batch_result.delay_sample_func is None:
                         self.future_map.store_to_map(future_indices, batch_result)
-                        batch_result.copy_to_cpu(return_logprob=batch.return_logprob)
+                        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
 
@@ -3222,7 +3227,12 @@ class Scheduler(
             _batch_result = batch_result.delay_sample_func()
             assert _batch_result is batch_result
             self.future_map.store_to_map(batch_result.future_indices, batch_result)
-            batch_result.copy_to_cpu(return_logprob=self.cur_batch.return_logprob)
+            self.copy_stream.wait_stream(self.forward_stream)
+            with self.copy_stream_ctx:
+                batch_result.copy_to_cpu(
+                    return_logprob=self.cur_batch.return_logprob,
+                    return_hidden_states=self.cur_batch.return_hidden_states,
+                )
 
         # Release the closure and large GPU tensors that are no longer needed.
         # The delay_sample_func closure captures forward_batch (which holds

python/sglang/srt/managers/utils.py

@@ -21,6 +21,23 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+def _async_d2h(t: torch.Tensor) -> torch.Tensor:
+    """Copy a tensor to CPU without blocking the caller when possible.
+
+    CUDA D2H to pageable host memory can synchronize the caller.  Use a pinned
+    destination and record the source tensor on the current stream so the CUDA
+    caching allocator cannot recycle it before the copy stream drains.
+    """
+
+    if not t.is_cuda:
+        return t.to("cpu", non_blocking=True)
+
+    cpu_t = torch.empty(t.shape, dtype=t.dtype, device="cpu", 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
@@ -49,43 +66,44 @@ class GenerationBatchResult:
     # metrics
     expert_distribution_metrics: Optional[ExpertDistributionMetrics] = None
 
-    def copy_to_cpu(self, return_logprob: bool):
+    @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,
         e.g., next_token_ids, logits outputs
         """
         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.to("cpu", non_blocking=True)
+                self.logits_output.next_token_logprobs = _async_d2h(
+                    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.to("cpu", non_blocking=True)
+                self.logits_output.input_token_logprobs = _async_d2h(
+                    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 self.logits_output.hidden_states is not None:
-            self.logits_output.hidden_states = self.logits_output.hidden_states.to(
-                "cpu", non_blocking=True
+        if return_hidden_states and self.logits_output.hidden_states is not None:
+            self.logits_output.hidden_states = _async_d2h(
+                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 (x := self.expert_distribution_metrics) is not None:
             x.copy_to_cpu()

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -11,6 +11,7 @@ from sglang.srt.model_executor.cuda_graph_runner import (
     CUDA_GRAPH_CAPTURE_FAILED_MSG,
     CudaGraphRunner,
     DeepEPCudaGraphRunnerAdapter,
+    _grouped_foreach_copy_,
     get_batch_sizes_to_capture,
     get_global_graph_memory_pool,
     log_input_buffer_sizes,
@@ -382,16 +383,28 @@ class EAGLEDraftCudaGraphRunner:
 
         num_tokens = bs * self.num_tokens_per_bs
 
-        # Common inputs
-        buffers.seq_lens[:raw_bs].copy_(forward_batch.seq_lens)
-        buffers.out_cache_loc[: raw_num_token * self.speculative_num_steps].copy_(
-            forward_batch.out_cache_loc
+        # Common inputs: batch small device copies by dtype to reduce replay
+        # launch overhead.  hidden_states is large and contiguous, so keep it on
+        # the direct copy_ path.
+        _grouped_foreach_copy_(
+            [
+                buffers.seq_lens[:raw_bs],
+                buffers.out_cache_loc[: raw_num_token * self.speculative_num_steps],
+                buffers.positions[:raw_num_token],
+                buffers.topk_p[:raw_bs],
+                buffers.topk_index[:raw_bs],
+                buffers.req_pool_indices[:raw_bs],
+            ],
+            [
+                forward_batch.seq_lens,
+                forward_batch.out_cache_loc,
+                forward_batch.positions,
+                forward_batch.spec_info.topk_p,
+                forward_batch.spec_info.topk_index,
+                forward_batch.req_pool_indices,
+            ],
         )
-        buffers.positions[:raw_num_token].copy_(forward_batch.positions)
-        buffers.topk_p[:raw_bs].copy_(forward_batch.spec_info.topk_p)
-        buffers.topk_index[:raw_bs].copy_(forward_batch.spec_info.topk_index)
         buffers.hidden_states[:raw_bs].copy_(forward_batch.spec_info.hidden_states)
-        buffers.req_pool_indices[:raw_bs].copy_(forward_batch.req_pool_indices)
 
         # TODO(ch-wan): support num_token_non_padded
         if self.require_gathered_buffer:
@@ -415,9 +428,9 @@ class EAGLEDraftCudaGraphRunner:
         # while replay metadata and graph kernels observe the padded fake rows.
         raw_seq_lens_sum = forward_batch.seq_lens_sum
         if bs != raw_bs and raw_seq_lens_sum is not None:
-            forward_batch.seq_lens_sum = raw_seq_lens_sum + (
-                bs - raw_bs
-            ) * self.seq_len_fill_value
+            forward_batch.seq_lens_sum = (
+                raw_seq_lens_sum + (bs - raw_bs) * self.seq_len_fill_value
+            )
 
         self.model_runner.draft_attn_backend.init_forward_metadata_replay_cuda_graph(
             forward_batch, bs

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -564,8 +564,12 @@ class EagleDraftWorker(BaseDraftWorker):
             num_tokens_for_logprob_per_req=self.speculative_num_steps + 1,
         )
         select_index = (
-            torch.arange(len(batch.seq_lens), device=self.device)
-            * self.speculative_num_draft_tokens
+            torch.arange(
+                0,
+                len(batch.seq_lens) * self.speculative_num_draft_tokens,
+                self.speculative_num_draft_tokens,
+                device=self.device,
+            )
             + batch_result.accept_lens
             - 1
         )
@@ -704,9 +708,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
         # allocator and kv cache pool are shared with target worker, which are cleared in scheduler
         pass
 
-    def _can_use_cp_draft_shared_kv(
-        self, model_worker_batch: ModelWorkerBatch
-    ) -> bool:
+    def _can_use_cp_draft_shared_kv(self, model_worker_batch: ModelWorkerBatch) -> bool:
         """Whether the CP-local draft-hidden side channel applies to this prefill.
 
         Mirrors ``EAGLEWorker._can_use_cp_draft_shared_kv`` (the legacy v1 path).