diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 10df1914a..143054cd6 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -3044,7 +3044,10 @@ 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) + batch_result.copy_to_cpu( + return_logprob=batch.return_logprob, + return_hidden_states=batch.return_hidden_states, + ) else: batch_result.future_indices = future_indices @@ -3149,7 +3152,10 @@ 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) + 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 a277f9b79..1a7224a74 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -59,7 +59,7 @@ class GenerationBatchResult: fpm_start_event: Optional[torch.cuda.Event] = None fpm_end_event: Optional[torch.cuda.Event] = None - def copy_to_cpu(self, return_logprob: bool): + 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 @@ -88,7 +88,7 @@ class GenerationBatchResult: v.to("cpu", non_blocking=True) 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: + if return_hidden_states and self.logits_output.hidden_states is not None: self.logits_output.hidden_states = self.logits_output.hidden_states.to( "cpu", non_blocking=True )