minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[perf] avoid hidden states d2h when return_hidden_states=false (#25155)

Browse Source
This commit is contained in:
Qiaolin Yu
2026-05-13 23:47:17 -07:00
committed by GitHub
parent 1f119f6a44
commit 421179c453
2 changed files with 10 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler.py
+8 -2
View File
@@ -3044,7 +3044,10 @@ class Scheduler(
batch_result.copy_done = self.device_module.Event() batch_result.copy_done = self.device_module.Event()
if batch_result.delay_sample_func is None: if batch_result.delay_sample_func is None:
self.future_map.store_to_map(future_indices, batch_result) 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: else:
batch_result.future_indices = future_indices batch_result.future_indices = future_indices
@@ -3149,7 +3152,10 @@ class Scheduler(
_batch_result = batch_result.delay_sample_func() _batch_result = batch_result.delay_sample_func()
assert _batch_result is batch_result assert _batch_result is batch_result
self.future_map.store_to_map(batch_result.future_indices, 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. # Release the closure and large GPU tensors that are no longer needed.
# The delay_sample_func closure captures forward_batch (which holds # The delay_sample_func closure captures forward_batch (which holds
python/sglang/srt/managers/utils.py
+2 -2
View File
@@ -59,7 +59,7 @@ class GenerationBatchResult:
fpm_start_event: Optional[torch.cuda.Event] = None fpm_start_event: Optional[torch.cuda.Event] = None
fpm_end_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. """Copy tensors to CPU in overlap scheduling.
Only the tensors which are needed for processing results are copied, Only the tensors which are needed for processing results are copied,
e.g., next_token_ids, logits outputs 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 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 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( self.logits_output.hidden_states = self.logits_output.hidden_states.to(
"cpu", non_blocking=True "cpu", non_blocking=True
) )