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