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),