minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Xinyuan Tong
2026-05-07 14:21:51 -07:00
committed by GitHub
parent af2a2ac618
commit 5b589ed2e7
14 changed files with 1713 additions and 430 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/constrained/base_grammar_backend.py
+45 -1
View File
@@ -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
+38
View File
@@ -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
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
+232 -38
View File
@@ -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):
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:
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):
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):
if self.grammar is not None:
return self.grammar.finished
return self._finished
@finished.setter
def finished(self, finished):
if self.grammar is not None:
self.grammar.finished = finished
else:
self._finished = finished
def try_jump_forward(self, tokenizer):
if self.grammar is not None:
return self.grammar.try_jump_forward(tokenizer)
return None
def jump_forward_str_state(self, 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
):
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
+63
View File
@@ -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
+175
View File
@@ -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
+60 -2
View File
@@ -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
+3
View File
@@ -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
+31 -1
View File
@@ -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
+9
View File
@@ -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",
test/registered/unit/constrained/test_base_grammar_backend.py
+15 -5
View File
@@ -269,10 +269,13 @@ class TestCreateGrammarBackend(unittest.TestCase):
GRAMMAR_BACKEND_REGISTRY.clear()
GRAMMAR_BACKEND_REGISTRY.update(self._saved)
def _make_server_args(self, backend="none", reasoning_parser=None):
def _make_server_args(
self, backend="none", reasoning_parser=None, enable_strict_thinking=False
):
args = MagicMock()
args.grammar_backend = backend
args.reasoning_parser = reasoning_parser
args.enable_strict_thinking = enable_strict_thinking
args.constrained_json_whitespace_pattern = None
args.constrained_json_disable_any_whitespace = False
return args
@@ -282,6 +285,11 @@ class TestCreateGrammarBackend(unittest.TestCase):
result = create_grammar_backend(args, None, 32000)
self.assertIsNone(result)
def test_none_backend_with_strict_thinking_raises(self):
args = self._make_server_args("none", enable_strict_thinking=True)
with self.assertRaisesRegex(ValueError, "enable-strict-thinking"):
create_grammar_backend(args, None, 32000)
def test_invalid_backend_raises(self):
args = self._make_server_args("nonexistent_backend")
with self.assertRaises(ValueError):
@@ -316,7 +324,7 @@ class TestCreateGrammarBackend(unittest.TestCase):
mock_inner = MagicMock(spec=BaseGrammarBackend)
register_grammar_backend("inner_r", lambda *a: mock_inner)
args = self._make_server_args("inner_r", reasoning_parser="deepseek")
args = self._make_server_args("inner_r", reasoning_parser="deepseek-r1")
tokenizer = MagicMock()
result = create_grammar_backend(args, tokenizer, 32000)
@@ -382,13 +390,15 @@ class TestCreateGrammarBackend(unittest.TestCase):
)
mock_backend = MagicMock(spec=BaseGrammarBackend)
mock_backend.is_support_token_filter = False
mock_outlines_cls.return_value = mock_backend
args = self._make_server_args("outlines", reasoning_parser="deepseek")
args = self._make_server_args("outlines", reasoning_parser="deepseek-r1")
tokenizer = MagicMock()
# encode must return a single-token list for think_start/end tokens
tokenizer.encode.return_value = [42]
result = create_grammar_backend(args, tokenizer, 32000, think_end_id=42)
self.assertIsInstance(result, ReasonerGrammarBackend)
self.assertEqual(result.think_end_id, 42)
self.assertIs(result.grammar_backend, mock_backend)
@patch("sglang.srt.constrained.outlines_backend.OutlinesGrammarBackend")
@@ -396,7 +406,7 @@ class TestCreateGrammarBackend(unittest.TestCase):
"""Without think_end_id passed in, no reasoner wrapping."""
mock_backend = MagicMock(spec=BaseGrammarBackend)
mock_outlines_cls.return_value = mock_backend
args = self._make_server_args("outlines", reasoning_parser="deepseek")
args = self._make_server_args("outlines", reasoning_parser="deepseek-r1")
tokenizer = MagicMock(spec=[]) # No think_end_id attribute
result = create_grammar_backend(args, tokenizer, 32000, think_end_id=None)
test/registered/unit/constrained/test_e2e_constrained_reasoning.py
+314
View File
@@ -0,0 +1,314 @@
"""
End-to-end tests for strict reasoning + constrained decoding.
Tests that the full pipeline works:
- AC-5.1: Strict reasoning + JSON schema constrained generation
- AC-5.2: Strict reasoning + tool call parsing (basic validation only)
These tests launch a real server with a small model and verify
the constrained decoding pipeline produces valid output.
"""
import json
import unittest
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
CustomTestCase,
popen_launch_server,
)
register_cuda_ci(est_time=120, suite="stage-b-test-1-gpu-small")
MODEL = "Qwen/Qwen3-0.6B"
BASE_URL = "http://127.0.0.1:39877"
API_KEY = "sk-test-1234"
class TestConstrainedReasoningE2E(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = MODEL
cls.base_url = BASE_URL
cls.api_key = API_KEY
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
api_key=cls.api_key,
other_args=[
"--reasoning-parser",
"qwen3",
],
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def _chat(self, **kwargs):
default = {
"model": self.model,
"messages": [
{
"role": "user",
"content": "What is 2+2? Answer with just the number.",
}
],
"temperature": 0,
"max_tokens": 256,
}
default.update(kwargs)
resp = requests.post(
f"{self.base_url}/v1/chat/completions",
headers={"Authorization": f"Bearer {self.api_key}"},
json=default,
timeout=60,
)
self.assertEqual(resp.status_code, 200, f"Request failed: {resp.text}")
return resp.json()
def test_reasoning_with_json_schema(self):
"""AC-5.1: Reasoning + JSON schema produces valid JSON output."""
schema = {
"type": "object",
"properties": {
"answer": {"type": "integer"},
},
"required": ["answer"],
}
data = self._chat(
response_format={
"type": "json_schema",
"json_schema": {
"name": "answer_schema",
"schema": schema,
},
},
chat_template_kwargs={"enable_thinking": True},
separate_reasoning=True,
)
choice = data["choices"][0]
content = choice["message"]["content"] or ""
# Content should be valid JSON conforming to schema when non-empty.
# With small models + separate_reasoning, content may be empty if the
# model puts everything in reasoning_content. That's acceptable.
if content.strip():
try:
parsed = json.loads(content)
self.assertIn("answer", parsed)
self.assertIsInstance(parsed["answer"], int)
except (json.JSONDecodeError, TypeError):
# Small models may produce imperfect JSON
self.assertTrue(
content.strip().startswith("{"),
f"Expected JSON-like output, got: {content!r}",
)
# Content should NOT contain <think> tags (those go to reasoning_content)
self.assertNotIn("<think>", content)
def test_reasoning_disabled_with_json_schema(self):
"""JSON schema still works when reasoning is explicitly disabled."""
schema = {
"type": "object",
"properties": {
"answer": {"type": "integer"},
},
"required": ["answer"],
}
data = self._chat(
response_format={
"type": "json_schema",
"json_schema": {
"name": "answer_schema",
"schema": schema,
},
},
chat_template_kwargs={"enable_thinking": False},
)
choice = data["choices"][0]
content = choice["message"]["content"]
# Should still produce valid JSON
parsed = json.loads(content)
self.assertIn("answer", parsed)
def test_reasoning_with_separate_output(self):
"""Reasoning content is correctly separated from normal content."""
data = self._chat(
chat_template_kwargs={"enable_thinking": True},
separate_reasoning=True,
)
choice = data["choices"][0]
content = choice["message"]["content"]
reasoning = choice["message"].get("reasoning_content")
# Content should not contain think tags
self.assertNotIn("<think>", content)
self.assertNotIn("</think>", content)
def test_tool_call_after_reasoning(self):
"""AC-5.2: Tool call parsing works with reasoning enabled."""
tools = [
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the current weather",
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string"},
},
"required": ["location"],
},
},
}
]
data = self._chat(
messages=[
{
"role": "user",
"content": "What's the weather in Paris?",
}
],
tools=tools,
chat_template_kwargs={"enable_thinking": True},
separate_reasoning=True,
)
choice = data["choices"][0]
# The model may or may not produce tool calls (depends on model capability)
# but the response should be well-formed (no crashes)
self.assertIn("message", choice)
self.assertIn("finish_reason", choice)
# finish_reason should be either "stop" or "tool_calls"
self.assertIn(choice["finish_reason"], ["stop", "tool_calls", "length"])
class TestStrictThinkingE2E(CustomTestCase):
"""E2E tests with --enable-strict-thinking flag.
Validates that the strict thinking flag is correctly propagated through
the full pipeline: server_args -> grammar_backend -> ReasonerGrammarBackend
-> token filtering during thinking phase.
"""
@classmethod
def setUpClass(cls):
cls.model = MODEL
cls.base_url = "http://127.0.0.1:39878"
cls.api_key = API_KEY
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
api_key=cls.api_key,
other_args=[
"--reasoning-parser",
"qwen3",
"--enable-strict-thinking",
],
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def _chat(self, **kwargs):
default = {
"model": self.model,
"messages": [
{
"role": "user",
"content": "What is 2+2? Answer with just the number.",
}
],
"temperature": 0,
"max_tokens": 256,
}
default.update(kwargs)
resp = requests.post(
f"{self.base_url}/v1/chat/completions",
headers={"Authorization": f"Bearer {self.api_key}"},
json=default,
timeout=60,
)
self.assertEqual(resp.status_code, 200, f"Request failed: {resp.text}")
return resp.json()
def test_strict_thinking_with_json_schema(self):
"""Strict thinking + JSON schema: server starts and produces valid output."""
schema = {
"type": "object",
"properties": {
"answer": {"type": "integer"},
},
"required": ["answer"],
}
data = self._chat(
response_format={
"type": "json_schema",
"json_schema": {
"name": "answer_schema",
"schema": schema,
},
},
chat_template_kwargs={"enable_thinking": True},
separate_reasoning=True,
)
choice = data["choices"][0]
content = choice["message"]["content"] or ""
if content.strip():
try:
parsed = json.loads(content)
self.assertIn("answer", parsed)
except (json.JSONDecodeError, TypeError):
self.assertTrue(
content.strip().startswith("{"),
f"Expected JSON-like output, got: {content!r}",
)
# Think tags must not leak into content
self.assertNotIn("<think>", content)
def test_strict_thinking_disabled_per_request(self):
"""When thinking is disabled per-request, strict server still works."""
data = self._chat(
chat_template_kwargs={"enable_thinking": False},
)
choice = data["choices"][0]
self.assertIn("message", choice)
self.assertIn("finish_reason", choice)
# Should complete normally without errors
self.assertIn(choice["finish_reason"], ["stop", "length"])
def test_strict_thinking_separate_reasoning(self):
"""Strict thinking with separate_reasoning produces well-formed output."""
data = self._chat(
chat_template_kwargs={"enable_thinking": True},
separate_reasoning=True,
)
choice = data["choices"][0]
content = choice["message"]["content"] or ""
# Think tags must not leak into content
self.assertNotIn("<think>", content)
self.assertNotIn("</think>", content)
if __name__ == "__main__":
unittest.main()
test/registered/unit/constrained/test_grammar_manager.py
+164 -5
View File
@@ -24,6 +24,7 @@ from sglang.srt.constrained.base_grammar_backend import (
InvalidGrammarObject,
)
from sglang.srt.constrained.grammar_manager import GrammarManager
from sglang.srt.constrained.reasoner_grammar_backend import ReasonerGrammarObject
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(2.0, "stage-a-test-cpu")
@@ -48,7 +49,12 @@ def _make_scheduler(grammar_backend_name="none", skip_tokenizer=False):
def _make_req(
json_schema=None, regex=None, ebnf=None, structural_tag=None, rid="req-1"
json_schema=None,
regex=None,
ebnf=None,
structural_tag=None,
rid="req-1",
custom_params=None,
):
"""Create a mock request with sampling params."""
req = MagicMock()
@@ -57,6 +63,7 @@ def _make_req(
req.sampling_params.regex = regex
req.sampling_params.ebnf = ebnf
req.sampling_params.structural_tag = structural_tag
req.sampling_params.custom_params = custom_params
req.require_reasoning = False
req.grammar = None
req.grammar_key = None
@@ -256,6 +263,39 @@ class TestProcessReqWithGrammar(unittest.TestCase):
self.assertTrue(mgr.has_waiting_grammars())
self.assertEqual(len(mgr), 1)
def test_cache_hit_applies_request_thinking_budget(self):
mgr = self._make_mgr()
grammar_obj = ReasonerGrammarObject(
grammar=None, think_end_id=0, max_think_tokens=99
)
mgr.grammar_backend.get_cached_or_future_value.return_value = (
grammar_obj,
True,
)
req = _make_req(
json_schema="schema",
custom_params={"thinking_budget": 7},
)
mgr.process_req_with_grammar(req)
self.assertEqual(req.grammar.max_think_tokens, 7)
def test_strict_reasoning_grammar_applies_request_thinking_budget(self):
mgr = self._make_mgr()
mgr._enable_strict_thinking = True
grammar_obj = ReasonerGrammarObject(
grammar=None, think_end_id=0, max_think_tokens=99
)
mgr.grammar_backend.init_strict_reasoning_grammar.return_value = grammar_obj
req = _make_req(custom_params={"thinking_budget": 3})
req.require_reasoning = True
mgr.process_req_with_grammar(req)
self.assertIs(req.grammar, grammar_obj)
self.assertEqual(req.grammar.max_think_tokens, 3)
class TestAbortRequests(unittest.TestCase):
"""Test abort_requests handling."""
@@ -494,8 +534,8 @@ class TestGetReadyGrammarRequests(unittest.TestCase):
req.set_finish_with_abort.assert_called_once()
self.assertIn("timed out", req.set_finish_with_abort.call_args[0][0])
def test_future_exception_propagates(self):
"""A future that raised an exception should propagate on .result()."""
def test_future_exception_creates_invalid_grammar_object(self):
"""A future that raised an exception should create InvalidGrammarObject, not crash."""
mgr = self._make_mgr()
future = Future()
@@ -506,8 +546,32 @@ class TestGetReadyGrammarRequests(unittest.TestCase):
req.grammar_key = ("json", "crash")
mgr.grammar_queue.append(req)
with self.assertRaises(RuntimeError):
mgr.get_ready_grammar_requests()
result = mgr.get_ready_grammar_requests()
self.assertEqual(len(result), 1)
self.assertIsInstance(result[0].grammar, InvalidGrammarObject)
req.set_finish_with_abort.assert_called_once()
def test_ready_future_applies_request_budget_without_polluting_cache(self):
mgr = self._make_mgr()
grammar_obj = ReasonerGrammarObject(
grammar=None, think_end_id=0, max_think_tokens=99
)
future = Future()
future.set_result(grammar_obj)
req = _make_req(json_schema="schema", custom_params={"thinking_budget": 4})
req.grammar = future
req.grammar_key = ("json", "schema")
mgr.grammar_queue.append(req)
result = mgr.get_ready_grammar_requests()
self.assertEqual(len(result), 1)
self.assertEqual(req.grammar.max_think_tokens, 4)
cached_key, cached_value = mgr.grammar_backend.set_cache.call_args[0]
self.assertEqual(cached_key, ("json", "schema"))
self.assertEqual(cached_value.max_think_tokens, 99)
@patch("sglang.srt.constrained.grammar_manager.torch.distributed.all_gather_object")
def test_multi_rank_sync_intersects_ready_unions_failed(self, mock_all_gather):
@@ -579,5 +643,100 @@ class TestGetReadyGrammarRequests(unittest.TestCase):
self.assertEqual(len(mgr.grammar_queue), 0)
class TestStrictReasoningPaths(unittest.TestCase):
"""Test _enable_strict_thinking code paths in GrammarManager."""
def _make_mgr(self):
scheduler = _make_scheduler()
scheduler.server_args.skip_tokenizer_init = True
mgr = GrammarManager(scheduler)
mgr.grammar_backend = MagicMock(spec=BaseGrammarBackend)
mgr._enable_strict_thinking = True
return mgr
def test_strict_unconstrained_request_gets_strict_grammar(self):
"""Request without json_schema/regex/ebnf should get strict-only grammar."""
mgr = self._make_mgr()
grammar_obj = MagicMock()
mgr.grammar_backend.init_strict_reasoning_grammar.return_value = grammar_obj
req = _make_req() # No constraint
req.require_reasoning = True
result = mgr.process_req_with_grammar(req)
self.assertFalse(result) # Not added to grammar queue
self.assertIs(req.grammar, grammar_obj)
mgr.grammar_backend.init_strict_reasoning_grammar.assert_called_once_with(True)
def test_strict_unconstrained_no_reasoning_flag(self):
"""Unconstrained request with require_reasoning=False still gets strict grammar."""
mgr = self._make_mgr()
grammar_obj = MagicMock()
mgr.grammar_backend.init_strict_reasoning_grammar.return_value = grammar_obj
req = _make_req()
req.require_reasoning = False
mgr.process_req_with_grammar(req)
self.assertIs(req.grammar, grammar_obj)
mgr.grammar_backend.init_strict_reasoning_grammar.assert_called_once_with(False)
def test_strict_unconstrained_none_grammar_is_fine(self):
"""If init_strict_reasoning_grammar returns None, req.grammar stays None."""
mgr = self._make_mgr()
mgr.grammar_backend.init_strict_reasoning_grammar.return_value = None
req = _make_req()
req.require_reasoning = True
mgr.process_req_with_grammar(req)
self.assertIsNone(req.grammar)
def test_strict_constrained_request_uses_normal_dispatch(self):
"""Request with json_schema should go through normal dispatch, not strict path."""
mgr = self._make_mgr()
future = MagicMock(spec=Future)
mgr.grammar_backend.get_cached_or_future_value.return_value = (future, False)
req = _make_req(json_schema='{"type": "object"}')
req.require_reasoning = True
result = mgr.process_req_with_grammar(req)
self.assertTrue(result) # Added to grammar queue
mgr.grammar_backend.init_strict_reasoning_grammar.assert_not_called()
def test_strict_not_set_skips_strict_path(self):
"""When _enable_strict_thinking=False, unconstrained requests get no grammar."""
mgr = self._make_mgr()
mgr._enable_strict_thinking = False
req = _make_req()
req.require_reasoning = True
mgr.process_req_with_grammar(req)
self.assertIsNone(req.grammar)
mgr.grammar_backend.init_strict_reasoning_grammar.assert_not_called()
def test_future_exception_creates_invalid_grammar(self):
"""Future.result() raising should create InvalidGrammarObject, not crash."""
mgr = self._make_mgr()
future = Future()
future.set_exception(RuntimeError("compilation failed"))
req = _make_req(json_schema='{"type": "object"}')
req.require_reasoning = True
req.grammar = future
req.grammar_key = ("json", '{"type": "object"}')
mgr.grammar_queue.append(req)
mgr.SGLANG_GRAMMAR_POLL_INTERVAL = 0.001
result = mgr.get_ready_grammar_requests()
self.assertEqual(len(result), 1)
self.assertIsInstance(result[0].grammar, InvalidGrammarObject)
req.set_finish_with_abort.assert_called_once()
if __name__ == "__main__":
unittest.main()
test/registered/unit/constrained/test_reasoner_grammar_backend.py
+407 -367
View File
@@ -1,413 +1,453 @@
"""
Unit tests for sglang.srt.constrained.reasoner_grammar_backend.
Test Coverage:
- ReasonerGrammarObject: state transitions, accept_token during thinking
vs post-thinking, rollback across think boundary, fill_vocab_mask gating,
copy semantics, finished delegation, delegation of jump methods
- ReasonerGrammarBackend: dispatch wrapping, invalid grammar passthrough,
None grammar passthrough, reasoning init on wrapped object
Usage:
python -m pytest test_reasoner_grammar_backend.py -v
"""
import os
import unittest
from unittest.mock import MagicMock, call
from types import SimpleNamespace
from unittest.mock import MagicMock
from sglang.srt.constrained.base_grammar_backend import (
BaseGrammarBackend,
BaseGrammarObject,
InvalidGrammarObject,
)
import torch
from sglang.srt.constrained.base_grammar_backend import BaseGrammarBackend
from sglang.srt.constrained.reasoner_grammar_backend import (
ReasonerGrammarBackend,
ReasonerGrammarObject,
)
from sglang.srt.constrained.torch_ops.token_filter_torch_ops import (
set_token_filter_torch,
)
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(2.0, "stage-a-test-cpu")
THINK_END_ID = 99
class _DummyTokenizer:
def __init__(self, token_map):
self._token_map = token_map
def encode(self, text, add_special_tokens=False):
return list(self._token_map.get(text, []))
class TestReasonerGrammarObjectStateTransitions(unittest.TestCase):
"""Test thinking state machine in ReasonerGrammarObject."""
class _DummyGrammarBackend(BaseGrammarBackend):
def __init__(self, support_token_filter=True):
super().__init__()
self._support_token_filter = support_token_filter
self._dispatch_result = None
def _make(self):
grammar = MagicMock(spec=BaseGrammarObject)
return ReasonerGrammarObject(grammar, THINK_END_ID), grammar
@property
def is_support_token_filter(self):
return self._support_token_filter
def test_initial_state_thinking(self):
obj, _ = self._make()
self.assertEqual(obj.tokens_after_think_end, -1)
@staticmethod
def allocate_vocab_mask(vocab_size, batch_size, device):
return torch.zeros((batch_size, (vocab_size + 31) // 32), dtype=torch.int32)
def test_transfer_state_during_thinking(self):
"""Regular tokens during thinking don't change state."""
obj, _ = self._make()
obj.transfer_state(10)
self.assertEqual(obj.tokens_after_think_end, -1)
@staticmethod
def move_vocab_mask(vocab_mask, device):
return vocab_mask
def test_transfer_state_think_end_token(self):
"""Think end token transitions from -1 to 0."""
obj, _ = self._make()
obj.transfer_state(THINK_END_ID)
self.assertEqual(obj.tokens_after_think_end, 0)
@staticmethod
def apply_vocab_mask(logits, vocab_mask):
return None
def test_transfer_state_increments_after_thinking(self):
"""After thinking ends, each token increments counter."""
obj, _ = self._make()
obj.tokens_after_think_end = 0
obj.transfer_state(10)
self.assertEqual(obj.tokens_after_think_end, 1)
obj.transfer_state(20)
self.assertEqual(obj.tokens_after_think_end, 2)
@staticmethod
def set_token_filter(
vocab_mask, token_ids, batch_idx, is_allowed=True, reset_vocab_mask=True
):
set_token_filter_torch(
vocab_mask, token_ids, batch_idx, is_allowed, reset_vocab_mask
)
def test_think_end_after_thinking_already_ended(self):
"""Second think_end_id after thinking ended just increments."""
obj, _ = self._make()
obj.tokens_after_think_end = 3
obj.transfer_state(THINK_END_ID)
self.assertEqual(obj.tokens_after_think_end, 4)
def test_rollback_state_from_post_thinking(self):
obj, _ = self._make()
obj.tokens_after_think_end = 3
obj.rollback_state()
self.assertEqual(obj.tokens_after_think_end, 2)
def test_rollback_state_at_boundary(self):
"""Rollback from 0 goes back to -1 (thinking)."""
obj, _ = self._make()
obj.tokens_after_think_end = 0
obj.rollback_state()
self.assertEqual(obj.tokens_after_think_end, -1)
def test_rollback_state_during_thinking(self):
"""Rollback during thinking stays at -1."""
obj, _ = self._make()
obj.rollback_state()
self.assertEqual(obj.tokens_after_think_end, -1)
def _init_value_dispatch(self, key, reasoning):
return self._dispatch_result
class TestReasonerGrammarObjectAcceptToken(unittest.TestCase):
"""Test accept_token behavior with thinking/post-thinking states."""
def _make(self):
grammar = MagicMock(spec=BaseGrammarObject)
return ReasonerGrammarObject(grammar, THINK_END_ID), grammar
def test_accept_during_thinking_skips_grammar(self):
"""During thinking phase, inner grammar should NOT receive tokens."""
obj, grammar = self._make()
obj.accept_token(10)
grammar.accept_token.assert_not_called()
# State should still be -1
self.assertEqual(obj.tokens_after_think_end, -1)
def test_accept_think_end_token(self):
"""Think end token transitions state but doesn't call inner grammar (state was -1 before transfer)."""
obj, grammar = self._make()
# tokens_after_think_end is -1, so grammar.accept_token is not called
# But wait: accept_token checks `>= 0` BEFORE transfer_state
# At call time tokens_after_think_end == -1, so grammar.accept_token skipped
obj.accept_token(THINK_END_ID)
grammar.accept_token.assert_not_called()
self.assertEqual(obj.tokens_after_think_end, 0)
def test_accept_after_thinking_calls_grammar(self):
"""After thinking ends, tokens go to inner grammar."""
obj, grammar = self._make()
obj.tokens_after_think_end = 0
obj.accept_token(42)
grammar.accept_token.assert_called_once_with(42)
self.assertEqual(obj.tokens_after_think_end, 1)
def test_accept_sequence_through_thinking_and_generation(self):
"""Full sequence: think tokens -> think_end -> generation tokens."""
obj, grammar = self._make()
# Thinking phase
obj.accept_token(1)
obj.accept_token(2)
self.assertEqual(grammar.accept_token.call_count, 0)
# Think end
obj.accept_token(THINK_END_ID)
self.assertEqual(grammar.accept_token.call_count, 0)
# Generation phase
obj.accept_token(10)
obj.accept_token(20)
self.assertEqual(grammar.accept_token.call_count, 2)
grammar.accept_token.assert_has_calls([call(10), call(20)])
def _allowed_token_ids(vocab_mask, token_ids):
allowed = []
for token_id in token_ids:
elem = token_id // 32
bit = token_id % 32
if int(vocab_mask[0, elem].item()) & (1 << bit):
allowed.append(token_id)
return allowed
class TestReasonerGrammarObjectRollback(unittest.TestCase):
"""Test rollback across thinking boundary."""
class TestReasonerGrammarObject(unittest.TestCase):
def _make_strict_object(self):
return ReasonerGrammarObject(
grammar=None,
think_end_id=7,
think_excluded_token_ids=[3, 5],
max_think_tokens=2,
enable_token_filter=True,
token_filter_fn=set_token_filter_torch,
allocate_vocab_mask_fn=lambda vocab_size, batch_size, device: torch.zeros(
(batch_size, (vocab_size + 31) // 32), dtype=torch.int32
),
move_vocab_mask_fn=lambda vocab_mask, device: vocab_mask,
apply_vocab_mask_fn=lambda logits, vocab_mask: None,
)
def _make(self):
grammar = MagicMock(spec=BaseGrammarObject)
return ReasonerGrammarObject(grammar, THINK_END_ID), grammar
def test_rollback_within_generation(self):
"""Rollback entirely within generation phase."""
obj, grammar = self._make()
obj.tokens_after_think_end = 5
obj.rollback(3)
grammar.rollback.assert_called_once_with(3)
self.assertEqual(obj.tokens_after_think_end, 2)
def test_rollback_across_boundary(self):
"""Rollback that crosses from generation back into thinking."""
obj, grammar = self._make()
obj.tokens_after_think_end = 2
obj.rollback(4)
# Only 2 tokens were post-thinking, so inner grammar rolls back 2
grammar.rollback.assert_called_once_with(2)
# After 4 rollback_state calls from 2: 2->1->0->-1->-1
self.assertEqual(obj.tokens_after_think_end, -1)
def test_rollback_during_thinking(self):
"""Rollback during thinking phase doesn't touch inner grammar."""
obj, grammar = self._make()
obj.rollback(3)
grammar.rollback.assert_not_called()
self.assertEqual(obj.tokens_after_think_end, -1)
def test_rollback_zero(self):
obj, grammar = self._make()
obj.tokens_after_think_end = 2
obj.rollback(0)
grammar.rollback.assert_not_called()
self.assertEqual(obj.tokens_after_think_end, 2)
def test_rollback_exactly_to_boundary(self):
"""Rollback exactly the number of post-thinking tokens."""
obj, grammar = self._make()
obj.tokens_after_think_end = 3
obj.rollback(3)
grammar.rollback.assert_called_once_with(3)
self.assertEqual(obj.tokens_after_think_end, 0)
def test_rollback_far_beyond_all_tokens(self):
"""Rollback k much larger than tokens_after_think_end clamps grammar rollback."""
obj, grammar = self._make()
obj.tokens_after_think_end = 2
obj.rollback(100)
# Inner grammar only rolls back the 2 post-thinking tokens
grammar.rollback.assert_called_once_with(2)
# State bottoms out at -1
self.assertEqual(obj.tokens_after_think_end, -1)
def test_accept_then_rollback_roundtrip(self):
"""Accept tokens then rollback should restore original state."""
obj, grammar = self._make()
obj.tokens_after_think_end = 0 # Just finished thinking
# Accept 3 generation tokens
obj.accept_token(10)
obj.accept_token(20)
obj.accept_token(30)
self.assertEqual(obj.tokens_after_think_end, 3)
self.assertEqual(grammar.accept_token.call_count, 3)
# Rollback all 3
obj.rollback(3)
self.assertEqual(obj.tokens_after_think_end, 0)
grammar.rollback.assert_called_once_with(3)
class TestReasonerGrammarObjectVocabMask(unittest.TestCase):
"""Test vocab mask gating based on thinking state."""
def _make(self):
grammar = MagicMock(spec=BaseGrammarObject)
return ReasonerGrammarObject(grammar, THINK_END_ID), grammar
def test_fill_during_thinking_skips(self):
obj, grammar = self._make()
obj.fill_vocab_mask("mask", 0)
grammar.fill_vocab_mask.assert_not_called()
def test_fill_after_thinking_delegates(self):
obj, grammar = self._make()
obj.tokens_after_think_end = 0
obj.fill_vocab_mask("mask", 0)
grammar.fill_vocab_mask.assert_called_once_with("mask", 0)
def test_fill_well_into_generation(self):
obj, grammar = self._make()
obj.tokens_after_think_end = 5
obj.fill_vocab_mask("mask", 2)
grammar.fill_vocab_mask.assert_called_once_with("mask", 2)
def test_fill_at_think_end_boundary(self):
"""After accepting think_end token, fill_vocab_mask should delegate."""
obj, grammar = self._make()
# Simulate: accept think_end, state goes from -1 to 0
obj.accept_token(THINK_END_ID)
self.assertEqual(obj.tokens_after_think_end, 0)
obj.fill_vocab_mask("mask", 0)
grammar.fill_vocab_mask.assert_called_once_with("mask", 0)
def test_allocate_delegates(self):
obj, grammar = self._make()
obj.allocate_vocab_mask(32000, 4, "cpu")
grammar.allocate_vocab_mask.assert_called_once_with(32000, 4, "cpu")
def test_move_delegates(self):
obj, grammar = self._make()
obj.move_vocab_mask("mask", "cuda")
grammar.move_vocab_mask.assert_called_once_with("mask", "cuda")
class TestReasonerGrammarObjectDelegation(unittest.TestCase):
"""Test that non-state methods delegate to inner grammar."""
def _make(self):
grammar = MagicMock(spec=BaseGrammarObject)
return ReasonerGrammarObject(grammar, THINK_END_ID), grammar
def test_is_terminated_delegates(self):
obj, grammar = self._make()
grammar.is_terminated.return_value = True
self.assertTrue(obj.is_terminated())
def test_finished_getter_delegates(self):
obj, grammar = self._make()
grammar.finished = True
self.assertTrue(obj.finished)
def test_finished_setter_delegates(self):
obj, grammar = self._make()
obj.finished = True
self.assertTrue(grammar.finished)
def test_try_jump_forward_delegates(self):
obj, grammar = self._make()
grammar.try_jump_forward.return_value = ([1, 2], "ab")
result = obj.try_jump_forward("tokenizer")
grammar.try_jump_forward.assert_called_once_with("tokenizer")
self.assertEqual(result, ([1, 2], "ab"))
def test_jump_forward_str_state_delegates(self):
obj, grammar = self._make()
grammar.jump_forward_str_state.return_value = ("str", 5)
result = obj.jump_forward_str_state("helper")
self.assertEqual(result, ("str", 5))
def test_jump_and_retokenize_delegates(self):
obj, grammar = self._make()
obj.jump_and_retokenize([1], [2], 3)
grammar.jump_and_retokenize.assert_called_once_with([1], [2], 3)
def test_apply_vocab_mask_property(self):
obj, grammar = self._make()
grammar.apply_vocab_mask = "mask_fn"
self.assertEqual(obj.apply_vocab_mask, "mask_fn")
def test_copy_creates_new_wrapper(self):
obj, grammar = self._make()
grammar_copy = MagicMock(spec=BaseGrammarObject)
grammar.copy.return_value = grammar_copy
copied = obj.copy()
self.assertIsInstance(copied, ReasonerGrammarObject)
self.assertIsNot(copied, obj)
self.assertIs(copied.grammar, grammar_copy)
self.assertEqual(copied.think_end_id, THINK_END_ID)
def test_copy_does_not_share_state(self):
"""Modifying copy's state should not affect the original."""
obj, grammar = self._make()
grammar_copy = MagicMock(spec=BaseGrammarObject)
grammar.copy.return_value = grammar_copy
copied = obj.copy()
copied.tokens_after_think_end = 5
self.assertEqual(obj.tokens_after_think_end, -1)
class TestReasonerGrammarObjectMaybeInitReasoning(unittest.TestCase):
"""Test maybe_init_reasoning state initialization."""
def test_reasoning_true_sets_thinking(self):
grammar = MagicMock(spec=BaseGrammarObject)
obj = ReasonerGrammarObject(grammar, THINK_END_ID)
def test_strict_thinking_phase_excludes_configured_tokens(self):
obj = self._make_strict_object()
obj.maybe_init_reasoning(True)
self.assertEqual(obj.tokens_after_think_end, -1)
mask = obj.allocate_vocab_mask(64, 1, "cpu")
def test_reasoning_false_skips_thinking(self):
grammar = MagicMock(spec=BaseGrammarObject)
obj = ReasonerGrammarObject(grammar, THINK_END_ID)
obj.maybe_init_reasoning(False)
self.assertEqual(obj.tokens_after_think_end, 0)
obj.fill_vocab_mask(mask, 0)
def test_reasoning_toggle(self):
"""Toggling reasoning resets state regardless of current position."""
grammar = MagicMock(spec=BaseGrammarObject)
obj = ReasonerGrammarObject(grammar, THINK_END_ID)
obj.tokens_after_think_end = 5 # Deep into generation
allowed = _allowed_token_ids(mask, [0, 1, 3, 5, 7, 8])
self.assertEqual(allowed, [0, 1, 7, 8])
def test_budget_exhaustion_allows_only_think_end(self):
obj = self._make_strict_object()
obj.maybe_init_reasoning(True)
self.assertEqual(obj.tokens_after_think_end, -1)
obj.accept_token(10)
obj.accept_token(11)
mask = obj.allocate_vocab_mask(64, 1, "cpu")
obj.maybe_init_reasoning(False)
self.assertEqual(obj.tokens_after_think_end, 0)
obj.fill_vocab_mask(mask, 0)
allowed = _allowed_token_ids(mask, [0, 1, 3, 5, 7, 8, 10, 11])
self.assertEqual(allowed, [7])
def test_strict_only_wrapper_exposes_backend_mask_hooks(self):
obj = self._make_strict_object()
mask = obj.allocate_vocab_mask(64, 2, "cpu")
self.assertEqual(mask.shape, (2, 2))
self.assertIs(obj.move_vocab_mask(mask, "cpu"), mask)
self.assertIsNotNone(obj.apply_vocab_mask)
class TestReasonerGrammarBackend(unittest.TestCase):
"""Test ReasonerGrammarBackend dispatch wrapping."""
def setUp(self):
self._prev_budget = os.environ.get("SGLANG_MAX_THINK_TOKENS")
def _make(self):
inner = MagicMock(spec=BaseGrammarBackend)
backend = ReasonerGrammarBackend(inner, THINK_END_ID)
return backend, inner
def tearDown(self):
if self._prev_budget is None:
os.environ.pop("SGLANG_MAX_THINK_TOKENS", None)
else:
os.environ["SGLANG_MAX_THINK_TOKENS"] = self._prev_budget
def test_wraps_valid_grammar(self):
backend, inner = self._make()
mock_grammar = MagicMock(spec=BaseGrammarObject)
inner._init_value_dispatch.return_value = mock_grammar
def _make_parser(self):
detector = SimpleNamespace(
think_start_token="<think>",
think_end_token="</think>",
think_excluded_tokens=["<tool_call>", "</tool_call>"],
)
return SimpleNamespace(detector=detector)
result = backend._init_value_dispatch(("json", "schema"), True)
self.assertIsInstance(result, ReasonerGrammarObject)
self.assertIs(result.grammar, mock_grammar)
self.assertEqual(result.think_end_id, THINK_END_ID)
def _make_tokenizer(self, start_ids=None, end_ids=None):
return _DummyTokenizer(
{
"<think>": [1] if start_ids is None else start_ids,
"</think>": [2] if end_ids is None else end_ids,
"<tool_call>": [3],
"</tool_call>": [4],
}
)
def test_passes_through_invalid_grammar(self):
backend, inner = self._make()
invalid = InvalidGrammarObject("bad grammar")
inner._init_value_dispatch.return_value = invalid
def test_init_strict_reasoning_grammar_uses_token_filter_and_budget(self):
os.environ["SGLANG_MAX_THINK_TOKENS"] = "2"
backend = _DummyGrammarBackend(support_token_filter=True)
reasoner = ReasonerGrammarBackend(
backend,
self._make_parser(),
self._make_tokenizer(),
enable_strict_thinking=True,
)
result = backend._init_value_dispatch(("json", "schema"), False)
self.assertIs(result, invalid)
self.assertIsInstance(result, InvalidGrammarObject)
obj = reasoner.init_strict_reasoning_grammar(reasoning=True)
def test_passes_through_none(self):
backend, inner = self._make()
inner._init_value_dispatch.return_value = None
self.assertIsInstance(obj, ReasonerGrammarObject)
self.assertTrue(obj.enable_token_filter)
self.assertEqual(obj.max_think_tokens, 2)
self.assertEqual(obj.think_excluded_token_ids, [3, 4])
result = backend._init_value_dispatch(("json", "schema"), False)
self.assertIsNone(result)
def test_init_strict_reasoning_grammar_none_when_strict_disabled(self):
backend = _DummyGrammarBackend(support_token_filter=True)
reasoner = ReasonerGrammarBackend(
backend,
self._make_parser(),
self._make_tokenizer(),
enable_strict_thinking=False,
)
def test_inits_reasoning_on_wrapped(self):
backend, inner = self._make()
mock_grammar = MagicMock(spec=BaseGrammarObject)
inner._init_value_dispatch.return_value = mock_grammar
self.assertIsNone(reasoner.init_strict_reasoning_grammar(reasoning=True))
result = backend._init_value_dispatch(("json", "schema"), True)
# reasoning=True → tokens_after_think_end should be -1
self.assertEqual(result.tokens_after_think_end, -1)
def test_wraps_inner_grammar_with_reasoning_state_machine(self):
os.environ["SGLANG_MAX_THINK_TOKENS"] = "1"
backend = _DummyGrammarBackend(support_token_filter=True)
inner_grammar = MagicMock()
backend._dispatch_result = inner_grammar
reasoner = ReasonerGrammarBackend(
backend,
self._make_parser(),
self._make_tokenizer(),
enable_strict_thinking=True,
)
def test_inits_no_reasoning_on_wrapped(self):
backend, inner = self._make()
mock_grammar = MagicMock(spec=BaseGrammarObject)
inner._init_value_dispatch.return_value = mock_grammar
wrapped = reasoner._init_value_dispatch(("json", "{}"), reasoning=True)
self.assertIsInstance(wrapped, ReasonerGrammarObject)
wrapped.accept_token(10)
inner_grammar.accept_token.assert_not_called()
wrapped.accept_token(2)
wrapped.accept_token(42)
inner_grammar.accept_token.assert_called_once_with(42)
result = backend._init_value_dispatch(("json", "schema"), False)
# reasoning=False → tokens_after_think_end should be 0
self.assertEqual(result.tokens_after_think_end, 0)
def test_accepts_multi_token_think_start_marker(self):
"""think_start_token can be multi-token (e.g., GPT-OSS) since it's not used."""
backend = _DummyGrammarBackend(support_token_filter=True)
reasoner = ReasonerGrammarBackend(
backend,
self._make_parser(),
self._make_tokenizer(start_ids=[1, 2]),
enable_strict_thinking=True,
)
self.assertIsNotNone(reasoner)
def test_rejects_multi_token_think_end_marker(self):
backend = _DummyGrammarBackend(support_token_filter=True)
with self.assertRaisesRegex(ValueError, "must encode to exactly one token"):
ReasonerGrammarBackend(
backend,
self._make_parser(),
self._make_tokenizer(end_ids=[2, 3]),
enable_strict_thinking=True,
)
def test_rejects_unencodable_excluded_token(self):
backend = _DummyGrammarBackend(support_token_filter=True)
parser = self._make_parser()
parser.detector.think_excluded_tokens = ["<unknown>"]
tokenizer = _DummyTokenizer(
{
"<think>": [1],
"</think>": [2],
}
)
with self.assertRaisesRegex(ValueError, "could not be encoded"):
ReasonerGrammarBackend(
backend,
parser,
tokenizer,
enable_strict_thinking=True,
)
def test_strict_mode_fails_when_backend_lacks_token_filter(self):
backend = _DummyGrammarBackend(support_token_filter=False)
with self.assertRaisesRegex(ValueError, "does not support token filtering"):
ReasonerGrammarBackend(
backend,
self._make_parser(),
self._make_tokenizer(),
enable_strict_thinking=True,
)
class TestReasonerGrammarObjectRollback(unittest.TestCase):
"""Tests for rollback correctness at the THINKING→GENERATION boundary."""
def _make_object_with_mock_grammar(self):
inner_grammar = MagicMock()
inner_grammar.is_terminated.return_value = False
obj = ReasonerGrammarObject(
grammar=inner_grammar,
think_end_id=7,
think_excluded_token_ids=[3, 5],
max_think_tokens=-1,
enable_token_filter=True,
token_filter_fn=set_token_filter_torch,
allocate_vocab_mask_fn=lambda vs, bs, d: torch.zeros(
(bs, (vs + 31) // 32), dtype=torch.int32
),
move_vocab_mask_fn=lambda vm, d: vm,
apply_vocab_mask_fn=lambda l, vm: None,
)
return obj, inner_grammar
def test_rollback_at_generation_boundary_returns_to_thinking(self):
obj, inner_grammar = self._make_object_with_mock_grammar()
obj.maybe_init_reasoning(True)
# Accept 3 thinking tokens then think_end_id
obj.accept_token(10)
obj.accept_token(11)
obj.accept_token(12)
obj.accept_token(7) # think_end_id → tokens_after_end = 0
self.assertTrue(obj._is_generation())
self.assertEqual(obj.tokens_after_end, 0)
# Rollback 1 step: should return to THINKING
obj.rollback(1)
self.assertTrue(obj._is_thinking())
self.assertEqual(obj.tokens_in_think, 3)
self.assertEqual(obj.tokens_after_end, -1)
# Grammar should not have been rolled back (no generation tokens were accepted)
inner_grammar.rollback.assert_not_called()
def test_rollback_spanning_both_phases(self):
obj, inner_grammar = self._make_object_with_mock_grammar()
obj.maybe_init_reasoning(True)
# 2 thinking tokens + think_end + 3 generation tokens
obj.accept_token(10) # think
obj.accept_token(11) # think
obj.accept_token(7) # think_end_id
obj.accept_token(20) # gen 1
obj.accept_token(21) # gen 2
obj.accept_token(22) # gen 3
self.assertEqual(obj.tokens_after_end, 3)
# Rollback 5: should roll back 3 generation tokens + think_end + 1 thinking token
obj.rollback(5)
self.assertTrue(obj._is_thinking())
self.assertEqual(obj.tokens_in_think, 1)
# Grammar should be rolled back by 3 (only generation tokens)
inner_grammar.rollback.assert_called_once_with(3)
def test_rollback_generation_tokens_only(self):
obj, inner_grammar = self._make_object_with_mock_grammar()
obj.maybe_init_reasoning(True)
obj.accept_token(10) # think
obj.accept_token(7) # think_end_id
obj.accept_token(20) # gen 1
obj.accept_token(21) # gen 2
# Rollback 1: should only roll back 1 generation token
obj.rollback(1)
self.assertTrue(obj._is_generation())
self.assertEqual(obj.tokens_after_end, 1)
inner_grammar.rollback.assert_called_once_with(1)
def test_rollback_thinking_tokens_does_not_touch_grammar(self):
obj, inner_grammar = self._make_object_with_mock_grammar()
obj.maybe_init_reasoning(True)
obj.accept_token(10)
obj.accept_token(11)
obj.accept_token(12)
obj.rollback(2)
self.assertTrue(obj._is_thinking())
self.assertEqual(obj.tokens_in_think, 1)
inner_grammar.rollback.assert_not_called()
inner_grammar.accept_token.assert_not_called()
def test_copy_preserves_state(self):
obj, inner_grammar = self._make_object_with_mock_grammar()
obj.maybe_init_reasoning(True)
obj.accept_token(10)
obj.accept_token(7) # think_end_id → GENERATION
obj.accept_token(20)
self.assertEqual(obj.tokens_in_think, 1)
self.assertEqual(obj.tokens_after_end, 1)
copy = obj.copy()
# State counters must be preserved for speculative decoding
self.assertEqual(copy.tokens_in_think, 1)
self.assertEqual(copy.tokens_after_end, 1)
self.assertTrue(copy._is_generation())
self.assertIsNotNone(copy.grammar)
inner_grammar.copy.assert_called_once()
def test_copy_preserves_thinking_state(self):
obj, inner_grammar = self._make_object_with_mock_grammar()
obj.maybe_init_reasoning(True)
obj.accept_token(10)
obj.accept_token(11)
copy = obj.copy()
self.assertEqual(copy.tokens_in_think, 2)
self.assertEqual(copy.tokens_after_end, -1)
self.assertTrue(copy._is_thinking())
class TestReasonerGrammarObjectFillVocabMask(unittest.TestCase):
"""Tests for fill_vocab_mask behavior in different states."""
def test_thinking_phase_does_not_consult_inner_grammar(self):
inner_grammar = MagicMock()
# Must return a real tensor for allocate_vocab_mask since fill_vocab_mask
# delegates to allocate_vocab_mask via self.grammar when grammar is not None
inner_grammar.allocate_vocab_mask.side_effect = lambda vs, bs, d: torch.zeros(
(bs, (vs + 31) // 32), dtype=torch.int32
)
obj = ReasonerGrammarObject(
grammar=inner_grammar,
think_end_id=7,
think_excluded_token_ids=[3, 5],
max_think_tokens=-1,
enable_token_filter=True,
token_filter_fn=set_token_filter_torch,
allocate_vocab_mask_fn=lambda vs, bs, d: torch.zeros(
(bs, (vs + 31) // 32), dtype=torch.int32
),
move_vocab_mask_fn=lambda vm, d: vm,
apply_vocab_mask_fn=lambda l, vm: None,
)
obj.maybe_init_reasoning(True)
mask = obj.allocate_vocab_mask(64, 1, "cpu")
obj.fill_vocab_mask(mask, 0)
inner_grammar.fill_vocab_mask.assert_not_called()
# Excluded tokens (3, 5) should be blocked
allowed = _allowed_token_ids(mask, [0, 1, 3, 5, 7, 8])
self.assertEqual(allowed, [0, 1, 7, 8])
def test_generation_phase_consults_inner_grammar(self):
inner_grammar = MagicMock()
inner_grammar.allocate_vocab_mask.side_effect = lambda vs, bs, d: torch.zeros(
(bs, (vs + 31) // 32), dtype=torch.int32
)
obj = ReasonerGrammarObject(
grammar=inner_grammar,
think_end_id=7,
think_excluded_token_ids=[3, 5],
max_think_tokens=-1,
enable_token_filter=True,
token_filter_fn=set_token_filter_torch,
allocate_vocab_mask_fn=lambda vs, bs, d: torch.zeros(
(bs, (vs + 31) // 32), dtype=torch.int32
),
move_vocab_mask_fn=lambda vm, d: vm,
apply_vocab_mask_fn=lambda l, vm: None,
)
obj.maybe_init_reasoning(True)
obj.accept_token(10)
obj.accept_token(7) # think_end_id → GENERATION
mask = obj.allocate_vocab_mask(64, 1, "cpu")
obj.fill_vocab_mask(mask, 0)
inner_grammar.fill_vocab_mask.assert_called_once_with(mask, 0)
def test_non_strict_thinking_is_noop(self):
inner_grammar = MagicMock()
obj = ReasonerGrammarObject(
grammar=inner_grammar,
think_end_id=7,
think_excluded_token_ids=None,
max_think_tokens=-1,
enable_token_filter=False,
token_filter_fn=None,
)
obj.maybe_init_reasoning(True)
mask = torch.zeros((1, 2), dtype=torch.int32)
obj.fill_vocab_mask(mask, 0)
inner_grammar.fill_vocab_mask.assert_not_called()
# Mask should remain all zeros (no filtering)
self.assertTrue(torch.all(mask == 0))
if __name__ == "__main__":
test/registered/unit/constrained/test_token_filter_ops.py
+146
View File
@@ -0,0 +1,146 @@
"""
Unit tests for token filter operations (Triton and Torch paths).
Verifies that both implementations produce identical bitmask output
for the same inputs, ensuring parity across GPU and CPU paths.
"""
import unittest
import torch
from sglang.srt.constrained.torch_ops.token_filter_torch_ops import (
set_token_filter_torch,
)
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(2.0, "stage-a-test-cpu")
# Conditionally import Triton path
_has_cuda = torch.cuda.is_available()
if _has_cuda:
from sglang.srt.constrained.triton_ops.token_filter_ops import (
set_token_filter_triton,
)
def _get_allowed_tokens(vocab_mask, batch_idx, max_token_id):
"""Extract allowed token IDs from a bitmask row."""
allowed = []
for token_id in range(max_token_id):
elem = token_id // 32
bit = token_id % 32
val = int(vocab_mask[batch_idx, elem].item())
if val & (1 << bit):
allowed.append(token_id)
return allowed
class TestSetTokenFilterTorch(unittest.TestCase):
"""Tests for the Torch token filter implementation."""
def test_allow_tokens_from_blank_mask(self):
vocab_mask = torch.zeros((1, 4), dtype=torch.int32) # 128 tokens
set_token_filter_torch(vocab_mask, [0, 5, 31, 32, 63], 0, is_allowed=True)
allowed = _get_allowed_tokens(vocab_mask, 0, 64)
self.assertEqual(allowed, [0, 5, 31, 32, 63])
def test_block_tokens_from_full_mask(self):
vocab_mask = torch.full((1, 4), -1, dtype=torch.int32) # all bits set
set_token_filter_torch(
vocab_mask, [3, 5], 0, is_allowed=False, reset_vocab_mask=False
)
allowed = _get_allowed_tokens(vocab_mask, 0, 64)
self.assertNotIn(3, allowed)
self.assertNotIn(5, allowed)
self.assertIn(0, allowed)
self.assertIn(1, allowed)
def test_reset_then_allow(self):
vocab_mask = torch.full((1, 2), -1, dtype=torch.int32)
set_token_filter_torch(
vocab_mask, [7], 0, is_allowed=True, reset_vocab_mask=True
)
allowed = _get_allowed_tokens(vocab_mask, 0, 64)
self.assertEqual(allowed, [7])
def test_reset_then_block(self):
vocab_mask = torch.zeros((1, 2), dtype=torch.int32)
set_token_filter_torch(
vocab_mask, [3, 5], 0, is_allowed=False, reset_vocab_mask=True
)
allowed = _get_allowed_tokens(vocab_mask, 0, 64)
self.assertNotIn(3, allowed)
self.assertNotIn(5, allowed)
# All other tokens should be allowed (reset to -1 for block mode)
self.assertIn(0, allowed)
self.assertIn(7, allowed)
def test_empty_token_list(self):
vocab_mask = torch.zeros((1, 2), dtype=torch.int32)
set_token_filter_torch(
vocab_mask, [], 0, is_allowed=True, reset_vocab_mask=True
)
allowed = _get_allowed_tokens(vocab_mask, 0, 64)
self.assertEqual(allowed, [])
def test_batch_indexing(self):
vocab_mask = torch.zeros((3, 2), dtype=torch.int32)
set_token_filter_torch(vocab_mask, [1], 0, is_allowed=True)
set_token_filter_torch(vocab_mask, [2], 1, is_allowed=True)
set_token_filter_torch(vocab_mask, [3], 2, is_allowed=True)
self.assertEqual(_get_allowed_tokens(vocab_mask, 0, 64), [1])
self.assertEqual(_get_allowed_tokens(vocab_mask, 1, 64), [2])
self.assertEqual(_get_allowed_tokens(vocab_mask, 2, 64), [3])
@unittest.skipUnless(_has_cuda, "CUDA not available")
class TestTritonTorchParity(unittest.TestCase):
"""Tests that Triton and Torch produce identical output."""
def _compare_outputs(self, token_ids, is_allowed, reset):
vocab_size = 128
num_elements = (vocab_size + 31) // 32
torch_mask = torch.zeros((1, num_elements), dtype=torch.int32)
triton_mask = torch.zeros((1, num_elements), dtype=torch.int32, device="cuda")
set_token_filter_torch(
torch_mask,
token_ids,
0,
is_allowed=is_allowed,
reset_vocab_mask=reset,
)
set_token_filter_triton(
triton_mask,
token_ids,
0,
is_allowed=is_allowed,
reset_vocab_mask=reset,
)
triton_cpu = triton_mask.cpu()
self.assertTrue(
torch.equal(torch_mask, triton_cpu),
f"Mismatch: torch={torch_mask} triton={triton_cpu}",
)
def test_parity_allow_tokens(self):
self._compare_outputs([0, 5, 31, 32, 63, 100], is_allowed=True, reset=True)
def test_parity_block_tokens(self):
self._compare_outputs([3, 5, 10], is_allowed=False, reset=True)
def test_parity_empty_tokens(self):
self._compare_outputs([], is_allowed=True, reset=True)
if __name__ == "__main__":
unittest.main()