From 4414db594cc2b06971de0c8f3a4d3743aed40876 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Sat, 27 Jun 2026 09:07:37 +0800 Subject: [PATCH] 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 | 14 +++++- python/sglang/srt/managers/utils.py | 44 +++++++++++++------ .../eagle_draft_cuda_graph_runner.py | 35 ++++++++++----- .../sglang/srt/speculative/eagle_worker_v2.py | 12 ++--- .../test_generation_batch_result_copy.py | 35 +++++++++++++++ 5 files changed, 109 insertions(+), 31 deletions(-) create mode 100644 test/registered/unit/managers/test_generation_batch_result_copy.py diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 5dccdfe11..7f3f53ef3 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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 diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index 67ee9786c..b584a8d17 100644 --- a/python/sglang/srt/managers/utils.py +++ b/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() diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index fccb286e7..d2bfa1dc2 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/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 diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 0bdfe7209..f68ec3b8d 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/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). diff --git a/test/registered/unit/managers/test_generation_batch_result_copy.py b/test/registered/unit/managers/test_generation_batch_result_copy.py new file mode 100644 index 000000000..a54c1d2a1 --- /dev/null +++ b/test/registered/unit/managers/test_generation_batch_result_copy.py @@ -0,0 +1,35 @@ +import torch + +from sglang.srt.layers.logits_processor import LogitsProcessorOutput +from sglang.srt.managers.utils import GenerationBatchResult + + +class _DummyEvent: + def __init__(self): + self.recorded = False + + def record(self): + self.recorded = True + + +class _HiddenStateSentinel: + def to(self, *args, **kwargs): + raise AssertionError("hidden_states should not be copied when disabled") + + +def test_copy_to_cpu_skips_hidden_states_when_not_requested(): + hidden_states = _HiddenStateSentinel() + copy_done = _DummyEvent() + result = GenerationBatchResult( + logits_output=LogitsProcessorOutput( + next_token_logits=None, + hidden_states=hidden_states, + ), + next_token_ids=torch.tensor([1], dtype=torch.int64), + ) + result.copy_done = copy_done + + result.copy_to_cpu(return_logprob=False, return_hidden_states=False) + + assert result.logits_output.hidden_states is hidden_states + assert copy_done.recorded