[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 "
-                "is valid but can return huge masks in the tail, blowing up "
-                "metadata, so we need a safety cap."
+                "return_sampling_mask cannot return the full vocabulary; set "
+                "top_p < 1 or a finite top_k."
             )
             req.set_finish_with_abort(error_msg)
             self.init_req_max_new_tokens(req)