From 6b46a9d7006260eeac6c6262d1bca90535d9f7d9 Mon Sep 17 00:00:00 2001 From: Vansh Verma Date: Fri, 17 Jul 2026 02:31:04 -0500 Subject: [PATCH] Validate stop_token_ids range in SamplingParams.verify User-supplied stop_token_ids reach the min_new_tokens penalizer's scatter_add_ over a [vocab_size + 1] tensor with no range check (SamplingParams.verify validates logit_bias keys but not stop_token_ids). An id >= vocab_size + 1 is an out-of-bounds scatter index: on CPU a RuntimeError, on CUDA a device-side assert that poisons the context and takes down the whole server process (every in-flight request) from a single malformed request. An id == vocab_size does not crash but is silently dropped by the trailing [:, :vocab_size] slice, so that stop token never fires. Reject out-of-range stop_token_ids in verify(), matching the existing logit_bias validation, so a bad request fails cleanly at the boundary. --- python/sglang/srt/sampling/sampling_params.py | 12 +++++++ .../unit/sampling/test_sampling_params.py | 31 +++++++++++++++++++ 2 files changed, 43 insertions(+) diff --git a/python/sglang/srt/sampling/sampling_params.py b/python/sglang/srt/sampling/sampling_params.py index b30af33bb55e..b42f84bec89c 100644 --- a/python/sglang/srt/sampling/sampling_params.py +++ b/python/sglang/srt/sampling/sampling_params.py @@ -221,6 +221,18 @@ def verify(self, vocab_size): f"logit_bias must has keys in [0, {vocab_size - 1}], got " f"{token_id}." ) + if self.stop_token_ids: + # Unchecked out-of-range ids reach the min_new_tokens penalizer's + # scatter_add_ over a [vocab_size + 1] tensor; an id >= vocab_size + 1 + # is an out-of-bounds index (a CUDA device-side assert that takes down + # the whole server), and id == vocab_size is silently dropped by the + # trailing [:, :vocab_size] slice. Validate like logit_bias keys. + for token_id in self.stop_token_ids: + if not 0 <= token_id < vocab_size: + raise ValueError( + f"stop_token_ids must be in [0, {vocab_size - 1}], got " + f"{token_id}." + ) grammars = [ self.json_schema, diff --git a/test/registered/unit/sampling/test_sampling_params.py b/test/registered/unit/sampling/test_sampling_params.py index f1985f6f470a..ebbae03ae54c 100644 --- a/test/registered/unit/sampling/test_sampling_params.py +++ b/test/registered/unit/sampling/test_sampling_params.py @@ -272,6 +272,37 @@ def test_logit_bias_valid_tokens(self): sp = self._make(logit_bias={"0": 1.0, "31999": -0.5}) sp.verify(self.VOCAB_SIZE) + # --- stop_token_ids --- + def test_stop_token_ids_exceeds_vocab_raises(self): + """verify() must reject an out-of-range stop_token_id. + + An id >= vocab_size + 1 reaches the min_new_tokens penalizer's + scatter_add_ over a [vocab_size + 1] tensor as an out-of-bounds index + (a CUDA device-side assert that crashes the whole server from one + request), so it must be rejected at validation like logit_bias keys. + """ + sp = self._make(stop_token_ids=[self.VOCAB_SIZE + 1000]) + with self.assertRaises(ValueError): + sp.verify(self.VOCAB_SIZE) + + def test_stop_token_ids_equals_vocab_raises(self): + """id == vocab_size is out of range; unchecked it is silently dropped by + the penalizer's trailing [:, :vocab_size] slice (stop token never fires).""" + sp = self._make(stop_token_ids=[self.VOCAB_SIZE]) + with self.assertRaises(ValueError): + sp.verify(self.VOCAB_SIZE) + + def test_stop_token_ids_negative_raises(self): + """verify() must reject a negative stop_token_id.""" + sp = self._make(stop_token_ids=[-1]) + with self.assertRaises(ValueError): + sp.verify(self.VOCAB_SIZE) + + def test_stop_token_ids_valid_tokens(self): + """stop_token_ids within [0, vocab_size) are accepted.""" + sp = self._make(stop_token_ids=[0, self.VOCAB_SIZE - 1]) + sp.verify(self.VOCAB_SIZE) + def test_multiple_grammars_raises(self): """Test that verify() rejects setting both json_schema and regex (mutually exclusive).""" sp = self._make(json_schema='{"type":"object"}', regex="abc")