From f31a7bd45c6ab86796aa012ebfd2378bc42e1a59 Mon Sep 17 00:00:00 2001 From: Yuxuan Zhang Date: Mon, 21 Sep 2026 06:31:31 +0800 Subject: [PATCH] Use pinned memory for asynchronous sampling metadata transfers (#39777) Co-authored-by: Xinyuan Tong Co-authored-by: hnyls2002 --- python/sglang/srt/layers/logprob_processor.py | 15 +- .../sampling/penaltylib/frequency_penalty.py | 8 +- .../srt/sampling/penaltylib/min_new_tokens.py | 29 ++- .../sampling/penaltylib/presence_penalty.py | 8 +- .../sampling/penaltylib/repetition_penalty.py | 8 +- .../srt/sampling/sampling_batch_info.py | 20 +- .../test_sampling_metadata_staging.py | 181 ++++++++++++++++++ 7 files changed, 245 insertions(+), 24 deletions(-) create mode 100644 test/registered/unit/sampling/test_sampling_metadata_staging.py diff --git a/python/sglang/srt/layers/logprob_processor.py b/python/sglang/srt/layers/logprob_processor.py index f2dcf7d7a..7844ed4e4 100644 --- a/python/sglang/srt/layers/logprob_processor.py +++ b/python/sglang/srt/layers/logprob_processor.py @@ -18,7 +18,7 @@ from sglang.srt.model_executor.runner_utils.pool import ( graph_pool_borrow_largest_run, ) from sglang.srt.runtime_context import get_exec -from sglang.srt.utils.common import async_d2h +from sglang.srt.utils.common import async_d2h, is_pin_memory_available if TYPE_CHECKING: from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessorOutput @@ -151,15 +151,16 @@ def get_token_ids_logprobs_raw( no_copy_to_cpu: bool = False, ): vals, idxs = [], [] + pin_memory = is_pin_memory_available(logprobs.device) if stage == LogprobStage.DECODE: for i, token_ids in enumerate(token_ids_logprobs_list): if token_ids is None: vals.append([]) idxs.append([]) else: - token_ids_tensor = torch.tensor(token_ids, dtype=torch.long).to( - logprobs.device, non_blocking=True - ) + token_ids_tensor = torch.tensor( + token_ids, dtype=torch.long, pin_memory=pin_memory + ).to(logprobs.device, non_blocking=True) row = logprobs[i, token_ids_tensor] vals.append(row if no_copy_to_cpu else row.tolist()) idxs.append(token_ids) @@ -178,9 +179,9 @@ def get_token_ids_logprobs_raw( idxs.append([]) pt += pruned_len continue - token_ids_tensor = torch.tensor(token_ids, dtype=torch.long).to( - logprobs.device, non_blocking=True - ) + token_ids_tensor = torch.tensor( + token_ids, dtype=torch.long, pin_memory=pin_memory + ).to(logprobs.device, non_blocking=True) pos_logprobs = logprobs[pt : pt + pruned_len, token_ids_tensor] vals.append(pos_logprobs if no_copy_to_cpu else pos_logprobs.tolist()) idxs.append([token_ids for _ in range(pruned_len)]) diff --git a/python/sglang/srt/sampling/penaltylib/frequency_penalty.py b/python/sglang/srt/sampling/penaltylib/frequency_penalty.py index 63d838574..5abeaf75c 100644 --- a/python/sglang/srt/sampling/penaltylib/frequency_penalty.py +++ b/python/sglang/srt/sampling/penaltylib/frequency_penalty.py @@ -1,6 +1,7 @@ import torch from sglang.srt.sampling.penaltylib.orchestrator import _BatchedPenalizer +from sglang.srt.utils.common import is_pin_memory_available class BatchedFrequencyPenalizer(_BatchedPenalizer): @@ -15,6 +16,7 @@ class BatchedFrequencyPenalizer(_BatchedPenalizer): ) def _prepare(self): + pin_memory = is_pin_memory_available(self.orchestrator.device) self.cumulated_frequency_penalties = torch.zeros( (len(self.orchestrator.reqs()), self.orchestrator.vocab_size), dtype=torch.float32, @@ -28,9 +30,11 @@ class BatchedFrequencyPenalizer(_BatchedPenalizer): for req in self.orchestrator.reqs() ], dtype=torch.float32, - device=self.orchestrator.device, + pin_memory=pin_memory, ) - ).unsqueeze_(1) + .to(self.orchestrator.device, non_blocking=True) + .unsqueeze_(1) + ) def _cumulate_output_tokens(self, output_ids: torch.Tensor): self.cumulated_frequency_penalties.scatter_add_( diff --git a/python/sglang/srt/sampling/penaltylib/min_new_tokens.py b/python/sglang/srt/sampling/penaltylib/min_new_tokens.py index bf47b92b4..51ad26cf4 100644 --- a/python/sglang/srt/sampling/penaltylib/min_new_tokens.py +++ b/python/sglang/srt/sampling/penaltylib/min_new_tokens.py @@ -1,6 +1,7 @@ import torch from sglang.srt.sampling.penaltylib.orchestrator import _BatchedPenalizer +from sglang.srt.utils.common import is_pin_memory_available class BatchedMinNewTokensPenalizer(_BatchedPenalizer): @@ -14,15 +15,21 @@ class BatchedMinNewTokensPenalizer(_BatchedPenalizer): ) def _prepare(self): - self.min_new_tokens = torch.tensor( - data=[ - req.sampling_params.min_new_tokens for req in self.orchestrator.reqs() - ], - dtype=torch.int32, - device=self.orchestrator.device, - ).unsqueeze_(1) + pin_memory = is_pin_memory_available(self.orchestrator.device) + self.min_new_tokens = ( + torch.tensor( + data=[ + req.sampling_params.min_new_tokens + for req in self.orchestrator.reqs() + ], + dtype=torch.int32, + pin_memory=pin_memory, + ) + .to(self.orchestrator.device, non_blocking=True) + .unsqueeze_(1) + ) - padded_stop_token_ids = torch.nn.utils.rnn.pad_sequence( + padded_stop_token_ids_cpu = torch.nn.utils.rnn.pad_sequence( sequences=[ torch.tensor( data=[ @@ -40,13 +47,17 @@ class BatchedMinNewTokensPenalizer(_BatchedPenalizer): if token_id is not None ], dtype=torch.int64, - device=self.orchestrator.device, ) for req in self.orchestrator.reqs() ], batch_first=True, padding_value=self.orchestrator.vocab_size, ) + if pin_memory: + padded_stop_token_ids_cpu = padded_stop_token_ids_cpu.pin_memory() + padded_stop_token_ids = padded_stop_token_ids_cpu.to( + self.orchestrator.device, non_blocking=True + ) self.stop_token_penalties = torch.zeros( size=(len(self.orchestrator.reqs()), self.orchestrator.vocab_size + 1), dtype=torch.float32, diff --git a/python/sglang/srt/sampling/penaltylib/presence_penalty.py b/python/sglang/srt/sampling/penaltylib/presence_penalty.py index 1c045039e..5d6004eb6 100644 --- a/python/sglang/srt/sampling/penaltylib/presence_penalty.py +++ b/python/sglang/srt/sampling/penaltylib/presence_penalty.py @@ -1,6 +1,7 @@ import torch from sglang.srt.sampling.penaltylib.orchestrator import _BatchedPenalizer +from sglang.srt.utils.common import is_pin_memory_available class BatchedPresencePenalizer(_BatchedPenalizer): @@ -15,6 +16,7 @@ class BatchedPresencePenalizer(_BatchedPenalizer): ) def _prepare(self): + pin_memory = is_pin_memory_available(self.orchestrator.device) self.cumulated_presence_penalties = torch.zeros( (len(self.orchestrator.reqs()), self.orchestrator.vocab_size), dtype=torch.float32, @@ -28,9 +30,11 @@ class BatchedPresencePenalizer(_BatchedPenalizer): for req in self.orchestrator.reqs() ], dtype=torch.float32, - device=self.orchestrator.device, + pin_memory=pin_memory, ) - ).unsqueeze_(1) + .to(self.orchestrator.device, non_blocking=True) + .unsqueeze_(1) + ) def _cumulate_output_tokens(self, output_ids: torch.Tensor): self.cumulated_presence_penalties.scatter_( diff --git a/python/sglang/srt/sampling/penaltylib/repetition_penalty.py b/python/sglang/srt/sampling/penaltylib/repetition_penalty.py index b9ad94f39..7d5891819 100644 --- a/python/sglang/srt/sampling/penaltylib/repetition_penalty.py +++ b/python/sglang/srt/sampling/penaltylib/repetition_penalty.py @@ -2,6 +2,7 @@ import torch from sglang.srt.sampling.penaltylib.orchestrator import _BatchedPenalizer from sglang.srt.utils import get_compiler_backend, is_npu +from sglang.srt.utils.common import is_pin_memory_available _is_npu = is_npu() @@ -29,6 +30,7 @@ class BatchedRepetitionPenalizer(_BatchedPenalizer): ) def _prepare(self): + pin_memory = is_pin_memory_available(self.orchestrator.device) self.cumulated_repetition_penalties = torch.ones( (len(self.orchestrator.reqs()), self.orchestrator.vocab_size), dtype=torch.float32, @@ -41,9 +43,11 @@ class BatchedRepetitionPenalizer(_BatchedPenalizer): for req in self.orchestrator.reqs() ], dtype=torch.float32, - device=self.orchestrator.device, + pin_memory=pin_memory, ) - ).unsqueeze_(1) + .to(self.orchestrator.device, non_blocking=True) + .unsqueeze_(1) + ) def _cumulate_output_tokens(self, output_ids: torch.Tensor): self.cumulated_repetition_penalties.scatter_( diff --git a/python/sglang/srt/sampling/sampling_batch_info.py b/python/sglang/srt/sampling/sampling_batch_info.py index 3d5c69de2..8440c426f 100644 --- a/python/sglang/srt/sampling/sampling_batch_info.py +++ b/python/sglang/srt/sampling/sampling_batch_info.py @@ -158,10 +158,26 @@ class SamplingBatchInfo: logit_bias = None if any(r.sampling_params.logit_bias is not None for r in reqs): logit_bias = torch.zeros(len(reqs), vocab_size, device=device) + rows, cols, vals = [], [], [] for i, r in enumerate(reqs): if r.sampling_params.logit_bias is not None: - for key, value in r.sampling_params.logit_bias.items(): - logit_bias[i, int(key)] = value + # Dedup on int(key) first ("1" and "01" collide): duplicate + # indices make the index_put below nondeterministic on CUDA. + row_bias = { + int(key): value + for key, value in r.sampling_params.logit_bias.items() + } + for token_id, value in row_bias.items(): + rows.append(i) + cols.append(token_id) + vals.append(value) + if rows: + logit_bias[ + _rows_to_device_indices(rows, device, _pin), + _rows_to_device_indices(cols, device, _pin), + ] = torch.tensor(vals, dtype=logit_bias.dtype, pin_memory=_pin).to( + device, non_blocking=True + ) # Check if any request has custom logit processor has_custom_logit_processor = ( diff --git a/test/registered/unit/sampling/test_sampling_metadata_staging.py b/test/registered/unit/sampling/test_sampling_metadata_staging.py new file mode 100644 index 000000000..f469ae0c3 --- /dev/null +++ b/test/registered/unit/sampling/test_sampling_metadata_staging.py @@ -0,0 +1,181 @@ +"""Sampling metadata built on the host and copied to the device in one transfer.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch +from torch.utils._python_dispatch import TorchDispatchMode + +from sglang.srt.sampling.penaltylib import ( + BatchedMinNewTokensPenalizer, + BatchedPenalizerOrchestrator, +) +from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo +from sglang.srt.sampling.sampling_params import SamplingParams +from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=15, suite="base-a-test-cpu") +register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +VOCAB_SIZE = 32 + + +class _Batch: + def __init__(self, reqs, device): + self.reqs = reqs + self.device = device + + +def _req(**sampling_params): + return SimpleNamespace( + sampling_params=SamplingParams(**sampling_params), + eos_token_ids=None, + tokenizer=SimpleNamespace(eos_token_id=None, additional_stop_token_ids=None), + custom_logit_processor=None, + return_sampling_mask=False, + ) + + +class _H2DCopies(TorchDispatchMode): + """Record the element count of every host-to-device copy.""" + + def __init__(self): + super().__init__() + self.numels = [] + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + kwargs = kwargs or {} + if func == torch.ops.aten._to_copy.default: + source = args[0] + target = kwargs.get("device") + if source.device.type == "cpu" and target and target.type == "cuda": + self.numels.append(source.numel()) + return func(*args, **kwargs) + + +class _SamplingMetadataTestBase(CustomTestCase): + device = "cpu" + + def setUp(self): + super().setUp() + exec_context = SimpleNamespace( + deterministic=SimpleNamespace(enable_deterministic_inference=False), + features=SimpleNamespace(enable_custom_logit_processor=True), + ) + context_patch = patch( + "sglang.srt.sampling.sampling_batch_info.get_exec", + return_value=exec_context, + ) + context_patch.start() + self.addCleanup(context_patch.stop) + + def assert_device_tensor(self, actual, expected, dtype=None): + self.assertEqual(actual.device, torch.device(self.device)) + if dtype is not None: + self.assertEqual(actual.dtype, dtype) + torch.testing.assert_close(actual.cpu(), expected, rtol=0, atol=0) + + +class TestSamplingMetadataCPU(_SamplingMetadataTestBase): + def test_min_tokens_pads_ragged_stop_sets(self): + reqs = [ + _req(min_new_tokens=2, stop_token_ids=[3]), + _req(min_new_tokens=0, stop_token_ids=[5]), + _req(min_new_tokens=1), + ] + # Row 0 unions four stop sources and must drop the None entries. + reqs[0].sampling_params.stop_token_ids.add(None) + reqs[0].eos_token_ids = {2} + reqs[0].tokenizer.additional_stop_token_ids = {4, None} + reqs[0].tokenizer.eos_token_id = 1 + orch = BatchedPenalizerOrchestrator( + VOCAB_SIZE, _Batch(reqs, self.device), {BatchedMinNewTokensPenalizer} + ) + self.assert_device_tensor( + orch.penalizers[BatchedMinNewTokensPenalizer].min_new_tokens, + torch.tensor([[2], [0], [1]], dtype=torch.int32), + torch.int32, + ) + for step in range(3): + logits = torch.zeros(3, VOCAB_SIZE, device=self.device) + orch.apply(logits) + expected = torch.zeros(3, VOCAB_SIZE) + if step < 2: + expected[0, [1, 2, 3, 4]] = -torch.inf + self.assert_device_tensor(logits, expected) + orch.cumulate_output_tokens( + torch.ones(3, dtype=torch.long, device=self.device) + ) + + def test_min_tokens_without_any_stop_tokens(self): + orch = BatchedPenalizerOrchestrator( + VOCAB_SIZE, + _Batch([_req(min_new_tokens=1)], self.device), + {BatchedMinNewTokensPenalizer}, + ) + logits = torch.zeros(1, VOCAB_SIZE, device=self.device) + orch.apply(logits) + self.assert_device_tensor(logits, torch.zeros(1, VOCAB_SIZE)) + + def test_sparse_logit_bias_keeps_last_value_for_colliding_keys(self): + reqs = [ + _req(logit_bias={"0": -100, "31": 100, "1": 2, "01": 3}), + _req(), + _req(logit_bias={}), + _req(logit_bias={"0": 0, "2": -1.25}), + ] + info = SamplingBatchInfo.from_schedule_batch( + _Batch(reqs, self.device), VOCAB_SIZE + ) + expected = torch.zeros(len(reqs), VOCAB_SIZE) + expected[0, 0], expected[0, 31], expected[0, 1] = -100, 100, 3 + expected[3, 2] = -1.25 + self.assert_device_tensor(info.logit_bias, expected, torch.float32) + + def test_logit_bias_is_none_without_any_bias(self): + info = SamplingBatchInfo.from_schedule_batch( + _Batch([_req(), _req()], self.device), VOCAB_SIZE + ) + self.assertIsNone(info.logit_bias) + + +@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") +class TestSamplingMetadataCUDA(_SamplingMetadataTestBase): + device = "cuda:0" + + def test_stop_token_copy_count_does_not_grow_with_batch_size(self): + copy_counts = {} + for batch_size in (1, 16): + reqs = [ + _req(min_new_tokens=2, stop_token_ids=[2, 3]) for _ in range(batch_size) + ] + with _H2DCopies() as copies: + orch = BatchedPenalizerOrchestrator( + VOCAB_SIZE, + _Batch(reqs, self.device), + {BatchedMinNewTokensPenalizer}, + ) + logits = torch.zeros(batch_size, VOCAB_SIZE, device=self.device) + orch.apply(logits) + torch.cuda.synchronize() + copy_counts[batch_size] = len(copies.numels) + self.assertTrue(torch.isneginf(logits[:, 2:4]).all().item()) + self.assertEqual(copy_counts[16], copy_counts[1]) + + def test_logit_bias_never_copies_a_dense_row(self): + reqs = [_req(logit_bias={"1": 2, "31": -1}), _req(), _req(logit_bias={})] + with _H2DCopies() as copies: + info = SamplingBatchInfo.from_schedule_batch( + _Batch(reqs, self.device), VOCAB_SIZE + ) + torch.cuda.synchronize() + self.assertLess(max(copies.numels), VOCAB_SIZE) + expected = torch.zeros(len(reqs), VOCAB_SIZE) + expected[0, 1], expected[0, 31] = 2, -1 + self.assert_device_tensor(info.logit_bias, expected, torch.float32) + + +if __name__ == "__main__": + unittest.main()