logits: remove blocking H2D copy (#24627)

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