diff --git a/python/sglang/srt/sampling/sampling_batch_info.py b/python/sglang/srt/sampling/sampling_batch_info.py index a8f17c754..7eafc7669 100644 --- a/python/sglang/srt/sampling/sampling_batch_info.py +++ b/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, diff --git a/test/registered/unit/sampling/test_sampling_batch_info.py b/test/registered/unit/sampling/test_sampling_batch_info.py index bc923932c..562da7995 100644 --- a/test/registered/unit/sampling/test_sampling_batch_info.py +++ b/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."""