chore: import upstream snapshot with attribution

python/sglang/srt/constrained/base_grammar_backend.py

@@ -0,0 +1,315 @@
+# Copyright 2023-2024 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.
+# ==============================================================================
+"""The baseclass of a backend for grammar-guided constrained decoding."""
+
+import logging
+import time
+from concurrent.futures import Future, ThreadPoolExecutor
+from dataclasses import dataclass, field
+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__)
+
+
+@dataclass
+class GrammarStats:
+    compilation_time: Optional[float] = None
+    schema_count: Optional[int] = None
+    ebnf_size: Optional[int] = None
+    is_cache_hit: bool = False
+    is_grammar_aborted: bool = False
+    tree_traversal_time: List[float] = field(default_factory=list)
+    dispatch_type: Optional[str] = None
+    num_timeout: int = 0
+
+
+class BaseGrammarObject:
+
+    def __init__(self):
+        self._finished = False
+        self.grammar_stats = None
+        self.current_token = None
+
+    def maybe_init_reasoning(self, reasoning: bool):
+        pass
+
+    def accept_token(self, token: int) -> None:
+        """
+        Accept a token in the grammar.
+        """
+        raise NotImplementedError()
+
+    def rollback(self, k: int):
+        raise NotImplementedError()
+
+    def is_terminated(self):
+        return False
+
+    def allocate_vocab_mask(
+        self, vocab_size: int, batch_size: int, device
+    ) -> torch.Tensor:
+        raise NotImplementedError()
+
+    def fill_vocab_mask(self, vocab_mask: torch.Tensor, idx: int) -> None:
+        raise NotImplementedError()
+
+    @staticmethod
+    def move_vocab_mask(vocab_mask: torch.Tensor, device) -> torch.Tensor:
+        raise NotImplementedError()
+
+    @staticmethod
+    def apply_vocab_mask(logits: torch.Tensor, vocab_mask: torch.Tensor) -> None:
+        raise NotImplementedError()
+
+    def copy(self) -> "BaseGrammarObject":
+        return self
+
+    @property
+    def finished(self):
+        return self._finished
+
+    @finished.setter
+    def finished(self, finished):
+        self._finished = finished
+
+    def try_jump_forward(self, tokenizer) -> Optional[Tuple[List[int], str]]:
+        """
+        Try to jump forward in the grammar.
+
+        Returns:
+            A jump forward helper which may be used in `jump_forward_str_state`.
+            None if the jump forward is not possible.
+        """
+        raise NotImplementedError()
+
+    def jump_forward_str_state(self, helper: Tuple[List[int], str]) -> Tuple[str, int]:
+        """
+        Jump forward for the grammar.
+
+        Returns:
+            A tuple of the jump forward string and the next state of the grammar
+            (which can be used in `jump_and_retokenize` if needed).
+        """
+        raise NotImplementedError()
+
+    def jump_and_retokenize(
+        self, old_output_ids: List[int], new_output_ids: List[int], next_state: int
+    ) -> None:
+        """
+        Jump forward occurs, and update the grammar state if needed.
+        """
+        raise NotImplementedError()
+
+
+class InvalidGrammarObject(BaseGrammarObject):
+    """Represents a grammar that failed to compile, carrying the original error message."""
+
+    def __init__(self, error_message: str = "Unknown grammar error"):
+        super().__init__()
+        self.error_message = error_message
+
+    def __repr__(self):
+        return f"InvalidGrammarObject(error_message={self.error_message!r})"
+
+
+class BaseGrammarBackend:
+    _enable_strict_thinking: bool = False
+
+    def __init__(self):
+        self.executor = ThreadPoolExecutor()
+        self.cache: Dict[Tuple[str, str], BaseGrammarObject] = {}
+
+    def _not_supported(self, key_type: str, key_string: str) -> BaseGrammarObject:
+        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.
+        """
+        raise ValueError(f"Invalid key_type: {key_type}={key_string}")
+
+    def dispatch_json(self, key_string: str) -> BaseGrammarObject:
+        return self._not_supported("json", key_string)
+
+    def dispatch_regex(self, key_string: str) -> BaseGrammarObject:
+        return self._not_supported("regex", key_string)
+
+    def dispatch_ebnf(self, key_string: str) -> BaseGrammarObject:
+        return self._not_supported("ebnf", key_string)
+
+    def dispatch_structural_tag(self, key_string: str) -> BaseGrammarObject:
+        return self._not_supported("structural_tag", key_string)
+
+    def _init_value_dispatch(
+        self, key: Tuple[str, str], require_reasoning: bool
+    ) -> BaseGrammarObject:
+        s = time.perf_counter()
+        key_type, key_string = key
+        if key_type == "json":
+            grammar = self.dispatch_json(key_string)
+        elif key_type == "regex":
+            grammar = self.dispatch_regex(key_string)
+        elif key_type == "ebnf":
+            grammar = self.dispatch_ebnf(key_string)
+        elif key_type == "structural_tag":
+            grammar = self.dispatch_structural_tag(key_string)
+        else:
+            grammar = self.dispatch_fallback(key_type, key_string)
+
+        if grammar is not None and grammar.grammar_stats is not None:
+            grammar.grammar_stats.compilation_time = time.perf_counter() - s
+        return grammar
+
+    def get_cached_or_future_value(
+        self, key: Tuple[str, str], require_reasoning: bool
+    ) -> Tuple[BaseGrammarObject | Future[BaseGrammarObject], bool]:
+        value = self.cache.get(key)
+        if value:
+            copied_value = value.copy()
+            copied_value.maybe_init_reasoning(require_reasoning)
+            return copied_value, True
+        value = self.executor.submit(self._init_value_dispatch, key, require_reasoning)
+        return value, False
+
+    def set_cache(self, key: Tuple[str, str], value: BaseGrammarObject):
+        self.cache[key] = value
+
+    def reset(self):
+        self.cache.clear()
+
+
+GRAMMAR_BACKEND_REGISTRY = {}
+
+
+def register_grammar_backend(name, init_func):
+    GRAMMAR_BACKEND_REGISTRY[name] = init_func
+
+
+def create_grammar_backend(
+    server_args: ServerArgs,
+    tokenizer,
+    vocab_size: int,
+    eos_token_ids: Optional[set] = None,
+    think_end_id: Optional[int] = None,
+) -> Optional[BaseGrammarBackend]:
+    name = server_args.grammar_backend
+
+    # Custom grammar backend has the highest priority
+    if name in GRAMMAR_BACKEND_REGISTRY:
+        return GRAMMAR_BACKEND_REGISTRY[name](
+            server_args, tokenizer, vocab_size, eos_token_ids
+        )
+
+    # Default grammar backends
+    if name == "outlines":
+        from sglang.srt.constrained.outlines_backend import OutlinesGrammarBackend
+
+        grammar_backend = OutlinesGrammarBackend(
+            tokenizer,
+            whitespace_pattern=server_args.constrained_json_whitespace_pattern,
+        )
+    elif name == "xgrammar":
+        from sglang.srt.constrained.xgrammar_backend import (
+            TokenizerNotSupportedError,
+            XGrammarGrammarBackend,
+        )
+
+        # Convert Set[int] to List[int] if needed
+        eos_list = list(eos_token_ids) if eos_token_ids else None
+
+        try:
+            grammar_backend = XGrammarGrammarBackend(
+                tokenizer,
+                vocab_size=vocab_size,
+                model_eos_token_ids=eos_list,
+                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'. "
+                "Structured outputs (JSON schema, regex, EBNF) will not be available."
+            )
+            server_args.override("grammar.import_fallback", grammar_backend="none")
+            return None
+    elif name == "llguidance":
+        from sglang.srt.constrained.llguidance_backend import GuidanceBackend
+
+        grammar_backend = GuidanceBackend(
+            tokenizer=tokenizer,
+            any_whitespace=not server_args.constrained_json_disable_any_whitespace,
+            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}")
+
+    if server_args.reasoning_parser and think_end_id is not None:
+        from sglang.srt.constrained.reasoner_grammar_backend import (
+            ReasonerGrammarBackend,
+        )
+
+        reasoning_parser = ReasoningParser(
+            model_type=server_args.reasoning_parser,
+            stream_reasoning=False,
+            tokenizer=tokenizer,
+        )
+
+        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

@@ -0,0 +1,311 @@
+from __future__ import annotations
+
+import logging
+import time
+from concurrent import futures
+from typing import TYPE_CHECKING, List
+
+import torch
+
+from sglang.srt.constrained.base_grammar_backend import (
+    InvalidGrammarObject,
+    create_grammar_backend,
+)
+from sglang.srt.constrained.reasoner_grammar_backend import ReasonerGrammarObject
+from sglang.srt.distributed.communication_tags import P2PTag
+from sglang.srt.environ import envs
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.io_struct import AbortReq
+    from sglang.srt.managers.schedule_batch import Req
+    from sglang.srt.managers.scheduler import Scheduler
+
+logger = logging.getLogger(__name__)
+
+
+class GrammarManager:
+    def __init__(self, scheduler: Scheduler):
+        self.scheduler = scheduler
+        self.server_args = scheduler.server_args
+        self.grammar_queue: List[Req] = []
+        if not self.server_args.skip_tokenizer_init:
+            self.grammar_backend = create_grammar_backend(
+                self.server_args,
+                scheduler.tokenizer,
+                scheduler.model_config.vocab_size,
+                scheduler.model_config.hf_eos_token_id,
+                think_end_id=scheduler.model_config.think_end_id,
+            )
+        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
+        self.is_grammar_sync_entry = scheduler.dp_tp_group.is_first_rank
+        self.pp_rank = scheduler.ps.pp_rank
+        self.pp_size = scheduler.ps.pp_size
+        self.pp_group = scheduler.pp_group
+        self.grammar_pp_sync_work_list = []
+
+        self.SGLANG_GRAMMAR_POLL_INTERVAL = envs.SGLANG_GRAMMAR_POLL_INTERVAL.get()
+        self.SGLANG_GRAMMAR_MAX_POLL_ITERATIONS = (
+            envs.SGLANG_GRAMMAR_MAX_POLL_ITERATIONS.get()
+        )
+
+    def __len__(self):
+        return len(self.grammar_queue)
+
+    def clear(self):
+        if self.grammar_backend:
+            self.grammar_backend.reset()
+
+    def has_waiting_grammars(self) -> bool:
+        return len(self.grammar_queue) > 0
+
+    def _drain_pp_sync_work(self):
+        for p2p_work in self.grammar_pp_sync_work_list:
+            p2p_work.work.wait()
+        self.grammar_pp_sync_work_list.clear()
+
+    def _pp_sync_ready_failed(
+        self,
+        ready_req_idxs: set[int],
+        failed_req_idxs: set[int],
+    ) -> tuple[set[int], set[int]]:
+        """
+        Synchronize ready/failed grammar request indexes across the PP pipeline.
+
+        PP0 provides the data. Each later PP rank receives it from the previous
+        rank and asynchronously forwards it to the next rank.
+        """
+        if self.pp_size <= 1 or self.pp_group is None:
+            return ready_req_idxs, failed_req_idxs
+
+        self._drain_pp_sync_work()
+        data = (ready_req_idxs, failed_req_idxs)
+        if self.pp_rank > 0:
+            data = self.pp_group.recv_object(
+                src=self.pp_rank - 1,
+                tag=P2PTag.GRAMMAR_PP_SYNC,
+            )
+        if self.pp_rank + 1 < self.pp_size:
+            self.grammar_pp_sync_work_list.extend(
+                self.pp_group.send_object(
+                    data,
+                    dst=self.pp_rank + 1,
+                    async_send=True,
+                    tag=P2PTag.GRAMMAR_PP_SYNC,
+                )
+            )
+        return data
+
+    def abort_requests(self, recv_req: AbortReq):
+        for req in self.grammar_queue:
+            if recv_req.abort_all or req.rid.startswith(recv_req.rid):
+                logger.debug(f"Abort grammar queue request. {req.rid=}")
+                if isinstance(req.grammar, futures.Future) and req.grammar:
+                    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
+        if (
+            req.sampling_params.json_schema is not None
+            or req.sampling_params.regex is not None
+            or req.sampling_params.ebnf is not None
+            or req.sampling_params.structural_tag is not None
+        ):
+            if self.grammar_backend is None:
+                error_msg = "Grammar-based generation (json_schema, regex, ebnf, structural_tag) is not supported when the server is launched with --grammar-backend none"
+                req.set_finish_with_abort(error_msg)
+            else:
+                if req.sampling_params.json_schema is not None:
+                    key = ("json", req.sampling_params.json_schema)
+                elif req.sampling_params.regex is not None:
+                    key = ("regex", req.sampling_params.regex)
+                elif req.sampling_params.ebnf is not None:
+                    key = ("ebnf", req.sampling_params.ebnf)
+                elif req.sampling_params.structural_tag:
+                    key = ("structural_tag", req.sampling_params.structural_tag)
+
+                value, cache_hit = self.grammar_backend.get_cached_or_future_value(
+                    key, req.require_reasoning
+                )
+                req.grammar = value
+
+                if not cache_hit:
+                    req.grammar_key = key
+                    add_to_grammar_queue = True
+                else:
+                    if isinstance(
+                        value, InvalidGrammarObject
+                    ):  # We hit a cached invalid grammar.
+                        error_msg = (
+                            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)
+
+        return add_to_grammar_queue
+
+    def get_ready_grammar_requests(self) -> List[Req]:
+        """
+        Move requests whose grammar objects are ready from grammar_queue to waiting_queue.
+
+        For PP0, DP/TP group rank i returns two sets ready_reqs_i,
+        failed_reqs_i. ready_reqs_all = all_gather(ready_reqs_i) within
+        PP0's DP/TP group. failed_reqs_all = all_gather(failed_reqs_i)
+        within PP0's DP/TP group.
+
+        ready_reqs = intersect(ready_reqs_all)
+        failed_reqs = union(failed_reqs_all)
+
+        PP0 then propagates the synced result to later PP ranks. Later PP
+        ranks receive and apply the propagated ready/failed decision.
+        """
+        assert self.grammar_backend
+        ready_req_idxs: set[int] = set()
+        failed_req_idxs: set[int] = set()
+
+        if self.pp_rank == 0:
+            # Poll for ready requests
+            start_time = time.perf_counter()
+            while time.perf_counter() - start_time < self.SGLANG_GRAMMAR_POLL_INTERVAL:
+                for i, req in enumerate(self.grammar_queue):
+                    if i in ready_req_idxs:
+                        continue
+
+                    if (
+                        req.finished() or req.grammar is None
+                    ):  # It is aborted by AbortReq
+                        ready_req_idxs.add(i)
+                        continue
+
+                    assert isinstance(req.grammar, futures.Future), f"{req=}"
+                    if req.grammar.done():
+                        ready_req_idxs.add(i)
+
+                if len(ready_req_idxs) == len(self.grammar_queue):
+                    break
+
+                # Sleep a bit to avoid busy waiting
+                time.sleep(self.SGLANG_GRAMMAR_POLL_INTERVAL / 10)
+
+            # Check failed requests
+            for i, req in enumerate(self.grammar_queue):
+                if i not in ready_req_idxs:
+                    # grammar_wait_ct is only updated on PP0; later PP ranks
+                    # receive PP0's ready/failed decision through PP sync.
+                    self.grammar_queue[i].grammar_wait_ct += 1
+                    if (
+                        self.grammar_queue[i].grammar_wait_ct
+                        >= self.SGLANG_GRAMMAR_MAX_POLL_ITERATIONS
+                    ):
+                        # Timeout after max poll iterations
+                        # The actual waiting time is SGLANG_GRAMMAR_MAX_POLL_ITERATIONS * max(SGLANG_GRAMMAR_POLL_INTERVAL, GPU_forward_batch_latency)
+                        failed_req_idxs.add(i)
+
+            # Sync ready and failed requests across all TP ranks in PP0.
+            if self.grammar_sync_size == 1:
+                synced_ready_req_idxs = ready_req_idxs
+                synced_failed_req_idxs = failed_req_idxs
+            else:
+                all_gather_output = [None] * self.grammar_sync_size
+                torch.distributed.all_gather_object(
+                    all_gather_output,
+                    (ready_req_idxs, failed_req_idxs),
+                    group=self.grammar_sync_group,
+                )
+                synced_ready_req_idxs = set.intersection(
+                    *[x[0] for x in all_gather_output]
+                )
+                synced_failed_req_idxs = set.union(*[x[1] for x in all_gather_output])
+        else:
+            synced_ready_req_idxs = ready_req_idxs
+            synced_failed_req_idxs = failed_req_idxs
+
+        # Propagate PP0's grammar queue decision to later PP ranks.
+        (
+            synced_ready_req_idxs,
+            synced_failed_req_idxs,
+        ) = self._pp_sync_ready_failed(
+            synced_ready_req_idxs,
+            synced_failed_req_idxs,
+        )
+
+        # Return ready requests
+        return_reqs: List[Req] = []
+        for i in synced_ready_req_idxs:
+            req = self.grammar_queue[i]
+            return_reqs.append(req)
+            if req.finished() or req.grammar is None:  # It is aborted by AbortReq
+                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)
+
+        # Return failed requests
+        for i in synced_failed_req_idxs:
+            req = self.grammar_queue[i]
+            return_reqs.append(req)
+
+            assert isinstance(req.grammar, futures.Future) and req.grammar_key
+            req.grammar.cancel()
+            self.grammar_backend.set_cache(
+                req.grammar_key, InvalidGrammarObject("Grammar preprocessing timed out")
+            )
+            error_msg = f"Grammar preprocessing timed out: {req.grammar_key=}"
+            req.set_finish_with_abort(error_msg)
+
+        # Remove finished requests from grammar_queue
+        self.grammar_queue = [
+            req
+            for i, req in enumerate(self.grammar_queue)
+            if i not in synced_ready_req_idxs and i not in synced_failed_req_idxs
+        ]
+        return return_reqs

python/sglang/srt/constrained/llguidance_backend.py

@@ -0,0 +1,190 @@
+# Copyright 2023-2024 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.
+# ==============================================================================
+"""Constrained decoding with llguidance backend."""
+
+import json
+import logging
+import os
+from typing import List, Optional, Tuple
+
+import torch
+from llguidance import LLMatcher, LLTokenizer, StructTag, grammar_from
+from llguidance.hf import from_tokenizer
+from llguidance.torch import (
+    allocate_token_bitmask,
+    apply_token_bitmask_inplace,
+    fill_next_token_bitmask,
+)
+
+from sglang.srt.constrained.base_grammar_backend import (
+    BaseGrammarBackend,
+    BaseGrammarObject,
+    InvalidGrammarObject,
+)
+from sglang.srt.constrained.utils import is_legacy_structural_tag
+
+logger = logging.getLogger(__name__)
+
+
+class GuidanceGrammar(BaseGrammarObject):
+
+    def __init__(self, llguidance_tokenizer: LLTokenizer, serialized_grammar: str):
+        super().__init__()
+        self.llguidance_tokenizer = llguidance_tokenizer
+        self.serialized_grammar = serialized_grammar
+
+        self.ll_matcher = LLMatcher(
+            self.llguidance_tokenizer,
+            self.serialized_grammar,
+            log_level=int(os.environ.get("LLGUIDANCE_LOG_LEVEL", "1")),
+        )
+        self._check_err()
+
+        self.eos_token = self.llguidance_tokenizer.eos_token
+
+    def accept_token(self, token: int):
+        if self.finished:
+            return
+        if self.ll_matcher.is_stopped() and token == self.eos_token:
+            self.finished = True
+            return
+        self.ll_matcher.consume_token(token)
+        self._check_err()
+
+    def rollback(self, num_tokens: int) -> None:
+        if num_tokens <= 0:
+            return
+        if self.finished:
+            self.finished = False
+            # EOS token after stop isn't tracked in ll_matcher
+            num_tokens -= 1
+        self.ll_matcher.rollback(num_tokens)
+        self._check_err()
+
+    def is_terminated(self):
+        return self.finished
+
+    def fill_vocab_mask(self, vocab_mask: torch.Tensor, idx: int) -> None:
+        fill_next_token_bitmask(self.ll_matcher, vocab_mask, idx)
+        self._check_err()
+
+    def allocate_vocab_mask(
+        self, vocab_size: int, batch_size: int, device
+    ) -> torch.Tensor:
+        return allocate_token_bitmask(batch_size, self.llguidance_tokenizer.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:
+        apply_token_bitmask_inplace(logits, vocab_mask)
+
+    def copy(self):
+        return GuidanceGrammar(
+            llguidance_tokenizer=self.llguidance_tokenizer,
+            serialized_grammar=self.serialized_grammar,
+        )
+
+    def try_jump_forward(self, tokenizer) -> Optional[Tuple[List[int], str]]:
+        ff_tokens = self.ll_matcher.compute_ff_tokens()
+        if ff_tokens:
+            return ff_tokens, ""
+        else:
+            return None
+
+    def jump_forward_str_state(self, helper: Tuple[List[int], str]) -> Tuple[str, int]:
+        return "", -1
+
+    def jump_and_retokenize(
+        self, old_output_ids: List[int], new_output_ids: List[int], next_state: int
+    ):
+        pass
+
+    def _check_err(self) -> None:
+        if self.ll_matcher.is_error():
+            raise ValueError(self.ll_matcher.get_error())
+
+
+class GuidanceBackend(BaseGrammarBackend):
+
+    def __init__(
+        self,
+        tokenizer,
+        any_whitespace: bool = True,
+        whitespace_pattern: Optional[str] = None,
+        n_vocab: Optional[int] = None,
+    ):
+        super().__init__()
+
+        self.tokenizer = tokenizer
+        self.any_whitespace = any_whitespace
+        self.whitespace_pattern = whitespace_pattern
+        self.llguidance_tokenizer = from_tokenizer(self.tokenizer, n_vocab)
+
+    def _from_serialized(self, serialized_grammar) -> BaseGrammarObject:
+        try:
+            return GuidanceGrammar(
+                llguidance_tokenizer=self.llguidance_tokenizer,
+                serialized_grammar=serialized_grammar,
+            )
+        except Exception as e:
+            logger.error(f"Hit invalid grammar: {serialized_grammar=}, {e=}")
+            return InvalidGrammarObject(str(e))
+
+    def dispatch_json(self, key_string: str) -> BaseGrammarObject:
+        try:
+            serialized_grammar = LLMatcher.grammar_from_json_schema(
+                key_string,
+                defaults={
+                    "whitespace_flexible": self.any_whitespace,
+                    "whitespace_pattern": self.whitespace_pattern,
+                },
+            )
+        except Exception as e:
+            logger.error(f"Hit invalid json_schema: {key_string=}, {e=}")
+            return InvalidGrammarObject(str(e))
+        return self._from_serialized(serialized_grammar)
+
+    def dispatch_regex(self, key_string: str) -> BaseGrammarObject:
+        serialized_grammar = grammar_from("regex", key_string)
+        return self._from_serialized(serialized_grammar)
+
+    def dispatch_ebnf(self, key_string: str) -> BaseGrammarObject:
+        try:
+            serialized_grammar = grammar_from("ebnf", key_string)
+            return self._from_serialized(serialized_grammar)
+        except ValueError as e:
+            logger.error(f"Hit invalid ebnf: {key_string=}, {e=}")
+            return InvalidGrammarObject(str(e))
+
+    def dispatch_structural_tag(self, key_string: str) -> BaseGrammarObject:
+        try:
+            structural_tag = json.loads(key_string)
+            assert is_legacy_structural_tag(structural_tag)
+            tags = [
+                StructTag(
+                    begin=structure["begin"],
+                    grammar=structure["schema"],
+                    end=structure["end"],
+                    trigger=structural_tag["triggers"][0],  # TODO?
+                )
+                for structure in structural_tag["structures"]
+            ]
+            g = StructTag.to_grammar(tags)
+            return self._from_serialized(g)
+        except Exception as e:
+            logger.error(f"Hit invalid structural_tag: {key_string=}, {e=}")
+            return InvalidGrammarObject(str(e))

python/sglang/srt/constrained/outlines_backend.py

@@ -0,0 +1,190 @@
+# Copyright 2023-2024 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.
+# ==============================================================================
+"""Constrained decoding with outlines backend."""
+
+import json
+import logging
+from typing import Dict, List, Optional, Tuple, Union
+
+import interegular
+import torch
+from outlines.fsm.guide import RegexGuide
+from outlines.models.transformers import TransformerTokenizer
+from pydantic import BaseModel
+
+from sglang.srt.constrained.base_grammar_backend import (
+    BaseGrammarBackend,
+    BaseGrammarObject,
+    InvalidGrammarObject,
+)
+from sglang.srt.constrained.outlines_jump_forward import OutlinesJumpForwardMap
+
+try:
+    from outlines.fsm.json_schema import build_regex_from_schema
+except ImportError:
+    from outlines_core.fsm.json_schema import build_regex_from_schema
+
+
+logger = logging.getLogger(__name__)
+
+
+class OutlinesGrammar(BaseGrammarObject):
+    def __init__(
+        self,
+        guide: RegexGuide,
+        jump_forward_map: Union[OutlinesJumpForwardMap, None],
+    ) -> None:
+        super().__init__()
+        self.guide = guide
+        self.jump_forward_map = jump_forward_map
+        self.state = 0
+
+    def accept_token(self, token: int):
+        self.state = self.guide.get_next_state(self.state, token)
+
+    def allocate_vocab_mask(
+        self, vocab_size: int, batch_size: int, device
+    ) -> torch.Tensor:
+        return torch.zeros(batch_size, vocab_size, dtype=torch.bool, device=device)
+
+    @staticmethod
+    def move_vocab_mask(vocab_mask: torch.Tensor, device) -> torch.Tensor:
+        return vocab_mask
+
+    def fill_vocab_mask(self, vocab_mask: torch.Tensor, idx: int) -> None:
+        tokens = torch.tensor(
+            self.guide.get_next_instruction(self.state).tokens, dtype=torch.int64
+        ).to(vocab_mask.device, non_blocking=True)
+        vocab_mask = vocab_mask[idx]
+        vocab_mask.fill_(1)
+        vocab_mask.scatter_(0, tokens, torch.zeros_like(tokens, dtype=torch.bool))
+
+    @staticmethod
+    def apply_vocab_mask(logits: torch.Tensor, vocab_mask: torch.Tensor):
+        logits.masked_fill_(vocab_mask, float("-inf"))
+
+    def copy(self):
+        return OutlinesGrammar(self.guide, self.jump_forward_map)
+
+    def try_jump_forward(self, tokenizer) -> Optional[Tuple]:
+        if not self.jump_forward_map:
+            return None
+
+        jump_forward_bytes = self.jump_forward_map.jump_forward_byte(self.state)
+        if jump_forward_bytes is None or len(jump_forward_bytes) <= 1:
+            return None
+
+        # preprocess the jump forward string
+        suffix_bytes = []
+        continuation_range = range(0x80, 0xC0)
+        cur_state = self.state
+        while (
+            len(jump_forward_bytes) and jump_forward_bytes[0][0] in continuation_range
+        ):
+            # continuation bytes
+            byte_edge = jump_forward_bytes.pop(0)
+            suffix_bytes.append(byte_edge[0])
+            cur_state = byte_edge[1]
+
+        suffix_tokens = [f"<0x{hex(b)[2:].upper()}>" for b in suffix_bytes]
+        suffix_ids = tokenizer.convert_tokens_to_ids(suffix_tokens)
+        return suffix_ids, cur_state
+
+    def jump_forward_str_state(self, helper: Tuple[List[int], str]) -> Tuple[str, int]:
+        _, cur_state = helper
+        return self.jump_forward_map.jump_forward_symbol(cur_state)
+
+    def jump_and_retokenize(
+        self, old_output_ids: List[int], new_output_ids: List[int], next_state: int
+    ):
+        self.state = next_state
+
+
+class OutlinesGrammarBackend(BaseGrammarBackend):
+    def __init__(
+        self,
+        tokenizer,
+        whitespace_pattern: str | None,
+    ):
+        super().__init__()
+
+        try:
+            self.outlines_tokenizer = TransformerTokenizer(tokenizer)
+        except AttributeError:
+            # FIXME: tmp fix for chatglm2 & chatglm3 (pad_token_id=0)
+            origin_pad_token_id = tokenizer.pad_token_id
+
+            def fset(self, value):
+                self._value = value
+
+            type(tokenizer).pad_token_id = property(
+                fget=type(tokenizer).pad_token_id.fget, fset=fset
+            )
+            self.outlines_tokenizer = TransformerTokenizer(tokenizer)
+            self.outlines_tokenizer.tokenizer.pad_token_id = origin_pad_token_id
+            self.outlines_tokenizer.pad_token_id = origin_pad_token_id
+            self.outlines_tokenizer.pad_token = (
+                self.outlines_tokenizer.tokenizer.pad_token
+            )
+            self.outlines_tokenizer.vocabulary = (
+                self.outlines_tokenizer.tokenizer.get_vocab()
+            )
+        self.whitespace_pattern = whitespace_pattern
+
+    def _compile_regex(self, regex: str) -> BaseGrammarObject:
+        try:
+            if hasattr(RegexGuide, "from_regex"):
+                # outlines >= 0.1.1
+                guide = RegexGuide.from_regex(regex, self.outlines_tokenizer)
+            else:
+                # outlines <= 0.0.46
+                guide = RegexGuide(regex, self.outlines_tokenizer)
+        except interegular.patterns.InvalidSyntax as e:
+            logger.error(f"Hit invalid regex schema: {regex=}, {e=}")
+            return InvalidGrammarObject(str(e))
+
+        jump_forward_map = None
+        return OutlinesGrammar(guide, jump_forward_map)
+
+    def dispatch_ebnf(self, key_string: str):
+        return super().dispatch_ebnf(key_string)
+
+    def dispatch_structural_tag(self, key_string: str):
+        return super().dispatch_structural_tag(key_string)
+
+    def dispatch_json(self, key_string: str):
+        try:
+            regex = build_regex_from_object(
+                key_string,
+                whitespace_pattern=self.whitespace_pattern,
+            )
+        except (NotImplementedError, json.decoder.JSONDecodeError, ValueError) as e:
+            logger.error(f"Hit invalid json_schema: {key_string=}, {e=}")
+            return InvalidGrammarObject(str(e))
+        return self._compile_regex(regex)
+
+    def dispatch_regex(self, key_string: str):
+        return self._compile_regex(key_string)
+
+
+def build_regex_from_object(
+    object: Union[str, BaseModel, Dict], whitespace_pattern: Optional[str] = None
+):
+    if isinstance(object, type(BaseModel)):
+        schema = json.dumps(object.model_json_schema())
+    elif isinstance(object, Dict):
+        schema = json.dumps(object)
+    else:
+        schema = object
+    return build_regex_from_schema(schema, whitespace_pattern)

python/sglang/srt/constrained/outlines_jump_forward.py

@@ -0,0 +1,200 @@
+# Copyright 2023-2024 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.
+# ==============================================================================
+"""
+Faster constrained decoding with jump forward decoding / compressed finite state machine.
+Reference: https://lmsys.org/blog/2024-02-05-compressed-fsm/
+"""
+
+import dataclasses
+import logging
+from collections import defaultdict
+from typing import Optional
+
+import interegular
+from interegular import InvalidSyntax
+from outlines.caching import cache
+
+from sglang.srt.utils import get_bool_env_var
+
+try:
+    # outlines >= 0.1.0
+    from outlines_core.fsm.outlines_core_rs import FSMInfo
+    from outlines_core.fsm.regex import make_byte_level_fsm, make_deterministic_fsm
+except ImportError:
+    # outlines <= 0.0.46
+    from outlines.fsm.regex import FSMInfo, make_byte_level_fsm, make_deterministic_fsm
+
+IP_REGEX = r"((25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}(25[0-5]|2[0-4]\d|[01]?\d\d?)"
+
+# Env var was set in sglang.srt.server_args.ServerArgs.__post_init__
+DISABLE_DISK_CACHE = get_bool_env_var("SGLANG_DISABLE_OUTLINES_DISK_CACHE", "true")
+
+logger = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass
+class JumpEdge:
+    symbol: str = None
+    symbol_next_state: int = None
+    byte: int = None
+    byte_next_state: int = None
+
+
+def disk_cache(expire: Optional[float] = None, typed=False, ignore=()):
+    if not DISABLE_DISK_CACHE:
+        return cache(expire, typed, ignore)
+    else:
+        return lambda fn: None
+
+
+@disk_cache()
+def init_state_to_jump_forward(regex_string):
+    try:
+        regex_pattern = interegular.parse_pattern(regex_string)
+    except InvalidSyntax as e:
+        logger.warning(f"skip invalid regex: {regex_string}, {e=}")
+        return
+
+    byte_fsm = make_byte_level_fsm(regex_pattern.to_fsm().reduce(), keep_utf8=True)
+    regex_fsm, _ = make_deterministic_fsm(byte_fsm)
+
+    fsm_info: FSMInfo = regex_fsm.fsm_info
+
+    symbol_to_id = fsm_info.alphabet_symbol_mapping
+    id_to_symbol = {}
+    for symbol, id_ in symbol_to_id.items():
+        id_to_symbol.setdefault(id_, []).append(symbol)
+
+    transitions = fsm_info.transitions
+
+    outgoings_ct = defaultdict(int)
+    # NOTE(lsyin): Final states can lead to terminate, so they have one outgoing edge naturally
+    for s in fsm_info.finals:
+        outgoings_ct[s] = 1
+
+    state_to_jump_forward = {}
+    for (state, id_), next_state in transitions.items():
+        if id_ == fsm_info.alphabet_anything_value:
+            # Arbitrarily symbol cannot be recognized as jump forward
+            continue
+
+        symbols = id_to_symbol[id_]
+        for c in symbols:
+            if len(c) > 1:
+                # Skip byte level transitions like c = "5E"
+                continue
+
+            outgoings_ct[state] += 1
+            if outgoings_ct[state] > 1:
+                if state in state_to_jump_forward:
+                    del state_to_jump_forward[state]
+                break
+
+            state_to_jump_forward[state] = JumpEdge(
+                symbol=c,
+                symbol_next_state=next_state,
+            )
+
+    # Process the byte level jump forward
+    outgoings_ct = defaultdict(int)
+    for s in fsm_info.finals:
+        outgoings_ct[s] = 1
+
+    for (state, id_), next_state in transitions.items():
+        if id_ == fsm_info.alphabet_anything_value:
+            continue
+        symbols = id_to_symbol[id_]
+        for c in symbols:
+            byte_ = None
+            if len(c) == 1 and ord(c) < 0x80:
+                # ASCII character
+                byte_ = ord(c)
+            elif len(c) > 1:
+                # FIXME: This logic is due to the leading \x00
+                # https://github.com/outlines-dev/outlines/pull/930
+                byte_ = int(symbols[0][1:], 16)
+
+            if byte_ is not None:
+                outgoings_ct[state] += 1
+                if outgoings_ct[state] > 1:
+                    if state in state_to_jump_forward:
+                        del state_to_jump_forward[state]
+                    break
+                e = state_to_jump_forward.get(state, JumpEdge())
+                e.byte = byte_
+                e.byte_next_state = next_state
+                state_to_jump_forward[state] = e
+
+    return state_to_jump_forward
+
+
+class OutlinesJumpForwardMap:
+    def __init__(self, regex_string):
+        self.state_to_jump_forward = init_state_to_jump_forward(regex_string)
+
+    def jump_forward_symbol(self, state):
+        jump_forward_str = ""
+        next_state = state
+        while state in self.state_to_jump_forward:
+            e = self.state_to_jump_forward[state]
+            if e.symbol is None:
+                break
+            jump_forward_str += e.symbol
+            next_state = e.symbol_next_state
+            state = next_state
+
+        return jump_forward_str, next_state
+
+    def jump_forward_byte(self, state):
+        if state not in self.state_to_jump_forward:
+            return None
+
+        jump_forward_bytes = []
+        next_state = None
+        while state in self.state_to_jump_forward:
+            e = self.state_to_jump_forward[state]
+            assert e.byte is not None and e.byte_next_state is not None
+            jump_forward_bytes.append((e.byte, e.byte_next_state))
+            next_state = e.byte_next_state
+            state = next_state
+
+        return jump_forward_bytes
+
+    def is_jump_forward_symbol_state(self, state):
+        return (
+            state in self.state_to_jump_forward
+            and self.state_to_jump_forward[state].symbol is not None
+        )
+
+
+def test_main(regex_string):
+    jump_forward_map = OutlinesJumpForwardMap(regex_string)
+    for state, e in jump_forward_map.state_to_jump_forward.items():
+        if e.symbol is not None:
+            jump_forward_str, next_state = jump_forward_map.jump_forward_symbol(state)
+            print(f"{state} -> {next_state}", jump_forward_str)
+        bytes_ = jump_forward_map.jump_forward_byte(state)
+        print(f"{state} -> {bytes_[-1][1]}", [hex(b) for b, _ in bytes_])
+
+
+if __name__ == "__main__":
+    import outlines
+
+    outlines.caching.clear_cache()
+    test_main(r"The google's DNS sever address is " + IP_REGEX)
+    test_main(r"霍格沃茨特快列车|霍比特人比尔博")
+    # 霍格: \xe9\x9c\x8d \xe6\xa0\xbc ...
+    # 霍比: \xe9\x9c\x8d \xe6\xaf\x94 ...
+
+    test_main(r"[-+]?[0-9]+[ ]*")

python/sglang/srt/constrained/reasoner_grammar_backend.py

@@ -0,0 +1,327 @@
+# Copyright 2023-2024 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.
+# ==============================================================================
+"""The baseclass of a backend for reasoner grammar-guided constrained decoding."""
+
+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,
+    BaseGrammarObject,
+    InvalidGrammarObject,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class ReasonerGrammarObject(BaseGrammarObject):
+    """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
+        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):
+        if reasoning:
+            self.tokens_in_think = 0
+            self.tokens_after_end = -1
+        else:
+            self.tokens_in_think = -1
+            self.tokens_after_end = 0
+
+    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._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):
+        # Track the last accepted token on the wrapper itself (mirroring
+        # XGrammarGrammar.accept_token). Disaggregation's process_prebuilt uses
+        # `grammar.current_token is None` to detect a retracted request whose
+        # token was already accepted and must not be re-accepted. Without this,
+        # a ReasonerGrammarObject's current_token stays None forever (the inner
+        # grammar's is updated, not the wrapper's), so the guard never fires and
+        # the token is accepted twice -> "Tokens not accepted" -> FINISH_ABORT.
+        self.current_token = token
+        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):
+        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 _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._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 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):
+        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, 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,
+        reasoning_parser: ReasoningParser,
+        tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
+        enable_strict_thinking: bool = False,
+    ):
+        super().__init__()
+        self.grammar_backend = grammar_backend
+        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)
+        if ret is None or isinstance(ret, InvalidGrammarObject):
+            return ret
+        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/utils.py

@@ -0,0 +1,12 @@
+from typing import Dict
+
+
+def is_legacy_structural_tag(obj: Dict) -> bool:
+    # test whether an object is a legacy structural tag
+    # see `StructuralTagResponseFormat` at `sglang.srt.entrypoints.openai.protocol`
+    if obj.get("structures", None) is not None:
+        assert obj.get("triggers", None) is not None
+        return True
+    else:
+        assert obj.get("format", None) is not None
+        return False

python/sglang/srt/constrained/xgrammar_backend.py

@@ -0,0 +1,419 @@
+# Copyright 2023-2024 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.
+# ==============================================================================
+"""Constrained decoding with xgrammar backend."""
+
+import dataclasses
+import json
+import logging
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+from xgrammar import (
+    CompiledGrammar,
+    GrammarCompiler,
+    GrammarMatcher,
+    StructuralTag,
+    StructuralTagItem,
+    TokenizerInfo,
+    allocate_token_bitmask,
+)
+
+from sglang.srt.constrained.base_grammar_backend import (
+    BaseGrammarBackend,
+    BaseGrammarObject,
+    GrammarStats,
+    InvalidGrammarObject,
+)
+from sglang.srt.constrained.utils import is_legacy_structural_tag
+from sglang.srt.utils import is_hip
+
+_is_hip = is_hip()
+
+if _is_hip:
+    from sgl_kernel import apply_token_bitmask_inplace_cuda
+else:
+    from sglang.kernels.ops.grammar.bitmask_ops import (
+        apply_token_bitmask_inplace_triton,
+    )
+
+from sglang.kernels.ops.grammar.token_filter_ops import set_token_filter_triton
+from sglang.srt.constrained.torch_ops.token_filter_torch_ops import (
+    set_token_filter_torch,
+)
+
+logger = logging.getLogger(__name__)
+MAX_ROLLBACK_TOKENS = 200
+
+
+class XGrammarGrammar(BaseGrammarObject):
+
+    def __init__(
+        self,
+        matcher: GrammarMatcher,
+        vocab_size: int,
+        ctx: CompiledGrammar,
+        override_stop_tokens: Optional[Union[List[int], int]],
+        key_string: Optional[str] = None,
+        grammar_stats: Optional[GrammarStats] = GrammarStats(),
+    ) -> None:
+        super().__init__()
+        self.matcher = matcher
+        self.vocab_size = vocab_size
+        self.ctx = ctx
+        self.override_stop_tokens = override_stop_tokens
+        self.accepted_tokens = []
+        self.key_string = key_string
+        self.grammar_stats = grammar_stats
+
+    def accept_token(self, token: int):
+        if not self.is_terminated():
+            self.current_token = token
+            accepted = self.matcher.accept_token(token)
+            if not accepted:
+                # log for debugging
+                raise ValueError(
+                    f"Tokens not accepted: {token}\n"
+                    f"Accepted tokens: {self.accepted_tokens}\n"
+                    f"Key string: {self.key_string}"
+                )
+            else:
+                self.accepted_tokens.append(token)
+
+    def rollback(self, k: int):
+        self.matcher.rollback(k)
+        self.accepted_tokens = self.accepted_tokens[:-k]
+
+    def is_terminated(self):
+        return self.matcher.is_terminated()
+
+    def allocate_vocab_mask(
+        self, vocab_size: int, batch_size: int, device
+    ) -> torch.Tensor:
+        return allocate_token_bitmask(batch_size, vocab_size)
+
+    def fill_vocab_mask(self, vocab_mask: torch.Tensor, idx: int) -> None:
+        self.matcher.fill_next_token_bitmask(vocab_mask, idx)
+
+    @staticmethod
+    def move_vocab_mask(vocab_mask: torch.Tensor, device) -> torch.Tensor:
+        return vocab_mask.to(device, non_blocking=True)
+
+    def apply_vocab_mask(self, logits: torch.Tensor, vocab_mask: torch.Tensor) -> None:
+        if logits.device.type in {"cuda", "xpu", "musa"}:
+            if _is_hip:
+                apply_token_bitmask_inplace_cuda(logits, vocab_mask)
+            else:
+                apply_token_bitmask_inplace_triton(logits, vocab_mask)
+        elif logits.device.type == "npu":
+            import sgl_kernel_npu  # noqa: F401
+
+            torch.ops.npu.apply_token_bitmask(logits, vocab_mask)
+        else:
+            raise RuntimeError(f"Unsupported device: {logits.device.type}")
+
+    def copy(self):
+        matcher = GrammarMatcher(
+            self.ctx,
+            max_rollback_tokens=MAX_ROLLBACK_TOKENS,
+            override_stop_tokens=self.override_stop_tokens,
+        )
+        if grammar_stats := self.grammar_stats:
+            grammar_stats = dataclasses.replace(
+                grammar_stats, is_cache_hit=True, tree_traversal_time=[]
+            )
+        return XGrammarGrammar(
+            matcher,
+            self.vocab_size,
+            self.ctx,
+            self.override_stop_tokens,
+            self.key_string,
+            grammar_stats,
+        )
+
+    def try_jump_forward(self, tokenizer) -> Optional[Tuple[List[int], str]]:
+        s = self.matcher.find_jump_forward_string()
+        if s:
+            return [], s
+        return None
+
+    def jump_forward_str_state(self, helper: Tuple[List[int], str]) -> Tuple[str, int]:
+        _, data = helper
+        return data, -1
+
+    def jump_and_retokenize(
+        self, old_output_ids: List[int], new_output_ids: List[int], next_state: int
+    ):
+        k = 0
+        for i, old_id in enumerate(old_output_ids):
+            if old_id == new_output_ids[i]:
+                k = i + 1
+            else:
+                break
+
+        # rollback to the last token that is the same
+        if k < len(old_output_ids):
+            self.matcher.rollback(len(old_output_ids) - k)
+
+        for i in range(k, len(new_output_ids)):
+            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=})"
+
+
+class TokenizerNotSupportedError(Exception):
+    """Raised when tokenizer is not supported by XGrammar backend."""
+
+    pass
+
+
+class XGrammarGrammarBackend(BaseGrammarBackend):
+    def __init__(
+        self,
+        tokenizer,
+        vocab_size: int,
+        model_eos_token_ids: Optional[List[int]] = None,
+        any_whitespace: bool = True,
+    ):
+        super().__init__()
+
+        if hasattr(tokenizer, "init_xgrammar"):
+            # For special tokenizer
+            tokenizer_info, override_stop_tokens = tokenizer.init_xgrammar()
+
+            if tokenizer_info is None:
+                # Not supported tokenizer
+                raise TokenizerNotSupportedError(
+                    f"Tokenizer type {type(tokenizer).__name__} is not supported by XGrammar"
+                )
+        else:
+            # Create TokenizerInfo with model's EOS tokens as the authoritative stop tokens
+            # This ensures consistency between what the model considers EOS and what XGrammar uses
+            try:
+                tokenizer_info = TokenizerInfo.from_huggingface(
+                    tokenizer, vocab_size=vocab_size, stop_token_ids=model_eos_token_ids
+                )
+                override_stop_tokens = None
+            except Exception as e:
+                raise TokenizerNotSupportedError(
+                    f"Failed to create XGrammar TokenizerInfo from tokenizer: {e}"
+                )
+
+        self.grammar_compiler = GrammarCompiler(tokenizer_info=tokenizer_info)
+        self.vocab_size = vocab_size
+        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."""
+        if not isinstance(structural_format, dict):
+            return
+
+        fmt_type = structural_format.get("type")
+        if fmt_type in {"json_schema", "qwen_xml_parameter"}:
+            if structural_format.get("json_schema") is None:
+                structural_format["json_schema"] = {}
+
+        if fmt_type == "tag":
+            XGrammarGrammarBackend._sanitize_structural_format(
+                structural_format.get("content")
+            )
+        elif fmt_type in {"sequence", "or"}:
+            for element in structural_format.get("elements", []):
+                XGrammarGrammarBackend._sanitize_structural_format(element)
+        elif fmt_type in {"triggered_tags", "tags_with_separator"}:
+            for tag in structural_format.get("tags", []):
+                XGrammarGrammarBackend._sanitize_structural_format(tag)
+
+    @staticmethod
+    def _sanitize_structural_tag_structures(structural_tag: Dict) -> None:
+        for structure in structural_tag.get("structures", []):
+            if structure.get("schema") is None:
+                structure["schema"] = {}
+
+    def _from_context(
+        self, ctx: CompiledGrammar, key_string: str, grammar_stats: GrammarStats
+    ) -> XGrammarGrammar:
+        matcher = GrammarMatcher(
+            ctx,
+            max_rollback_tokens=MAX_ROLLBACK_TOKENS,
+            override_stop_tokens=self.override_stop_tokens,
+        )
+        return XGrammarGrammar(
+            matcher,
+            self.vocab_size,
+            ctx,
+            self.override_stop_tokens,
+            key_string,
+            grammar_stats,
+        )
+
+    def dispatch_json(self, key_string: str) -> BaseGrammarObject:
+        try:
+            if key_string == "$$ANY$$":
+                # Note: This builtin JSON grammar includes *all* valid JSON (including, for example, arrays at the root)
+                ctx = self.grammar_compiler.compile_builtin_json_grammar()
+            else:
+                ctx = self.grammar_compiler.compile_json_schema(
+                    schema=key_string, any_whitespace=self.any_whitespace
+                )
+
+        except (RuntimeError, json.decoder.JSONDecodeError, UnicodeDecodeError) as e:
+            logger.error(f"Hit invalid json_schema: {key_string=}, {e=}")
+            return InvalidGrammarObject(str(e))
+        return self._from_context(ctx, key_string, GrammarStats(dispatch_type="json"))
+
+    def dispatch_ebnf(self, key_string: str) -> BaseGrammarObject:
+        try:
+            ctx = self.grammar_compiler.compile_grammar(key_string)
+        except RuntimeError as e:
+            logger.error(f"Hit invalid ebnf: {key_string=}, {e=}")
+            return InvalidGrammarObject(str(e))
+        return self._from_context(ctx, key_string, GrammarStats(dispatch_type="ebnf"))
+
+    def dispatch_regex(self, key_string: str) -> BaseGrammarObject:
+        try:
+            ctx = self.grammar_compiler.compile_regex(key_string)
+        except RuntimeError as e:
+            logger.error(f"Hit invalid regex: {key_string=}, {e=}")
+            return InvalidGrammarObject(str(e))
+        return self._from_context(ctx, key_string, GrammarStats(dispatch_type="regex"))
+
+    def dispatch_structural_tag(self, key_string: str) -> BaseGrammarObject:
+        try:
+            # TODO(dark): it's REALLY stupid to construct object from string and decode it again
+            structural_tag = json.loads(key_string)
+            if is_legacy_structural_tag(structural_tag):
+                self._sanitize_structural_tag_structures(structural_tag)
+                tags = [
+                    StructuralTagItem(
+                        begin=structure["begin"],
+                        schema=json.dumps(structure["schema"]),
+                        end=structure["end"],
+                    )
+                    for structure in structural_tag["structures"]
+                ]
+                new_tag = StructuralTag.from_legacy_structural_tag(
+                    tags, structural_tag["triggers"]
+                )
+                new_tag.format.at_least_one = structural_tag.get("at_least_one", False)
+                ctx = self.grammar_compiler.compile_structural_tag(new_tag)
+            else:
+                format_dict = structural_tag.get("format")
+                if isinstance(format_dict, dict):
+                    self._sanitize_structural_format(format_dict)
+                    structural_tag["format"] = format_dict
+                    key_string = json.dumps(structural_tag)
+                ctx = self.grammar_compiler.compile_structural_tag(key_string)
+        except (RuntimeError, json.decoder.JSONDecodeError) as e:
+            logger.error(f"Hit invalid structural_tag: {key_string=}, {e=}")
+            return InvalidGrammarObject(str(e))
+        return self._from_context(
+            ctx, key_string, GrammarStats(dispatch_type="structural_tag")
+        )
+
+    def reset(self):
+        super().reset()
+        self.grammar_compiler.clear_cache()
+
+
+def demo_test():
+    from transformers import AutoConfig, AutoTokenizer
+
+    from sglang.test.test_utils import DEFAULT_MODEL_NAME_FOR_TEST
+
+    tokenizer = AutoTokenizer.from_pretrained(DEFAULT_MODEL_NAME_FOR_TEST)
+    hf_config = AutoConfig.from_pretrained(DEFAULT_MODEL_NAME_FOR_TEST)
+
+    # Should use vocab size from model config
+    vocab_size = hf_config.vocab_size
+    eos_token_id = tokenizer.eos_token_id
+
+    backend = XGrammarGrammarBackend(
+        tokenizer, vocab_size=vocab_size, model_eos_token_ids=[eos_token_id]
+    )
+    regex = r"hello (world|there)"
+    grammar = backend.dispatch_regex(regex)
+    tokens = [
+        tokenizer.encode(t, add_special_tokens=False)[0] for t in ["hello", " world"]
+    ]
+
+    # Test termination
+    grammar.accept_token(tokens[0])  # accept "hello"
+    grammar.accept_token(tokens[1])  # accept " world"
+    grammar.accept_token(eos_token_id)  # accept EOS
+    assert grammar.is_terminated()
+
+    # Test rollback the terminated state
+    grammar.rollback(1)
+    assert not grammar.is_terminated()
+
+
+if __name__ == "__main__":
+    demo_test()