perf(sampling): avoid GPU syncs when applying custom logit processors (#39234)

2026-09-18 17:09:30 -07:00
parent b876213548
commit 8ea0ee300d
7 changed files with 395 additions and 158 deletions
@@ -27,7 +27,10 @@ from sglang.srt.sampling.custom_logit_processor import (
    Qwen3ThinkingBudgetLogitProcessor,
    _cache_from_str,
 )
-from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
+from sglang.srt.sampling.sampling_batch_info import (
+    ProcessorEntry,
+    SamplingBatchInfo,
+)
 from sglang.test.test_utils import CustomTestCase


@@ -66,7 +69,11 @@ class TestApplyCustomLogitProcessor(CustomTestCase):
            vocab_size=4,
            has_custom_logit_processor=True,
            custom_params=params,
-            custom_logit_processor={0: (processor, torch.tensor([True, False, True]))},
+            custom_logit_processor={
+                0: ProcessorEntry(
+                    processor=processor, rows=[0, 2], indices=torch.tensor([0, 2])
+                )
+            },
            device="cpu",
        )
        logits = torch.zeros(batch_size * num_tokens, 4)