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

logits: remove blocking H2D copy (#24627)

Browse Source
This commit is contained in:
Yilong Zhao
2026-05-08 13:22:13 -07:00
committed by GitHub
parent 672f778512
commit f30d1d0b0a
1 changed files with 10 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/logits_processor.py
+10 -6
View File
@@ -535,12 +535,15 @@ class LogitsProcessor(nn.Module):
pruned_states_before_norm = torch.cat(pruned_states_before_norm_list) pruned_states_before_norm = torch.cat(pruned_states_before_norm_list)
if aux_pruned_states_lists is not None: if aux_pruned_states_lists is not None:
aux_pruned_states = [torch.cat(lst) for lst in aux_pruned_states_lists] aux_pruned_states = [torch.cat(lst) for lst in aux_pruned_states_lists]
# Build the index tensors via pinned host memory + non-blocking H2D
# so the small copy doesn't drain the stream.
sample_indices = torch.tensor( sample_indices = torch.tensor(
sample_indices, device=pruned_states.device, dtype=torch.int64 sample_indices, dtype=torch.int64, pin_memory=True
) ).to(pruned_states.device, non_blocking=True)
input_logprob_indices = torch.tensor( input_logprob_indices = torch.tensor(
input_logprob_indices, device=pruned_states.device, dtype=torch.int64 input_logprob_indices, dtype=torch.int64, pin_memory=True
) ).to(pruned_states.device, non_blocking=True)
return ( return (
pruned_states, pruned_states,
@@ -609,8 +612,9 @@ class LogitsProcessor(nn.Module):
): ):
pruned_lens = torch.tensor( pruned_lens = torch.tensor(
logits_metadata.extend_logprob_pruned_lens_cpu, logits_metadata.extend_logprob_pruned_lens_cpu,
device=device, dtype=torch.int64,
) pin_memory=True,
).to(device, non_blocking=True)
if logits_metadata.temp_scaled_logprobs: if logits_metadata.temp_scaled_logprobs:
logits_metadata.temperature = torch.repeat_interleave( logits_metadata.temperature = torch.repeat_interleave(
logits_metadata.temperature.view(-1), logits_metadata.temperature.view(-1),