feat(constrained): two-phase reasoning grammar + --enable-strict-thinking (#23953)

python/sglang/srt/constrained/base_grammar_backend.py

@@ -21,6 +21,7 @@ from typing import Dict, List, Optional, Tuple
 
 import torch
 
+from sglang.srt.parser.reasoning_parser import ReasoningParser
 from sglang.srt.server_args import ServerArgs
 
 logger = logging.getLogger(__name__)
@@ -128,6 +129,8 @@ class InvalidGrammarObject(BaseGrammarObject):
 
 
 class BaseGrammarBackend:
+    _enable_strict_thinking: bool = False
+
     def __init__(self):
         self.executor = ThreadPoolExecutor()
         self.cache: Dict[Tuple[str, str], BaseGrammarObject] = {}
@@ -136,6 +139,24 @@ class BaseGrammarBackend:
         logger.warning(f"Skip unsupported {key_type=}, {key_string=}")
         return InvalidGrammarObject()
 
+    @property
+    def enable_strict_thinking(self):
+        return self._enable_strict_thinking
+
+    @property
+    def is_support_token_filter(self):
+        return False
+
+    def set_token_filter(
+        self, vocab_mask, token_ids, batch_idx, is_allowed=True, reset_vocab_mask=True
+    ):
+        """Set or clear specific tokens in the vocab mask. No-op by default."""
+        pass
+
+    def init_strict_reasoning_grammar(self, reasoning: bool):
+        """Create a grammar object for strict token filtering only. Returns None by default."""
+        return None
+
     def dispatch_fallback(self, key_type: str, key_string: str) -> BaseGrammarObject:
         """
         This function should not be reached in any case.
@@ -239,6 +260,13 @@ def create_grammar_backend(
                 any_whitespace=not server_args.constrained_json_disable_any_whitespace,
             )
         except TokenizerNotSupportedError as e:
+            if server_args.enable_strict_thinking:
+                raise ValueError(
+                    f"--enable-strict-thinking requires a grammar backend with "
+                    f"token filtering support, but XGrammar failed to initialize: "
+                    f"{e}. Cannot fall back to grammar_backend='none' with strict "
+                    f"thinking enabled."
+                ) from e
             logger.warning(
                 f"Grammar backend disabled because tokenizer is not supported by XGrammar: {e}. "
                 "Falling back to grammar_backend='none'. "
@@ -255,6 +283,13 @@ def create_grammar_backend(
             whitespace_pattern=server_args.constrained_json_whitespace_pattern,
         )
     elif name == "none":
+        if server_args.enable_strict_thinking:
+            raise ValueError(
+                "--enable-strict-thinking requires a grammar backend that supports "
+                "token filtering, but grammar_backend='none' was specified. Use "
+                "--grammar-backend xgrammar or another backend that supports token "
+                "filtering."
+            )
         return None
     else:
         raise ValueError(f"Invalid grammar backend: {name}")
@@ -264,6 +299,15 @@ def create_grammar_backend(
             ReasonerGrammarBackend,
         )
 
-        grammar_backend = ReasonerGrammarBackend(grammar_backend, think_end_id)
+        reasoning_parser = ReasoningParser(
+            model_type=server_args.reasoning_parser, stream_reasoning=False
+        )
+
+        grammar_backend = ReasonerGrammarBackend(
+            grammar_backend,
+            reasoning_parser,
+            tokenizer,
+            enable_strict_thinking=server_args.enable_strict_thinking,
+        )
 
     return grammar_backend

python/sglang/srt/constrained/grammar_manager.py

@@ -11,6 +11,7 @@ from sglang.srt.constrained.base_grammar_backend import (
     InvalidGrammarObject,
     create_grammar_backend,
 )
+from sglang.srt.constrained.reasoner_grammar_backend import ReasonerGrammarObject
 from sglang.srt.environ import envs
 
 if TYPE_CHECKING:
@@ -37,6 +38,12 @@ class GrammarManager:
         else:
             self.grammar_backend = None
 
+        self._enable_strict_thinking = (
+            self.grammar_backend.enable_strict_thinking
+            if self.grammar_backend is not None
+            else False
+        )
+
         self.grammar_sync_group = scheduler.dp_tp_cpu_group
         self.grammar_sync_size = scheduler.dp_tp_group.world_size
         self.grammar_sync_entry = scheduler.dp_tp_group.first_rank
@@ -65,6 +72,20 @@ class GrammarManager:
                     req.grammar.cancel()
                 req.set_finish_with_abort("Aborted by AbortReq.")
 
+    def _get_request_thinking_budget(self, req: Req) -> int | None:
+        custom_params = req.sampling_params.custom_params
+        if not isinstance(custom_params, dict):
+            return None
+        thinking_budget = custom_params.get("thinking_budget")
+        return thinking_budget if isinstance(thinking_budget, int) else None
+
+    def _apply_request_reasoning_budget(self, req: Req) -> None:
+        thinking_budget = self._get_request_thinking_budget(req)
+        if thinking_budget is None:
+            return
+        if isinstance(req.grammar, ReasonerGrammarObject):
+            req.grammar.max_think_tokens = thinking_budget
+
     def process_req_with_grammar(self, req: Req) -> bool:
         # Init grammar cache for this request
         add_to_grammar_queue = False
@@ -103,6 +124,15 @@ class GrammarManager:
                             f"Failed to compile {key[0]} grammar: {value.error_message}"
                         )
                         req.set_finish_with_abort(error_msg)
+                    else:
+                        self._apply_request_reasoning_budget(req)
+        elif self._enable_strict_thinking:
+            grammar_obj = self.grammar_backend.init_strict_reasoning_grammar(
+                req.require_reasoning
+            )
+            if grammar_obj is not None:
+                req.grammar = grammar_obj
+                self._apply_request_reasoning_budget(req)
 
         if add_to_grammar_queue:
             self.grammar_queue.append(req)
@@ -177,8 +207,16 @@ class GrammarManager:
                 continue
 
             assert isinstance(req.grammar, futures.Future) and req.grammar_key
-            req.grammar = req.grammar.result()
+            try:
+                req.grammar = req.grammar.result()
+            except Exception as e:
+                logger.error(
+                    f"Grammar compilation raised an exception: {e}, "
+                    f"grammar_key={req.grammar_key}"
+                )
+                req.grammar = InvalidGrammarObject(f"Grammar compilation failed: {e}")
             self.grammar_backend.set_cache(req.grammar_key, req.grammar.copy())
+            self._apply_request_reasoning_budget(req)
             if isinstance(req.grammar, InvalidGrammarObject):
                 error_msg = f"Failed to compile {req.grammar_key[0]} grammar: {req.grammar.error_message}"
                 req.set_finish_with_abort(error_msg)

python/sglang/srt/constrained/reasoner_grammar_backend.py

@@ -13,9 +13,14 @@
 # ==============================================================================
 """The baseclass of a backend for reasoner grammar-guided constrained decoding."""
 
-from typing import List, Optional, Tuple
+import logging
+from typing import List, Optional, Tuple, Union
 
 import torch
+from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
+
+from sglang.srt.environ import envs
+from sglang.srt.parser.reasoning_parser import ReasoningParser
 
 from .base_grammar_backend import (
     BaseGrammarBackend,
@@ -23,102 +28,291 @@ from .base_grammar_backend import (
     InvalidGrammarObject,
 )
 
+logger = logging.getLogger(__name__)
+
 
 class ReasonerGrammarObject(BaseGrammarObject):
-    def __init__(self, grammar: BaseGrammarObject, think_end_id: int):
+    """Wraps a grammar object to handle reasoning (think/generation) phases.
+
+    State machine (must call maybe_init_reasoning before use):
+      THINKING (tokens_in_think >= 0, tokens_after_end == -1)
+        -> grammar not consulted, optional token filtering
+      GENERATION (tokens_after_end >= 0)
+        -> grammar consulted for accept/fill/rollback
+
+    When enable_token_filter=True (strict mode), fill_vocab_mask filters
+    excluded tokens during THINKING and enforces max_think_tokens budget.
+    When the budget is exhausted, only think_end_id is allowed, forcing the
+    model to exit the thinking phase.
+    When enable_token_filter=False (non-strict mode), fill_vocab_mask is
+    a no-op during THINKING.
+    """
+
+    def __init__(
+        self,
+        grammar: Optional[BaseGrammarObject],
+        think_end_id: int,
+        think_excluded_token_ids: Optional[List[int]] = None,
+        max_think_tokens: int = -1,
+        enable_token_filter: bool = False,
+        token_filter_fn=None,
+        allocate_vocab_mask_fn=None,
+        move_vocab_mask_fn=None,
+        apply_vocab_mask_fn=None,
+    ):
         super().__init__()
         self.grammar = grammar
         self.think_end_id = think_end_id
-        # -1    means thinking has not ended yet
-        # 0     means just ended thinking in the last token
-        # +     means number of tokens after thinking ended
-        self.tokens_after_think_end = -1
+        self.think_excluded_token_ids = think_excluded_token_ids
+        self.max_think_tokens = max_think_tokens
+        self.enable_token_filter = enable_token_filter
+        self.token_filter_fn = token_filter_fn
+        self.allocate_vocab_mask_fn = allocate_vocab_mask_fn
+        self.move_vocab_mask_fn = move_vocab_mask_fn
+        self.apply_vocab_mask_fn = apply_vocab_mask_fn
+        self._think_end_id_list = [think_end_id]
+
+        self.tokens_in_think = -1
+        self.tokens_after_end = -1
 
     def maybe_init_reasoning(self, reasoning: bool):
-        self.tokens_after_think_end = -1 if reasoning else 0
+        if reasoning:
+            self.tokens_in_think = 0
+        else:
+            self.tokens_in_think = -1
+            self.tokens_after_end = 0
 
-    def transfer_state(self, token: int) -> int:
-        if self.tokens_after_think_end == -1 and token == self.think_end_id:
-            self.tokens_after_think_end = 0
-        elif self.tokens_after_think_end >= 0:
-            self.tokens_after_think_end += 1
+    def _is_thinking(self):
+        return self.tokens_in_think >= 0 and self.tokens_after_end == -1
+
+    def _is_generation(self):
+        return self.tokens_after_end >= 0
+
+    def transfer_state(self, token: int) -> None:
+        if self._is_thinking():
+            if token == self.think_end_id:
+                self.tokens_after_end = 0
+            else:
+                self.tokens_in_think += 1
+        elif self._is_generation():
+            self.tokens_after_end += 1
 
     def rollback_state(self):
-        if self.tokens_after_think_end == 0:
-            self.tokens_after_think_end = -1
-        elif self.tokens_after_think_end > 0:
-            self.tokens_after_think_end -= 1
+        if self._is_thinking():
+            if self.tokens_in_think > 0:
+                self.tokens_in_think -= 1
+        elif self._is_generation():
+            if self.tokens_after_end == 0:
+                self.tokens_after_end = -1
+            elif self.tokens_after_end > 0:
+                self.tokens_after_end -= 1
 
     def accept_token(self, token: int):
-        if self.tokens_after_think_end >= 0:
+        if self._is_generation() and self.grammar is not None:
             self.grammar.accept_token(token)
         self.transfer_state(token)
 
     def is_terminated(self):
-        return self.grammar.is_terminated()
+        if self.grammar is not None:
+            return self.grammar.is_terminated()
+        return False
 
     def rollback(self, k):
-        steps_after_think = min(k, self.tokens_after_think_end)
-        if steps_after_think > 0:
-            self.grammar.rollback(steps_after_think)
-
+        if self.grammar is not None:
+            steps_after = min(k, max(0, self.tokens_after_end))
+            if steps_after > 0:
+                self.grammar.rollback(steps_after)
         for _ in range(k):
             self.rollback_state()
 
-    def allocate_vocab_mask(
-        self, vocab_size: int, batch_size: int, device
-    ) -> torch.Tensor:
-        return self.grammar.allocate_vocab_mask(vocab_size, batch_size, device)
+    def _can_think_more(self):
+        return self.max_think_tokens < 0 or self.tokens_in_think < self.max_think_tokens
+
+    def _do_token_filter(self, vocab_mask, token_ids, idx, is_allowed=True):
+        if self.token_filter_fn is not None:
+            self.token_filter_fn(vocab_mask, token_ids, idx, is_allowed)
 
     def fill_vocab_mask(self, vocab_mask: torch.Tensor, idx: int) -> None:
-        if self.tokens_after_think_end >= 0:
+        if self._is_thinking():
+            if not self.enable_token_filter:
+                return
+            if self._can_think_more():
+                self._do_token_filter(
+                    vocab_mask, self.think_excluded_token_ids, idx, is_allowed=False
+                )
+            else:
+                self._do_token_filter(
+                    vocab_mask, self._think_end_id_list, idx, is_allowed=True
+                )
+            return
+        if self._is_generation() and self.grammar is not None:
             self.grammar.fill_vocab_mask(vocab_mask, idx)
 
-    def move_vocab_mask(self, vocab_mask: torch.Tensor, device) -> torch.Tensor:
-        return self.grammar.move_vocab_mask(vocab_mask, device)
+    def allocate_vocab_mask(self, vocab_size, batch_size, device):
+        if self.grammar is not None:
+            return self.grammar.allocate_vocab_mask(vocab_size, batch_size, device)
+        if self.allocate_vocab_mask_fn is not None:
+            return self.allocate_vocab_mask_fn(vocab_size, batch_size, device)
+        return None
+
+    def move_vocab_mask(self, vocab_mask, device):
+        if self.grammar is not None:
+            return self.grammar.move_vocab_mask(vocab_mask, device)
+        if self.move_vocab_mask_fn is not None:
+            return self.move_vocab_mask_fn(vocab_mask, device)
+        return vocab_mask
 
     @property
     def apply_vocab_mask(self):
-        return self.grammar.apply_vocab_mask
+        if self.grammar is not None:
+            return self.grammar.apply_vocab_mask
+        return self.apply_vocab_mask_fn
 
-    def copy(self) -> BaseGrammarObject:
-        return ReasonerGrammarObject(self.grammar.copy(), self.think_end_id)
+    def copy(self):
+        new_obj = ReasonerGrammarObject(
+            self.grammar.copy() if self.grammar is not None else None,
+            self.think_end_id,
+            self.think_excluded_token_ids,
+            self.max_think_tokens,
+            self.enable_token_filter,
+            self.token_filter_fn,
+            self.allocate_vocab_mask_fn,
+            self.move_vocab_mask_fn,
+            self.apply_vocab_mask_fn,
+        )
+        new_obj.tokens_in_think = self.tokens_in_think
+        new_obj.tokens_after_end = self.tokens_after_end
+        new_obj._finished = self._finished
+        return new_obj
 
     @property
     def finished(self):
-        return self.grammar.finished
+        if self.grammar is not None:
+            return self.grammar.finished
+        return self._finished
 
     @finished.setter
     def finished(self, finished):
-        self.grammar.finished = finished
+        if self.grammar is not None:
+            self.grammar.finished = finished
+        else:
+            self._finished = finished
 
     def try_jump_forward(self, tokenizer):
-        return self.grammar.try_jump_forward(tokenizer)
+        if self.grammar is not None:
+            return self.grammar.try_jump_forward(tokenizer)
+        return None
 
     def jump_forward_str_state(self, helper):
-        return self.grammar.jump_forward_str_state(helper)
+        if self.grammar is not None:
+            return self.grammar.jump_forward_str_state(helper)
+        return None
 
-    def jump_and_retokenize(
-        self, old_output_ids: List[int], new_output_ids: List[int], next_state: int
-    ):
-        return self.grammar.jump_and_retokenize(
-            old_output_ids, new_output_ids, next_state
-        )
+    def jump_and_retokenize(self, old_output_ids, new_output_ids, next_state):
+        if self.grammar is not None:
+            return self.grammar.jump_and_retokenize(
+                old_output_ids, new_output_ids, next_state
+            )
 
 
 class ReasonerGrammarBackend(BaseGrammarBackend):
-    def __init__(self, grammar_backend: BaseGrammarBackend, think_end_id):
+    def __init__(
+        self,
+        grammar_backend: BaseGrammarBackend,
+        reasoning_parser: ReasoningParser,
+        tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
+        enable_strict_thinking: bool = False,
+    ):
         super().__init__()
         self.grammar_backend = grammar_backend
-        self.think_end_id = think_end_id
+        think_end_ids = tokenizer.encode(
+            reasoning_parser.detector.think_end_token, add_special_tokens=False
+        )
+        if not think_end_ids:
+            raise ValueError(
+                f"think_end_token '{reasoning_parser.detector.think_end_token}' "
+                f"could not be encoded by the tokenizer."
+            )
+        if len(think_end_ids) != 1:
+            raise ValueError(
+                f"think_end_token '{reasoning_parser.detector.think_end_token}' "
+                "must encode to exactly one token for constrained reasoning."
+            )
+        self.think_end_id = think_end_ids[0]
+        self._enable_strict_thinking = enable_strict_thinking
+        self.think_excluded_token_ids = self._get_think_excluded_token_ids(
+            reasoning_parser, tokenizer
+        )
+        self.max_think_tokens = envs.SGLANG_MAX_THINK_TOKENS.get()
+        if (
+            self.enable_strict_thinking
+            and self.think_excluded_token_ids is not None
+            and not self.grammar_backend.is_support_token_filter
+        ):
+            raise ValueError(
+                "Strict reasoning format requested but the grammar backend does not "
+                "support token filtering. Use a grammar backend that supports token "
+                "filtering (e.g., xgrammar) or disable strict reasoning mode."
+            )
+        self.enable_token_filter = (
+            self.enable_strict_thinking
+            and self.think_excluded_token_ids is not None
+            and self.grammar_backend.is_support_token_filter
+        )
+        self._token_filter_fn = (
+            self.grammar_backend.set_token_filter if self.enable_token_filter else None
+        )
+
+    def _get_think_excluded_token_ids(
+        self,
+        reasoning_parser: ReasoningParser,
+        tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
+    ) -> Optional[List[int]]:
+        excluded_ids = []
+        if (not self.enable_strict_thinking) or (
+            not reasoning_parser.detector.think_excluded_tokens
+        ):
+            return None
+        for token in reasoning_parser.detector.think_excluded_tokens:
+            new_ids = tokenizer.encode(token, add_special_tokens=False)
+            if not new_ids:
+                raise ValueError(
+                    f"think_excluded_token '{token}' could not be encoded by the "
+                    f"tokenizer. All excluded tokens must be encodable for strict "
+                    f"reasoning mode to function correctly."
+                )
+            excluded_ids += new_ids
+        return excluded_ids
+
+    def _make_grammar_object(
+        self, grammar: Optional[BaseGrammarObject], reasoning: bool
+    ) -> ReasonerGrammarObject:
+        obj = ReasonerGrammarObject(
+            grammar=grammar,
+            think_end_id=self.think_end_id,
+            think_excluded_token_ids=self.think_excluded_token_ids,
+            max_think_tokens=self.max_think_tokens,
+            enable_token_filter=self.enable_token_filter,
+            token_filter_fn=self._token_filter_fn,
+            allocate_vocab_mask_fn=self.grammar_backend.allocate_vocab_mask,
+            move_vocab_mask_fn=self.grammar_backend.move_vocab_mask,
+            apply_vocab_mask_fn=self.grammar_backend.apply_vocab_mask,
+        )
+        obj.maybe_init_reasoning(reasoning)
+        return obj
+
+    def init_strict_reasoning_grammar(
+        self, reasoning: bool
+    ) -> Optional[BaseGrammarObject]:
+        """Create a grammar object for strict token filtering only (no inner grammar)."""
+        if not self.enable_strict_thinking:
+            return None
+        return self._make_grammar_object(None, reasoning)
 
     def _init_value_dispatch(
         self, key: Tuple[str, str], reasoning: bool
     ) -> Optional[BaseGrammarObject]:
         ret = self.grammar_backend._init_value_dispatch(key, reasoning)
-        # avoid wrapping invalid grammar, so that the scheduler can detect it
         if ret is None or isinstance(ret, InvalidGrammarObject):
             return ret
-        obj = ReasonerGrammarObject(ret, self.think_end_id)
-        obj.maybe_init_reasoning(reasoning)
-        return obj
+        return self._make_grammar_object(ret, reasoning)

python/sglang/srt/constrained/torch_ops/token_filter_torch_ops.py

@@ -0,0 +1,63 @@
+# Copyright 2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Torch fallback for token filter operations (non-CUDA devices and HIP).
+
+Sets or clears specific bits in an int32 bitmask by token ID.  The token list
+is typically tiny (< 10 entries); aggregation is done in Python with the actual
+bitmask operations using torch tensor indexing.
+"""
+
+import ctypes
+from typing import List
+
+import torch
+
+
+def set_token_filter_torch(
+    vocab_mask: torch.Tensor,
+    token_ids: List[int],
+    batch_idx: int,
+    is_allowed: bool = True,
+    reset_vocab_mask: bool = True,
+):
+    if reset_vocab_mask:
+        vocab_mask[batch_idx].fill_(-1 if (not is_allowed) else 0)
+
+    if not token_ids:
+        return
+
+    # Aggregate bit masks per int32 element to handle duplicate indices.
+    aggregated: dict[int, int] = {}
+    for token_id in token_ids:
+        element_idx = token_id // 32
+        bit_idx = token_id % 32
+        aggregated[element_idx] = aggregated.get(element_idx, 0) | (1 << bit_idx)
+
+    row = vocab_mask[batch_idx]
+    element_indices = torch.tensor(
+        list(aggregated.keys()), dtype=torch.long, device=row.device
+    )
+    bitmasks = torch.tensor(
+        [
+            ctypes.c_int32(mask if is_allowed else ~mask).value
+            for mask in aggregated.values()
+        ],
+        dtype=row.dtype,
+        device=row.device,
+    )
+
+    if is_allowed:
+        row[element_indices] = torch.bitwise_or(row[element_indices], bitmasks)
+    else:
+        row[element_indices] = torch.bitwise_and(row[element_indices], bitmasks)

python/sglang/srt/constrained/triton_ops/token_filter_ops.py

@@ -0,0 +1,175 @@
+# Copyright 2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Triton kernels for token filter operations."""
+
+from collections import OrderedDict
+from typing import List
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.utils import get_device_core_count
+
+
+@triton.jit
+def reset_vocab_mask_kernel(
+    vocab_mask_ptr,
+    batch_idx: int,
+    num_elements: int,
+    reset_value: tl.constexpr,
+):
+    """Reset the vocab mask for a specific batch index to a given value.
+
+    Parameters
+    ----------
+    vocab_mask_ptr : tl.tensor
+        Pointer to the vocab mask tensor.
+
+    batch_idx : int
+        The batch index to reset.
+
+    num_elements : int
+        Number of int32 elements in the vocab mask for each batch.
+
+    reset_value : int
+        The value to reset the vocab mask to (typically -1 or 0).
+    """
+    pid = tl.program_id(0)
+    num_threads = tl.num_programs(0)
+
+    for i in tl.range(pid, num_elements, num_threads):
+        offset = batch_idx * num_elements + i
+        tl.store(vocab_mask_ptr + offset, reset_value)
+
+
+@triton.jit
+def set_token_filter_batch_kernel(
+    vocab_mask_ptr,
+    token_ids_ptr,
+    batch_idx: int,
+    num_tokens: int,
+    num_elements: int,
+    is_allowed: tl.constexpr,
+):
+    """Set or clear specific tokens in the vocab mask for a batch.
+
+    Each token ID maps to a specific bit in the int32 bitmask array.
+    The kernel sets or clears those bits using atomic operations.
+
+    Parameters
+    ----------
+    vocab_mask_ptr : tl.tensor
+        Pointer to the vocab mask tensor.
+
+    token_ids_ptr : tl.tensor
+        Pointer to the token IDs to set/clear.
+
+    batch_idx : int
+        The batch index to modify.
+
+    num_tokens : int
+        Number of tokens to process.
+
+    num_elements : int
+        Number of int32 elements in the vocab mask for each batch.
+
+    is_allowed : bool
+        If True, set the bit to 1 (allow token).
+        If False, clear the bit to 0 (block token).
+    """
+    pid = tl.program_id(0)
+    num_threads = tl.num_programs(0)
+
+    for i in tl.range(pid, num_tokens, num_threads):
+        token_id = tl.load(token_ids_ptr + i)
+        element_idx = token_id // 32
+        bit_idx = token_id % 32
+
+        offset = batch_idx * num_elements + element_idx
+
+        if is_allowed:
+            tl.atomic_or(vocab_mask_ptr + offset, 1 << bit_idx)
+        else:
+            tl.atomic_and(vocab_mask_ptr + offset, ~(1 << bit_idx))
+
+
+_cached_num_sms = None
+_cached_token_id_tensors: OrderedDict[tuple[int, tuple[int, ...]], torch.Tensor] = (
+    OrderedDict()
+)
+_MAX_TOKEN_ID_TENSOR_CACHE_SIZE = 32
+
+
+def _compute_grid(work_items: int):
+    global _cached_num_sms
+    if _cached_num_sms is None:
+        _cached_num_sms = get_device_core_count()
+    if _cached_num_sms > 0:
+        return (min(_cached_num_sms, work_items),)
+    return (work_items,)
+
+
+def _get_cached_token_ids_tensor(
+    token_ids: List[int], device: torch.device
+) -> torch.Tensor:
+    key = (device.index or 0, tuple(token_ids))
+    cached = _cached_token_id_tensors.get(key)
+    if cached is not None:
+        _cached_token_id_tensors.move_to_end(key)
+        return cached
+
+    token_ids_tensor = torch.tensor(token_ids, dtype=torch.int32, device=device)
+    _cached_token_id_tensors[key] = token_ids_tensor
+    if len(_cached_token_id_tensors) > _MAX_TOKEN_ID_TENSOR_CACHE_SIZE:
+        _cached_token_id_tensors.popitem(last=False)
+    return token_ids_tensor
+
+
+def set_token_filter_triton(
+    vocab_mask: torch.Tensor,
+    token_ids: List[int],
+    batch_idx: int,
+    is_allowed: bool = True,
+    reset_vocab_mask: bool = True,
+):
+    """Set or clear specific tokens in the vocab mask using Triton."""
+    assert vocab_mask.device.type == "cuda"
+
+    num_elements = vocab_mask.shape[1]
+
+    if reset_vocab_mask:
+        reset_value = 0 if is_allowed else -1
+        reset_vocab_mask_kernel[_compute_grid(num_elements)](
+            vocab_mask,
+            batch_idx,
+            num_elements,
+            reset_value,
+            num_warps=4,
+        )
+
+    if not token_ids:
+        return
+
+    num_tokens = len(token_ids)
+    token_ids_tensor = _get_cached_token_ids_tensor(token_ids, vocab_mask.device)
+    set_token_filter_batch_kernel[_compute_grid(num_tokens)](
+        vocab_mask,
+        token_ids_tensor,
+        batch_idx,
+        num_tokens,
+        num_elements,
+        is_allowed,
+        num_warps=4,
+    )

python/sglang/srt/constrained/xgrammar_backend.py

@@ -49,6 +49,10 @@ else:
         apply_token_bitmask_inplace_triton,
     )
 
+from sglang.srt.constrained.torch_ops.token_filter_torch_ops import (
+    set_token_filter_torch,
+)
+from sglang.srt.constrained.triton_ops.token_filter_ops import set_token_filter_triton
 
 logger = logging.getLogger(__name__)
 MAX_ROLLBACK_TOKENS = 200
@@ -62,7 +66,7 @@ class XGrammarGrammar(BaseGrammarObject):
         vocab_size: int,
         ctx: CompiledGrammar,
         override_stop_tokens: Optional[Union[List[int], int]],
-        key_string: Optional[str] = None,  # TODO (sk): for debugging, remove later
+        key_string: Optional[str] = None,
         grammar_stats: Optional[GrammarStats] = GrammarStats(),
     ) -> None:
         super().__init__()
@@ -162,7 +166,14 @@ class XGrammarGrammar(BaseGrammarObject):
             self.matcher.rollback(len(old_output_ids) - k)
 
         for i in range(k, len(new_output_ids)):
-            assert self.matcher.accept_token(new_output_ids[i])
+            if not self.matcher.accept_token(new_output_ids[i]):
+                raise ValueError(
+                    f"Token not accepted during retokenization: {new_output_ids[i]} "
+                    f"at position {i}\n"
+                    f"Old output IDs: {old_output_ids}\n"
+                    f"New output IDs: {new_output_ids}\n"
+                    f"Key string: {self.key_string}"
+                )
 
     def __repr__(self):
         return f"XGrammarGrammar({self.key_string=}, {self.accepted_tokens=}, {self.current_token=})"
@@ -211,6 +222,53 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
         self.override_stop_tokens = override_stop_tokens
         self.any_whitespace = any_whitespace
 
+    @property
+    def is_support_token_filter(self):
+        return True
+
+    @staticmethod
+    def allocate_vocab_mask(vocab_size: int, batch_size: int, device) -> torch.Tensor:
+        return allocate_token_bitmask(batch_size, vocab_size)
+
+    @staticmethod
+    def move_vocab_mask(vocab_mask: torch.Tensor, device) -> torch.Tensor:
+        return vocab_mask.to(device, non_blocking=True)
+
+    @staticmethod
+    def apply_vocab_mask(logits: torch.Tensor, vocab_mask: torch.Tensor) -> None:
+        if logits.device.type in {"cuda", "npu", "xpu", "musa"}:
+            if _is_hip:
+                apply_token_bitmask_inplace_cuda(logits, vocab_mask)
+            else:
+                apply_token_bitmask_inplace_triton(logits, vocab_mask)
+        else:
+            raise RuntimeError(f"Unsupported device: {logits.device.type}")
+
+    @staticmethod
+    def set_token_filter(
+        vocab_mask: torch.Tensor,
+        token_ids: List[int],
+        batch_idx: int,
+        is_allowed: bool = True,
+        reset_vocab_mask: bool = True,
+    ):
+        if _is_hip or (vocab_mask.device.type != "cuda"):
+            set_token_filter_torch(
+                vocab_mask,
+                token_ids,
+                batch_idx,
+                is_allowed=is_allowed,
+                reset_vocab_mask=reset_vocab_mask,
+            )
+        else:
+            set_token_filter_triton(
+                vocab_mask,
+                token_ids,
+                batch_idx,
+                is_allowed=is_allowed,
+                reset_vocab_mask=reset_vocab_mask,
+            )
+
     @staticmethod
     def _sanitize_structural_format(structural_format):
         """Recursively replace missing json_schema fields with an empty schema."""

python/sglang/srt/environ.py

@@ -507,6 +507,9 @@ class Envs:
     # Tool-Call behavior
     SGLANG_TOOL_STRICT_LEVEL = EnvInt(ToolStrictLevel.OFF)
 
+    # Think tokens budget: negative means unlimited, >= 0 caps thinking tokens
+    SGLANG_MAX_THINK_TOKENS = EnvInt(-1)
+
     # Ngram
     SGLANG_NGRAM_FORCE_GREEDY_VERIFY = EnvBool(False)
 

python/sglang/srt/parser/reasoning_parser.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional, Tuple, Type
+from typing import Dict, List, Optional, Tuple, Type
 
 from sglang.srt.entrypoints.openai.protocol import ChatCompletionRequest
 from sglang.srt.parser.harmony_parser import HarmonyParser
@@ -23,6 +23,7 @@ class BaseReasoningFormatDetector:
         self,
         think_start_token: str,
         think_end_token: str,
+        think_excluded_tokens: Optional[List[str]] = None,
         force_reasoning: bool = False,
         stream_reasoning: bool = True,
         tool_start_token: Optional[str] = None,
@@ -33,6 +34,7 @@ class BaseReasoningFormatDetector:
     ):
         self.think_start_token = think_start_token
         self.think_end_token = think_end_token
+        self.think_excluded_tokens = think_excluded_tokens
         self.tool_start_token = tool_start_token
         self.force_reasoning = force_reasoning
         self._in_reasoning = force_reasoning
@@ -242,9 +244,16 @@ class Qwen3Detector(BaseReasoningFormatDetector):
         continue_final_message: bool = False,
         previous_content: str = "",
     ):
+        think_excluded_tokens = [
+            "<tool_call>",
+            "</tool_call>",
+            "<|im_end|>",
+            "<|endoftext|>",
+        ]
         super().__init__(
             "<think>",
             "</think>",
+            think_excluded_tokens=think_excluded_tokens,
             force_reasoning=force_reasoning,
             stream_reasoning=stream_reasoning,
             continue_final_message=continue_final_message,
@@ -297,9 +306,22 @@ class KimiK2Detector(BaseReasoningFormatDetector):
         continue_final_message: bool = False,
         previous_content: str = "",
     ):
+        think_excluded_tokens = [
+            "<think>",
+            "<|tool_calls_section_begin|>",
+            "<|tool_call_begin|>",
+            "<|tool_call_argument_begin|>",
+            "<|tool_call_section_end|>",
+            "<|tool_call_end|>",
+            "[EOS]",
+            "<|im_end|>",
+            "<|end_header_id|>",
+            "[EOT]",
+        ]
         super().__init__(
             "<think>",
             "</think>",
+            think_excluded_tokens=think_excluded_tokens,
             force_reasoning=force_reasoning,
             stream_reasoning=stream_reasoning,
             tool_start_token="<|tool_calls_section_begin|>",
@@ -323,9 +345,17 @@ class Glm45Detector(BaseReasoningFormatDetector):
     """
 
     def __init__(self, stream_reasoning: bool = True, force_reasoning: bool = False):
+        think_excluded_tokens = [
+            "<tool_call>",
+            "</tool_call>",
+            "<eop>",
+            "<|user|>",
+            "<|endoftext|>",
+        ]
         super().__init__(
             "<think>",
             "</think>",
+            think_excluded_tokens=think_excluded_tokens,
             force_reasoning=force_reasoning,
             stream_reasoning=stream_reasoning,
             tool_start_token="<tool_call>",

python/sglang/srt/server_args.py

@@ -495,6 +495,7 @@ class ServerArgs:
     enable_cache_report: bool = False
     reasoning_parser: Optional[str] = None
     strip_thinking_cache: bool = False
+    enable_strict_thinking: bool = False
     tool_call_parser: Optional[str] = None
     tool_server: Optional[str] = None
     sampling_defaults: str = "model"
@@ -5127,6 +5128,14 @@ class ServerArgs:
             "radix tree on finish; keep only the prompt prefix. Opt-in: changes "
             "cache contents.",
         )
+        parser.add_argument(
+            "--enable-strict-thinking",
+            action="store_true",
+            default=ServerArgs.enable_strict_thinking,
+            help="Enable strict token filtering during the thinking phase. "
+            "Blocks model-specific excluded tokens (e.g., tool call markers) "
+            "during reasoning. Requires a grammar backend that supports token filtering.",
+        )
         tool_call_parser_choices = list(FunctionCallParser.ToolCallParserEnum.keys())
         parser.add_argument(
             "--tool-call-parser",