Fuse SamplingBatchInfo tensor construction into one pass

from_schedule_batch built temperature/top_p/top_k/min_p (+seed) with
4-5 separate list comprehensions and one synchronous H2D copy each,
plus 4 more passes for the is_all_greedy/need_* flags. Collect
everything in a single pass over reqs and upload the float params as
one pinned non-blocking H2D copy (disjoint device views of one
buffer; filter/merge only index and cat, producing fresh tensors, so
the shared buffer is safe), int32 top_k and optional int64 seeds as
their own pinned copies.

B300 (torch 2.11 cu130), scheduler-thread blocking time per call:
  bs=8: 38.8 -> 18.2 us (2.1x); bs=32: 1.9x; bs=200: 1.2x
CPU-only construction at bs=200: 2890 -> 1387 us (2.1x).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

python/sglang/srt/sampling/sampling_batch_info.py

@@ -10,6 +10,7 @@ import sglang.srt.sampling.penaltylib as penaltylib
 from sglang.srt.sampling.custom_logit_processor import CustomLogitProcessor
 from sglang.srt.sampling.sampling_params import TOP_K_ALL
 from sglang.srt.server_args import get_global_server_args
+from sglang.srt.utils.common import is_pin_memory_available
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import ScheduleBatch
@@ -73,39 +74,59 @@ class SamplingBatchInfo:
 
         reqs = batch.reqs
         device = batch.device
-        temperatures = torch.tensor(
-            [r.sampling_params.temperature for r in reqs],
-            dtype=torch.float,
-            device=device,
-        ).view(-1, 1)
-        top_ps = torch.tensor(
-            [r.sampling_params.top_p for r in reqs], dtype=torch.float, device=device
-        )
-        top_ks = torch.tensor(
-            [r.sampling_params.top_k for r in reqs], dtype=torch.int32, device=device
-        )
-        min_ps = torch.tensor(
-            [r.sampling_params.min_p for r in reqs], dtype=torch.float, device=device
+        bs = len(reqs)
+        pin = is_pin_memory_available(device)
+
+        # Collect every per-request sampling param and the aggregate flags in a
+        # single pass, then upload the float params with one pinned H2D copy
+        # instead of one synchronous copy per parameter. The device tensors are
+        # disjoint slices of one buffer; filter/merge only index and cat them,
+        # which always produces fresh tensors, so sharing a buffer is safe.
+        fused_float = [0.0] * (3 * bs)
+        top_ks_cpu = [0] * bs
+        seeds_cpu = [0] * bs if enable_deterministic else None
+        is_all_greedy = True
+        need_top_p_sampling = False
+        need_top_k_sampling = False
+        need_min_p_sampling = False
+        has_logit_bias = False
+        for i, r in enumerate(reqs):
+            sp = r.sampling_params
+            top_k = sp.top_k
+            min_p = sp.min_p
+            fused_float[i] = sp.temperature
+            fused_float[bs + i] = sp.top_p
+            fused_float[2 * bs + i] = min_p
+            top_ks_cpu[i] = top_k
+            if enable_deterministic:
+                seeds_cpu[i] = (
+                    sp.sampling_seed if sp.sampling_seed is not None else 42
+                )
+            is_all_greedy &= top_k <= 1
+            need_top_p_sampling |= sp.top_p != 1.0
+            need_top_k_sampling |= top_k != TOP_K_ALL
+            need_min_p_sampling |= min_p > 0
+            has_logit_bias |= sp.logit_bias is not None
+
+        fused_float_device = torch.tensor(
+            fused_float, dtype=torch.float, pin_memory=pin
+        ).to(device, non_blocking=True)
+        temperatures = fused_float_device[:bs].view(-1, 1)
+        top_ps = fused_float_device[bs : 2 * bs]
+        min_ps = fused_float_device[2 * bs :]
+        top_ks = torch.tensor(top_ks_cpu, dtype=torch.int32, pin_memory=pin).to(
+            device, non_blocking=True
         )
         sampling_seed = (
-            torch.tensor(
-                [
-                    (
-                        r.sampling_params.sampling_seed
-                        if r.sampling_params.sampling_seed is not None
-                        else 42
-                    )
-                    for r in reqs
-                ],
-                dtype=torch.int64,
-                device=device,
+            torch.tensor(seeds_cpu, dtype=torch.int64, pin_memory=pin).to(
+                device, non_blocking=True
             )
             if enable_deterministic
             else None
         )
 
         logit_bias = None
-        if any(r.sampling_params.logit_bias is not None for r in reqs):
+        if has_logit_bias:
             logit_bias = torch.zeros(len(reqs), vocab_size, device=device)
             for i, r in enumerate(reqs):
                 if r.sampling_params.logit_bias is not None:
@@ -168,10 +189,10 @@ class SamplingBatchInfo:
             top_ks=top_ks,
             min_ps=min_ps,
             sampling_seed=sampling_seed,
-            is_all_greedy=all(r.sampling_params.top_k <= 1 for r in reqs),
-            need_top_p_sampling=any(r.sampling_params.top_p != 1.0 for r in reqs),
-            need_top_k_sampling=any(r.sampling_params.top_k != TOP_K_ALL for r in reqs),
-            need_min_p_sampling=any(r.sampling_params.min_p > 0 for r in reqs),
+            is_all_greedy=is_all_greedy,
+            need_top_p_sampling=need_top_p_sampling,
+            need_top_k_sampling=need_top_k_sampling,
+            need_min_p_sampling=need_min_p_sampling,
             vocab_size=vocab_size,
             penalizer_orchestrator=penalizer_orchestrator,
             has_custom_logit_processor=has_custom_logit_processor,

test/registered/unit/sampling/test_sampling_batch_info.py

@@ -464,6 +464,47 @@ class TestFromScheduleBatch(CustomTestCase):
         self.assertAlmostEqual(info.top_ps[0].item(), 0.9, places=5)
         self.assertEqual(info.top_ks[0].item(), 50)
 
+    @patch("sglang.srt.sampling.sampling_batch_info.get_global_server_args")
+    def test_mixed_batch_tensor_values_dtypes_shapes(self, mock_server_args):
+        """All per-request params land in the right slot with the right dtype,
+        and the fused float tensors are independent after filtering."""
+        mock_server_args.return_value.enable_deterministic_inference = False
+        mock_server_args.return_value.enable_custom_logit_processor = False
+
+        params = [
+            dict(temp=0.5, top_p=0.7, top_k=10, min_p=0.2),
+            dict(temp=1.0, top_p=1.0, top_k=1, min_p=0.0),
+            dict(temp=1.3, top_p=0.95, top_k=TOP_K_ALL, min_p=0.05),
+        ]
+        batch = MagicMock()
+        batch.reqs = [self._make_req(**p) for p in params]
+        batch.device = DEVICE
+        info = SamplingBatchInfo.from_schedule_batch(batch, VOCAB_SIZE)
+
+        self.assertEqual(info.temperatures.shape, (3, 1))
+        self.assertEqual(info.temperatures.dtype, torch.float32)
+        self.assertEqual(info.top_ps.shape, (3,))
+        self.assertEqual(info.top_ps.dtype, torch.float32)
+        self.assertEqual(info.min_ps.shape, (3,))
+        self.assertEqual(info.top_ks.dtype, torch.int32)
+        for i, p in enumerate(params):
+            self.assertAlmostEqual(info.temperatures[i, 0].item(), p["temp"], places=5)
+            self.assertAlmostEqual(info.top_ps[i].item(), p["top_p"], places=5)
+            self.assertEqual(info.top_ks[i].item(), p["top_k"])
+            self.assertAlmostEqual(info.min_ps[i].item(), p["min_p"], places=5)
+        self.assertFalse(info.is_all_greedy)
+        self.assertTrue(info.need_top_p_sampling)
+        self.assertTrue(info.need_top_k_sampling)
+        self.assertTrue(info.need_min_p_sampling)
+
+        # Filtering must yield independent tensors with the right values even
+        # though construction shares one fused buffer.
+        keep = torch.tensor([2, 0], dtype=torch.int64)
+        info.filter_batch([2, 0], keep)
+        self.assertAlmostEqual(info.temperatures[0, 0].item(), 1.3, places=5)
+        self.assertAlmostEqual(info.top_ps[1].item(), 0.7, places=5)
+        self.assertEqual(info.top_ks[0].item(), TOP_K_ALL)
+
     @patch("sglang.srt.sampling.sampling_batch_info.get_global_server_args")
     def test_greedy_detection(self, mock_server_args):
         """Test that top_k=1 sets is_all_greedy=True."""