[Sampling] Restore finite top-k requirement for sampling masks (#35205)

2026-08-20 16:55:25 -07:00
parent 14795dcb1a
commit f825d72936
2 changed files with 25 additions and 25 deletions
@@ -2572,13 +2572,11 @@ class Scheduler(
            self._add_request_to_queue(req)
            return

-        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:
+        if req.return_sampling_mask and req.sampling_params.top_k == TOP_K_ALL:
            error_msg = (
-                "return_sampling_mask cannot return the full vocabulary; set "
-                "top_p < 1 or a finite top_k."
+                "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."
            )
            req.set_finish_with_abort(error_msg)
            self.init_req_max_new_tokens(req)