diff --git a/python/sglang/srt/constrained/base_grammar_backend.py b/python/sglang/srt/constrained/base_grammar_backend.py index 0b6c78416..4f9dc0fb3 100644 --- a/python/sglang/srt/constrained/base_grammar_backend.py +++ b/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( - self, key_type: str, key_string: str - ) -> Optional[BaseGrammarObject]: + 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) -> 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() diff --git a/python/sglang/srt/constrained/grammar_manager.py b/python/sglang/srt/constrained/grammar_manager.py index c2487c8d1..a1ea6d73f 100644 --- a/python/sglang/srt/constrained/grammar_manager.py +++ b/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=}" diff --git a/python/sglang/srt/constrained/llguidance_backend.py b/python/sglang/srt/constrained/llguidance_backend.py index d600a07f3..6029fb909 100644 --- a/python/sglang/srt/constrained/llguidance_backend.py +++ b/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) diff --git a/python/sglang/srt/constrained/outlines_backend.py b/python/sglang/srt/constrained/outlines_backend.py index 28831ab86..eaed4bafe 100644 --- a/python/sglang/srt/constrained/outlines_backend.py +++ b/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 diff --git a/python/sglang/srt/constrained/xgrammar_backend.py b/python/sglang/srt/constrained/xgrammar_backend.py index 3bc306521..56b053fa0 100644 --- a/python/sglang/srt/constrained/xgrammar_backend.py +++ b/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( - self.grammar_stats, is_cache_hit=True, tree_traversal_time=[] - ), + grammar_stats, ) 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) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index e156950d9..5744bb8e5 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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: Optional[BaseGrammarObject] = None + self.grammar_key: Optional[Tuple[str, str]] = 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