diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index aca197603..ac32f696d 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -985,8 +985,9 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): if isinstance(obj, EmbeddingReqInput): self._validate_for_matryoshka_dim(obj) - # Validate custom logit processor + # Validate generation-specific fields if isinstance(obj, GenerateReqInput): + self._validate_token_ids_logprob(obj) if ( obj.return_hidden_states and not self.server_args.enable_return_hidden_states @@ -1047,6 +1048,26 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): f"Provided dimensions are greater than max embedding dimension: {self.model_config.hidden_size}" ) + def _validate_token_ids_logprob(self, obj: GenerateReqInput) -> None: + # Batch requests are split into per-request sub-objects before this + # runs (normalize_batch_and_arguments + __getitem__), so the only + # legal shape here is the per-request contract of + # TokenizedGenerateReqInput.token_ids_logprob: a flat list of ints. + token_ids_logprob = obj.token_ids_logprob + if not token_ids_logprob: + return + if not isinstance(token_ids_logprob, list): + raise ValueError("token_ids_logprob must be a flat list of integers.") + vocab_size = self.model_config.vocab_size + for token_id in token_ids_logprob: + if not isinstance(token_id, int): + raise ValueError("token_ids_logprob must be a flat list of integers.") + if token_id < 0 or token_id >= vocab_size: + raise ValueError( + f"token_ids_logprob contains out-of-vocabulary token id " + f"{token_id}; valid range is [0, {vocab_size})." + ) + def _validate_input_ids_in_vocab( self, input_ids: Union[List[int], List[List[int]]], vocab_size: int ) -> None: diff --git a/test/registered/openai_server/validation/test_request_length_validation.py b/test/registered/openai_server/validation/test_request_length_validation.py index 1ee3ae9ad..3236a6630 100644 --- a/test/registered/openai_server/validation/test_request_length_validation.py +++ b/test/registered/openai_server/validation/test_request_length_validation.py @@ -1,6 +1,7 @@ import unittest import openai +import requests from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci @@ -104,6 +105,70 @@ class TestRequestLengthValidation(CustomTestCase): str(cm.exception), ) + def test_token_ids_logprob_out_of_vocabulary(self): + headers = {"Authorization": f"Bearer {self.api_key}"} + for token_ids_logprob in ([-1], [2_000_000_000]): + response = requests.post( + f"{self.base_url}/generate", + headers=headers, + json={ + "text": "hi", + "sampling_params": {"max_new_tokens": 1}, + "return_logprob": True, + "token_ids_logprob": token_ids_logprob, + }, + ) + self.assertEqual(response.status_code, 400) + self.assertIn("out-of-vocabulary", response.text) + + def test_token_ids_logprob_rejects_nested_list(self): + # Nested lists are a batch-level wire format; a single request must + # pass a flat list of ints. A ragged nested list with in-vocab ids + # would otherwise crash the scheduler in the sampler gather. + headers = {"Authorization": f"Bearer {self.api_key}"} + for token_ids_logprob in ([[0]], [[0], [1, 2]]): + response = requests.post( + f"{self.base_url}/generate", + headers=headers, + json={ + "text": "hi", + "sampling_params": {"max_new_tokens": 1}, + "return_logprob": True, + "token_ids_logprob": token_ids_logprob, + }, + ) + self.assertEqual(response.status_code, 400) + self.assertIn("flat list of integers", response.text) + + def test_token_ids_logprob_batch_with_one_oov(self): + headers = {"Authorization": f"Bearer {self.api_key}"} + response = requests.post( + f"{self.base_url}/generate", + headers=headers, + json={ + "text": ["hi", "hi"], + "sampling_params": {"max_new_tokens": 1}, + "return_logprob": True, + "token_ids_logprob": [[0], [2_000_000_000]], + }, + ) + self.assertEqual(response.status_code, 400) + self.assertIn("out-of-vocabulary", response.text) + + def test_token_ids_logprob_valid(self): + headers = {"Authorization": f"Bearer {self.api_key}"} + response = requests.post( + f"{self.base_url}/generate", + headers=headers, + json={ + "text": "hi", + "sampling_params": {"max_new_tokens": 1}, + "return_logprob": True, + "token_ids_logprob": [0], + }, + ) + self.assertEqual(response.status_code, 200) + if __name__ == "__main__": unittest.main()