[Sampling] Allow sampling-mask replay with DisallowedTokensLogitsProcessor (#38279)

Co-authored-by: Byron Hsu <24364830+ByronHsu@users.noreply.github.com>

python/sglang/srt/managers/tokenizer_manager.py

@@ -143,6 +143,7 @@ from sglang.srt.runtime_context import (
     get_serving,
     get_spec,
 )
+from sglang.srt.sampling.custom_logit_processor import supports_sampling_mask
 from sglang.srt.sampling.sampling_params import SamplingParams
 from sglang.srt.server_args import (
     PortArgs,
@@ -1259,6 +1260,17 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
                     "The server is not configured to enable custom logit processor. "
                     "Please set `--enable-custom-logit-processor` to enable this feature."
                 )
+            if (
+                obj.return_sampling_mask
+                and obj.custom_logit_processor
+                and not supports_sampling_mask(obj.custom_logit_processor)
+            ):
+                # Reject before scheduling so aborted requests cannot execute
+                # unsupported processors during sampling batch preparation.
+                raise ValueError(
+                    "return_sampling_mask only supports DisallowedTokensLogitsProcessor "
+                    "among custom logit processors."
+                )
 
     def _validate_mm_limits(
         self, obj: Union[GenerateReqInput, EmbeddingReqInput]

python/sglang/srt/sampling/custom_logit_processor.py

@@ -58,6 +58,17 @@ class DisallowedTokensLogitsProcessor(CustomLogitProcessor):
         return logits
 
 
+def supports_sampling_mask(serialized_processor: str) -> bool:
+    """Hard exclusion preserves the relative logits needed for mask-based replay."""
+    try:
+        return isinstance(
+            CustomLogitProcessor.from_str(serialized_processor),
+            DisallowedTokensLogitsProcessor,
+        )
+    except Exception:
+        return False
+
+
 def _open_thinking_start(ids: list[int], start_id: int, end_id: int) -> int:
     """Return the index of the start token of the currently open thinking block, or -1."""
     for idx in reversed(range(len(ids))):