diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 0dc1727a8..ea77294d2 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -143,6 +143,7 @@ from sglang.srt.runtime_context import ( get_serving, get_spec, ) +from sglang.srt.sampling.custom_logit_processor import supports_sampling_mask from sglang.srt.sampling.sampling_params import SamplingParams from sglang.srt.server_args import ( PortArgs, @@ -1259,6 +1260,17 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): "The server is not configured to enable custom logit processor. " "Please set `--enable-custom-logit-processor` to enable this feature." ) + if ( + obj.return_sampling_mask + and obj.custom_logit_processor + and not supports_sampling_mask(obj.custom_logit_processor) + ): + # Reject before scheduling so aborted requests cannot execute + # unsupported processors during sampling batch preparation. + raise ValueError( + "return_sampling_mask only supports DisallowedTokensLogitsProcessor " + "among custom logit processors." + ) def _validate_mm_limits( self, obj: Union[GenerateReqInput, EmbeddingReqInput] diff --git a/python/sglang/srt/sampling/custom_logit_processor.py b/python/sglang/srt/sampling/custom_logit_processor.py index d66d872bd..48ea39661 100644 --- a/python/sglang/srt/sampling/custom_logit_processor.py +++ b/python/sglang/srt/sampling/custom_logit_processor.py @@ -58,6 +58,17 @@ class DisallowedTokensLogitsProcessor(CustomLogitProcessor): return logits +def supports_sampling_mask(serialized_processor: str) -> bool: + """Hard exclusion preserves the relative logits needed for mask-based replay.""" + try: + return isinstance( + CustomLogitProcessor.from_str(serialized_processor), + DisallowedTokensLogitsProcessor, + ) + except Exception: + return False + + def _open_thinking_start(ids: list[int], start_id: int, end_id: int) -> int: """Return the index of the start token of the currently open thinking block, or -1.""" for idx in reversed(range(len(ids))): diff --git a/test/registered/sampling/test_sampling_mask.py b/test/registered/sampling/test_sampling_mask.py index aeb9ec6b8..be27da24d 100644 --- a/test/registered/sampling/test_sampling_mask.py +++ b/test/registered/sampling/test_sampling_mask.py @@ -9,6 +9,11 @@ import torch from sglang.srt.layers import sampler as sampler_module from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.layers.sampler import Sampler +from sglang.srt.sampling.custom_logit_processor import ( + DisallowedTokensLogitsProcessor, + Qwen3ThinkingBudgetLogitProcessor, +) +from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo from sglang.srt.utils import is_hip, kill_process_tree from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.test_utils import ( @@ -30,6 +35,7 @@ _SAMPLING_SEED = 1234 _SERVER_ARGS = ( "--mem-fraction-static", "0.7", + "--enable-custom-logit-processor", ) _INVALID_SAMPLING_MASK_ERROR = ( "top_p-only sampling is valid but can return huge masks in the tail" @@ -41,6 +47,64 @@ class TestSamplingMaskCapture(CustomTestCase): self.sampler = Sampler.__new__(Sampler) torch.nn.Module.__init__(self.sampler) + def test_hard_exclusion_replay_in_mixed_batch(self): + backends = ["pytorch"] if is_hip() else ["pytorch", "flashinfer"] + for backend in backends: + with self.subTest(backend=backend): + logits = ( + torch.tensor([[0.3, 0.2, 0.5, 0.15, 0.1]], device="cuda") + .log() + .repeat(2, 1) + ) + original = logits.clone() + info = SamplingBatchInfo( + temperatures=torch.ones(2, 1, device="cuda"), + top_ps=torch.full((2,), 0.9, device="cuda"), + top_ks=torch.full((2,), 3, dtype=torch.int32, device="cuda"), + min_ps=torch.zeros(2, device="cuda"), + is_all_greedy=False, + is_any_greedy=False, + need_top_p_sampling=True, + need_top_k_sampling=True, + need_min_p_sampling=False, + vocab_size=5, + has_custom_logit_processor=True, + custom_params=[{"token_ids": [2]}, None], + custom_logit_processor={ + 0: ( + DisallowedTokensLogitsProcessor(), + torch.tensor([True, False], device="cuda"), + ) + }, + return_sampling_masks=[True, True], + ) + logits = self.sampler._preprocess_logits(logits, info) + with patch( + "sglang.srt.layers.sampler.get_exec", + return_value=SimpleNamespace( + kernel=SimpleNamespace(sampling_backend=backend) + ), + ): + sampled, capture = self.sampler._sample_from_probs( + logits.softmax(-1), + info, + positions=torch.zeros(2, dtype=torch.int64, device="cuda"), + simple_sampling_case=False, + return_sampling_mask=True, + ) + output = LogitsProcessorOutput(next_token_logits=None) + self.sampler._attach_sampling_mask_to_output( + output, info, sampled, capture + ) + support = output.next_token_sampling_mask_idx[0] + self.assertEqual(set(support), {0, 1, 3}) + self.assertIn(int(sampled[0]), support) + expected = original[0, sampled[0]] - original[0, support].logsumexp(0) + self.assertAlmostEqual( + output.next_token_sampling_logprobs[0], expected.item(), places=5 + ) + self.assertIn(2, output.next_token_sampling_mask_idx[1]) + @unittest.skipIf(is_hip(), "FlashInfer is not available on ROCm") def test_flashinfer_joint_cutoff_ties_match_capture(self): batch_size = 256 @@ -232,12 +296,15 @@ class SamplingMaskTestMixin: return_sampling_mask=True, return_logprob=False, top_logprobs_num=0, + custom_logit_processor=None, ): payload = { "text": "The capital of France is", "sampling_params": sampling_params, "return_sampling_mask": return_sampling_mask, } + if custom_logit_processor is not None: + payload["custom_logit_processor"] = custom_logit_processor if return_logprob: payload["return_logprob"] = True payload["top_logprobs_num"] = top_logprobs_num @@ -276,6 +343,55 @@ class TestSamplingMask(SamplingMaskTestMixin, CustomTestCase): def setUpClass(cls): cls._launch_server() + def test_disallowed_tokens_with_replay(self): + params = { + "temperature": 1.0, + "top_k": _TOP_K, + "top_p": _TOP_P, + "max_new_tokens": 1, + "ignore_eos": True, + } + baseline = self._post_generate(params) + self.assertEqual(baseline.status_code, 200, baseline.text) + # Exclude tokens that actually belong to the unmodified sampling support. + blocked = baseline.json()["meta_info"]["output_token_sampling_mask"][0][:2] + self.assertTrue(blocked) + response = self._post_generate( + {**params, "custom_params": {"token_ids": blocked}}, + return_logprob=True, + top_logprobs_num=_TOP_LOGPROBS_NUM, + custom_logit_processor=DisallowedTokensLogitsProcessor.to_str(), + ) + self.assertEqual(response.status_code, 200, response.text) + output = response.json() + meta = output["meta_info"] + token = output["output_ids"][0] + mask = meta["output_token_sampling_mask"][0] + self.assertTrue(set(mask).isdisjoint(blocked)) + self.assertIn(token, mask) + probs = { + int(tid): math.exp(lp) for lp, tid, _ in meta["output_top_logprobs"][0] + } + expected = math.log(probs[token] / sum(probs[tid] for tid in mask)) + self.assertAlmostEqual( + meta["output_token_sampling_logprobs"][0], expected, delta=1e-2 + ) + + def test_rejected_processors_do_not_break_generation(self): + params = {"top_k": _TOP_K, "max_new_tokens": 1} + for processor in ( + Qwen3ThinkingBudgetLogitProcessor.to_str(), + "invalid processor", + ): + with self.subTest(processor=processor): + response = self._post_generate(params, custom_logit_processor=processor) + self.assertEqual(response.status_code, 400, response.text) + self.assertIn( + "only supports DisallowedTokensLogitsProcessor", response.text + ) + recovery = self._post_generate(params) + self.assertEqual(recovery.status_code, 200, recovery.text) + def test_generate_returns_sampling_mask(self): top_p_sampling_masks = self._generate_sampling_masks( { diff --git a/test/registered/unit/managers/test_sampling_mask_validation.py b/test/registered/unit/managers/test_sampling_mask_validation.py new file mode 100644 index 000000000..a41dfe548 --- /dev/null +++ b/test/registered/unit/managers/test_sampling_mask_validation.py @@ -0,0 +1,83 @@ +import unittest +from unittest.mock import patch + +from sglang.srt.managers.io_struct import GenerateReqInput +from sglang.srt.managers.tokenizer_manager import TokenizerManager +from sglang.srt.runtime_context import get_context +from sglang.srt.sampling.custom_logit_processor import ( + CustomLogitProcessor, + DisallowedTokensLogitsProcessor, + Qwen3ThinkingBudgetLogitProcessor, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + + +class TestSamplingMaskValidation(CustomTestCase): + def setUp(self): + override = get_context().override_server_args( + enable_custom_logit_processor=True + ) + override.install() + self.addCleanup(override.restore) + self.manager = TokenizerManager.__new__(TokenizerManager) + self.manager.context_len = 128 + self.manager.num_reserved_tokens = 0 + self.manager.allow_auto_truncate = False + self.manager.validate_total_tokens = False + self.manager.is_generation = True + + def _validate(self, processor, return_sampling_mask=True): + req = GenerateReqInput( + input_ids=[1, 2, 3], + sampling_params={"top_k": 10}, + custom_logit_processor=processor, + return_sampling_mask=return_sampling_mask, + ) + self.manager._validate_one_request(req, req.input_ids) + + def test_accepts_hard_exclusion(self): + class BoundMask(DisallowedTokensLogitsProcessor): + def __call__(self, logits, custom_param_list=None): + logits[..., [2]] = -float("inf") + return logits + + for processor in ( + None, + DisallowedTokensLogitsProcessor.to_str(), + BoundMask.to_str(), + ): + with self.subTest(processor=processor): + self._validate(processor) + + def test_rejects_unsupported_or_malformed_processors(self): + for processor in ( + Qwen3ThinkingBudgetLogitProcessor.to_str(), + "invalid processor", + '{"callable": "00"}', + ): + with self.subTest(processor=processor): + with self.assertRaisesRegex( + ValueError, "only supports DisallowedTokensLogitsProcessor" + ): + self._validate(processor) + + def test_disabled_processors_are_not_deserialized(self): + with ( + get_context().override_server_args(enable_custom_logit_processor=False), + patch.object(CustomLogitProcessor, "from_str") as deserialize, + ): + with self.assertRaisesRegex(ValueError, "--enable-custom-logit-processor"): + self._validate(DisallowedTokensLogitsProcessor.to_str()) + deserialize.assert_not_called() + + def test_other_processors_still_work_without_sampling_masks(self): + self._validate( + Qwen3ThinkingBudgetLogitProcessor.to_str(), return_sampling_mask=False + ) + + +if __name__ == "__main__": + unittest.main()