[Minor] fix type annotations and invalid method calls in constrained … (#20132)

python/sglang/srt/constrained/base_grammar_backend.py

@@ -15,9 +15,8 @@
 
 import logging
 import time
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import Future, ThreadPoolExecutor
 from dataclasses import dataclass, field
-from threading import Event
 from typing import Dict, List, Optional, Tuple
 
 import torch
@@ -120,43 +119,36 @@ class BaseGrammarObject:
 INVALID_GRAMMAR_OBJ = BaseGrammarObject()
 
 
-@dataclass
-class CacheEntry:
-    value: BaseGrammarObject
-    event: Event
-
-
 class BaseGrammarBackend:
     def __init__(self):
         self.executor = ThreadPoolExecutor()
-        self.cache: Dict[Tuple[str, str], CacheEntry] = {}
+        self.cache: Dict[Tuple[str, str], BaseGrammarObject] = {}
 
-    def _not_supported(self, key_type: str, key_string: str) -> None:
+    def _not_supported(self, key_type: str, key_string: str) -> BaseGrammarObject:
         logger.warning(f"Skip unsupported {key_type=}, {key_string=}")
+        return INVALID_GRAMMAR_OBJ
 
-    def dispatch_fallback(
+    def dispatch_fallback(self, key_type: str, key_string: str) -> BaseGrammarObject:
-        self, key_type: str, key_string: str
-    ) -> Optional[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) -> Optional[BaseGrammarObject]:
+    def dispatch_json(self, key_string: str) -> BaseGrammarObject:
         return self._not_supported("json", key_string)
 
-    def dispatch_regex(self, key_string: str) -> Optional[BaseGrammarObject]:
+    def dispatch_regex(self, key_string: str) -> BaseGrammarObject:
         return self._not_supported("regex", key_string)
 
-    def dispatch_ebnf(self, key_string: str) -> Optional[BaseGrammarObject]:
+    def dispatch_ebnf(self, key_string: str) -> BaseGrammarObject:
         return self._not_supported("ebnf", key_string)
 
-    def dispatch_structural_tag(self, key_string: str) -> Optional[BaseGrammarObject]:
+    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
-    ) -> Optional[BaseGrammarObject]:
+    ) -> BaseGrammarObject:
         s = time.perf_counter()
         key_type, key_string = key
         if key_type == "json":
@@ -167,10 +159,6 @@ class BaseGrammarBackend:
             grammar = self.dispatch_ebnf(key_string)
         elif key_type == "structural_tag":
             grammar = self.dispatch_structural_tag(key_string)
-        elif key_type == "structural_pattern":
-            grammar = self.dispatch_structural_pattern(key_string)
-        elif key_type == "structural_pattern_v2":
-            grammar = self.dispatch_structural_pattern_v2(key_string)
         else:
             grammar = self.dispatch_fallback(key_type, key_string)
 
@@ -180,7 +168,7 @@ class BaseGrammarBackend:
 
     def get_cached_or_future_value(
         self, key: Tuple[str, str], require_reasoning: bool
-    ) -> Optional[BaseGrammarObject]:
+    ) -> Tuple[BaseGrammarObject | Future[BaseGrammarObject], bool]:
         value = self.cache.get(key)
         if value:
             copied_value = value.copy()

python/sglang/srt/constrained/grammar_manager.py

@@ -60,7 +60,7 @@ class GrammarManager:
         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 req.grammar:
+                if isinstance(req.grammar, futures.Future) and req.grammar:
                     req.grammar.cancel()
                 req.set_finish_with_abort("Aborted by AbortReq.")
 
@@ -115,6 +115,7 @@ class GrammarManager:
         ready_reqs = intersect(ready_reqs_all)
         failed_reqs = union(failed_reqs_all)
         """
+        assert self.grammar_backend
         ready_req_idxs: set[int] = set()
         failed_req_idxs: set[int] = set()
 
@@ -170,6 +171,7 @@ class GrammarManager:
             if req.finished() or req.grammar is None:  # It is aborted by AbortReq
                 continue
 
+            assert isinstance(req.grammar, futures.Future) and req.grammar_key
             req.grammar = req.grammar.result()
             self.grammar_backend.set_cache(req.grammar_key, req.grammar.copy())
             if req.grammar is INVALID_GRAMMAR_OBJ:
@@ -181,6 +183,7 @@ class GrammarManager:
             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, INVALID_GRAMMAR_OBJ)
             error_msg = f"Grammar preprocessing timed out: {req.grammar_key=}"

python/sglang/srt/constrained/llguidance_backend.py

@@ -122,7 +122,7 @@ class GuidanceBackend(BaseGrammarBackend):
         self.whitespace_pattern = whitespace_pattern
         self.llguidance_tokenizer = from_tokenizer(self.tokenizer, n_vocab)
 
-    def _from_serialized(self, serialized_grammar) -> Optional[GuidanceGrammar]:
+    def _from_serialized(self, serialized_grammar) -> BaseGrammarObject:
         try:
             return GuidanceGrammar(
                 llguidance_tokenizer=self.llguidance_tokenizer,
@@ -132,7 +132,7 @@ class GuidanceBackend(BaseGrammarBackend):
             logger.error(f"Hit invalid grammar: {serialized_grammar=}, {e=}")
             return INVALID_GRAMMAR_OBJ
 
-    def dispatch_json(self, key_string: str) -> Optional[GuidanceGrammar]:
+    def dispatch_json(self, key_string: str) -> BaseGrammarObject:
         try:
             serialized_grammar = LLMatcher.grammar_from_json_schema(
                 key_string,
@@ -146,11 +146,11 @@ class GuidanceBackend(BaseGrammarBackend):
             return INVALID_GRAMMAR_OBJ
         return self._from_serialized(serialized_grammar)
 
-    def dispatch_regex(self, key_string: str) -> Optional[GuidanceGrammar]:
+    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) -> Optional[GuidanceGrammar]:
+    def dispatch_ebnf(self, key_string: str) -> BaseGrammarObject:
         try:
             serialized_grammar = grammar_from("ebnf", key_string)
             return self._from_serialized(serialized_grammar)
@@ -158,7 +158,7 @@ class GuidanceBackend(BaseGrammarBackend):
             logger.error(f"Hit invalid ebnf: {key_string=}, {e=}")
             return INVALID_GRAMMAR_OBJ
 
-    def dispatch_structural_tag(self, key_string: str) -> Optional[GuidanceGrammar]:
+    def dispatch_structural_tag(self, key_string: str) -> BaseGrammarObject:
         try:
             structural_tag = json.loads(key_string)
             assert is_legacy_structural_tag(structural_tag)

python/sglang/srt/constrained/outlines_backend.py

@@ -142,7 +142,7 @@ class OutlinesGrammarBackend(BaseGrammarBackend):
             )
         self.whitespace_pattern = whitespace_pattern
 
-    def _compile_regex(self, regex: str) -> Optional[OutlinesGrammar]:
+    def _compile_regex(self, regex: str) -> BaseGrammarObject:
         try:
             if hasattr(RegexGuide, "from_regex"):
                 # outlines >= 0.1.1

python/sglang/srt/constrained/xgrammar_backend.py

@@ -113,8 +113,6 @@ class XGrammarGrammar(BaseGrammarObject):
                 apply_token_bitmask_inplace_cuda(logits, vocab_mask)
             else:
                 apply_token_bitmask_inplace_triton(logits, vocab_mask)
-        elif logits.device.type == "cpu" and self.apply_vocab_mask_cpu:
-            self.apply_vocab_mask_cpu(logits, vocab_mask)
         else:
             raise RuntimeError(f"Unsupported device: {logits.device.type}")
 
@@ -124,15 +122,17 @@ class XGrammarGrammar(BaseGrammarObject):
             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,
-            dataclasses.replace(
+            grammar_stats,
-                self.grammar_stats, is_cache_hit=True, tree_traversal_time=[]
-            ),
         )
 
     def try_jump_forward(self, tokenizer) -> Optional[Tuple[List[int], str]]:
@@ -254,7 +254,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
             grammar_stats,
         )
 
-    def dispatch_json(self, key_string: str) -> Optional[XGrammarGrammar]:
+    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)
@@ -269,7 +269,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
             return INVALID_GRAMMAR_OBJ
         return self._from_context(ctx, key_string, GrammarStats(dispatch_type="json"))
 
-    def dispatch_ebnf(self, key_string: str) -> Optional[XGrammarGrammar]:
+    def dispatch_ebnf(self, key_string: str) -> BaseGrammarObject:
         try:
             ctx = self.grammar_compiler.compile_grammar(key_string)
         except RuntimeError as e:
@@ -277,7 +277,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
             return INVALID_GRAMMAR_OBJ
         return self._from_context(ctx, key_string, GrammarStats(dispatch_type="ebnf"))
 
-    def dispatch_regex(self, key_string: str) -> Optional[XGrammarGrammar]:
+    def dispatch_regex(self, key_string: str) -> BaseGrammarObject:
         try:
             ctx = self.grammar_compiler.compile_regex(key_string)
         except RuntimeError as e:
@@ -285,7 +285,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
             return INVALID_GRAMMAR_OBJ
         return self._from_context(ctx, key_string, GrammarStats(dispatch_type="regex"))
 
-    def dispatch_structural_tag(self, key_string: str) -> Optional[XGrammarGrammar]:
+    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)

python/sglang/srt/managers/schedule_batch.py

@@ -39,6 +39,7 @@ import copy
 import dataclasses
 import logging
 import re
+from concurrent.futures import Future
 from enum import Enum, auto
 from functools import lru_cache
 from http import HTTPStatus
@@ -725,8 +726,10 @@ class Req(ReqDllmMixin):
         self.embedding = None
 
         # Constrained decoding
-        self.grammar_key: Optional[str] = None
+        self.grammar_key: Optional[Tuple[str, str]] = None
-        self.grammar: Optional[BaseGrammarObject] = None
+        self.grammar: Optional[Union[BaseGrammarObject, Future[BaseGrammarObject]]] = (
+            None
+        )
         self.grammar_wait_ct = 0
 
         # The number of cached tokens that were already cached in the KV cache