Use pinned memory for asynchronous sampling metadata transfers (#39777)

Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
Co-authored-by: hnyls2002 <lsyincs@gmail.com>

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(
+                token_ids_tensor = torch.tensor(
-                    logprobs.device, non_blocking=True
+                    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(
+            token_ids_tensor = torch.tensor(
-                logprobs.device, non_blocking=True
+                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)])

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_(

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(
+        pin_memory = is_pin_memory_available(self.orchestrator.device)
-            data=[
+        self.min_new_tokens = (
-                req.sampling_params.min_new_tokens for req in self.orchestrator.reqs()
+            torch.tensor(
-            ],
+                data=[
-            dtype=torch.int32,
+                    req.sampling_params.min_new_tokens
-            device=self.orchestrator.device,
+                    for req in self.orchestrator.reqs()
-        ).unsqueeze_(1)
+                ],
+                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,

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_(

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_(

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():
+                    # Dedup on int(key) first ("1" and "01" collide): duplicate
-                        logit_bias[i, int(key)] = value
+                    # 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 = (

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