[RL] Expose top-p-only sampling masks (#33593)

python/sglang/srt/entrypoints/openai/protocol.py

@@ -795,6 +795,7 @@ class ChatCompletionRequest(BaseModel):
     return_prompt_token_ids: bool = False
     return_token_ids: bool = False
     return_meta_info: bool = False
+    return_sampling_mask: bool = False
     reasoning_effort: ReasoningEffortType = Field(
         default=None,
         description="Constrains effort on reasoning for reasoning models. "

python/sglang/srt/entrypoints/openai/serving_chat.py

@@ -816,6 +816,9 @@ class OpenAIServingChat(OpenAIServingBase):
         if not request.messages:
             return "Messages cannot be empty."
 
+        if request.return_sampling_mask and not request.return_meta_info:
+            return "return_sampling_mask requires return_meta_info=true."
+
         media_error = self._validate_media_content(request)
         if media_error:
             return media_error
@@ -1004,6 +1007,7 @@ class OpenAIServingChat(OpenAIServingBase):
             return_logprob=request.logprobs,
             logprob_start_len=-1,
             top_logprobs_num=request.top_logprobs or 0,
+            return_sampling_mask=request.return_sampling_mask,
             stream=request.stream,
             return_text_in_logprobs=True,
             modalities=processed_messages.modalities,

python/sglang/srt/managers/scheduler.py

@@ -2528,11 +2528,13 @@ class Scheduler(
             self._add_request_to_queue(req)
             return
 
-        if req.return_sampling_mask and req.sampling_params.top_k == TOP_K_ALL:
+        uses_top_k_or_top_p_truncation = (
+            req.sampling_params.top_k != TOP_K_ALL or req.sampling_params.top_p < 1.0
+        )
+        if req.return_sampling_mask and not uses_top_k_or_top_p_truncation:
             error_msg = (
-                "return_sampling_mask requires finite top_k; top_p-only sampling "
+                "return_sampling_mask cannot return the full vocabulary; set "
-                "is valid but can return huge masks in the tail, blowing up "
+                "top_p < 1 or a finite top_k."
-                "metadata, so we need a safety cap."
             )
             req.set_finish_with_abort(error_msg)
             self.init_req_max_new_tokens(req)

test/registered/sampling/test_sampling_mask.py

@@ -18,15 +18,13 @@ register_amd_ci(est_time=320, suite="stage-b-test-1-gpu-small-amd")
 
 _MAX_NEW_TOKENS = 4
 _TOP_P = 0.99
+_TOP_P_SMALL = 1e-5
 _TOP_K = 10
 _SAMPLING_SEED = 1234
 _SERVER_ARGS = (
     "--mem-fraction-static",
     "0.7",
 )
-_INVALID_SAMPLING_MASK_ERROR = (
-    "top_p-only sampling is valid but can return huge masks in the tail"
-)
 
 
 class SamplingMaskTestMixin:
@@ -81,11 +79,6 @@ class SamplingMaskTestMixin:
             self.assertIn(output_id, sampling_mask)
         return sampling_masks
 
-    def _assert_rejects_unbounded_sampling_mask(self, sampling_params):
-        response = self._post_generate(sampling_params)
-        self.assertEqual(response.status_code, 400, response.text)
-        self.assertIn(_INVALID_SAMPLING_MASK_ERROR, response.text)
-
 
 class TestSamplingMask(SamplingMaskTestMixin, CustomTestCase):
     @classmethod
@@ -191,16 +184,48 @@ class TestSamplingMask(SamplingMaskTestMixin, CustomTestCase):
             expected_logprob = math.log(probs[output_id] / support_mass)
             self.assertAlmostEqual(mask_logprob, expected_logprob, delta=1e-2)
 
-    def test_generate_rejects_unbounded_sampling_mask(self):
+    def test_generate_returns_top_p_only_sampling_mask(self):
-        self._assert_rejects_unbounded_sampling_mask(
+        self._generate_sampling_masks(
             {
                 "temperature": 1.0,
-                "top_p": _TOP_P,
+                "top_p": _TOP_P_SMALL,
                 "max_new_tokens": _MAX_NEW_TOKENS,
                 "ignore_eos": True,
             }
         )
-        self._assert_rejects_unbounded_sampling_mask(
+
+    def test_chat_completions_returns_top_p_only_sampling_mask(self):
+        response = requests.post(
+            self.base_url + "/v1/chat/completions",
+            json={
+                "model": self.model,
+                "messages": [{"role": "user", "content": "Name a capital city."}],
+                "temperature": 1.0,
+                "top_p": _TOP_P_SMALL,
+                "max_tokens": _MAX_NEW_TOKENS,
+                "ignore_eos": True,
+                "return_sampling_mask": True,
+                "return_meta_info": True,
+                "return_token_ids": True,
+            },
+            timeout=60,
+        )
+        self.assertEqual(response.status_code, 200, response.text)
+
+        choice = response.json()["choices"][0]
+        output_ids = choice["token_ids"]
+        meta_info = choice["meta_info"]
+        sampling_masks = meta_info["output_token_sampling_mask"]
+        sampling_logprobs = meta_info["output_token_sampling_logprobs"]
+
+        self.assertEqual(len(output_ids), _MAX_NEW_TOKENS)
+        self.assertEqual(len(sampling_masks), len(output_ids))
+        self.assertEqual(len(sampling_logprobs), len(output_ids))
+        for output_id, sampling_mask in zip(output_ids, sampling_masks):
+            self.assertIn(output_id, sampling_mask)
+
+    def test_generate_rejects_full_vocabulary_sampling_mask(self):
+        response = self._post_generate(
             {
                 "temperature": 1.0,
                 "top_p": 1.0,
@@ -208,6 +233,10 @@ class TestSamplingMask(SamplingMaskTestMixin, CustomTestCase):
                 "ignore_eos": True,
             }
         )
+        self.assertEqual(response.status_code, 400, response.text)
+        self.assertIn(
+            "return_sampling_mask cannot return the full vocabulary", response.text
+        )
 
 
 class TestSamplingMaskDeterministic(SamplingMaskTestMixin, CustomTestCase):

test/registered/unit/entrypoints/openai/test_protocol.py

@@ -124,6 +124,7 @@ class TestChatCompletionRequest(unittest.TestCase):
         self.assertEqual(request.messages[0].content, "Hello")
         self.assertEqual(request.temperature, None)  # default
         self.assertFalse(request.stream)  # default
+        self.assertFalse(request.return_sampling_mask)
         self.assertEqual(request.tool_choice, "none")  # default when no tools
 
     def test_image_content_hash_validation(self):

test/registered/unit/entrypoints/openai/test_serving_chat.py

@@ -278,9 +278,12 @@ class ServingChatTestCase(unittest.TestCase):
                 None,
             )
 
+            self.basic_req.return_sampling_mask = True
+            self.basic_req.return_meta_info = True
             adapted, processed = self.chat._convert_to_internal_request(self.basic_req)
             self.assertIsInstance(adapted, GenerateReqInput)
             self.assertFalse(adapted.stream)
+            self.assertTrue(adapted.return_sampling_mask)
             self.assertEqual(adapted.session_id, "session-1")
             self.assertEqual(processed, self.basic_req)
 
@@ -295,6 +298,18 @@ class ServingChatTestCase(unittest.TestCase):
             with self.subTest(field=field), self.assertRaisesRegex(ValueError, field):
                 self.chat._convert_to_internal_request(req, self.fastapi_request)
 
+    def test_validate_request_rejects_sampling_mask_without_meta_info(self):
+        req = ChatCompletionRequest(
+            model="x",
+            messages=[{"role": "user", "content": "Hi?"}],
+            return_sampling_mask=True,
+        )
+
+        self.assertEqual(
+            self.chat._validate_request(req),
+            "return_sampling_mask requires return_meta_info=true.",
+        )
+
     def test_convert_to_internal_request_rejects_stream_return_meta_info(self):
         req = ChatCompletionRequest(
             model="x",