From f30d1d0b0a5584f25cbfaa7b7671621af867e695 Mon Sep 17 00:00:00 2001 From: Yilong Zhao <74357408+happierpig@users.noreply.github.com> Date: Fri, 8 May 2026 13:22:13 -0700 Subject: [PATCH] logits: remove blocking H2D copy (#24627) --- python/sglang/srt/layers/logits_processor.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 0f13dc9eb..e14895e97 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -535,12 +535,15 @@ class LogitsProcessor(nn.Module): pruned_states_before_norm = torch.cat(pruned_states_before_norm_list) if aux_pruned_states_lists is not None: 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, 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, device=pruned_states.device, dtype=torch.int64 - ) + input_logprob_indices, dtype=torch.int64, pin_memory=True + ).to(pruned_states.device, non_blocking=True) return ( pruned_states, @@ -609,8 +612,9 @@ class LogitsProcessor(nn.Module): ): pruned_lens = torch.tensor( 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: logits_metadata.temperature = torch.repeat_interleave( logits_metadata.temperature.view(-1),