chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Improve structured outputs: fix race condition, server crash, metrics and style (#6188)

Browse Source
This commit is contained in:
Lianmin Zheng
2025-05-11 08:36:16 -07:00
committed by GitHub
parent 94d42b6794
commit 01bdbf7f80
13 changed files with 568 additions and 258 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/constrained/xgrammar_backend.py
+69 -43
View File
@@ -18,7 +18,6 @@ import logging
from typing import List, Optional, Tuple, Union
import torch
import xgrammar
from xgrammar import (
CompiledGrammar,
GrammarCompiler,
@@ -35,7 +34,6 @@ from sglang.srt.constrained.base_grammar_backend import (
from sglang.srt.constrained.triton_ops.bitmask_ops import (
apply_token_bitmask_inplace_triton,
)
from sglang.srt.utils import get_bool_env_var
logger = logging.getLogger(__name__)
@@ -51,49 +49,35 @@ class XGrammarGrammar(BaseGrammarObject):
vocab_size: int,
ctx: CompiledGrammar,
override_stop_tokens: Optional[Union[List[int], int]],
key_string: Optional[str] = None, # TODO (sk): for debugging, remove later
) -> None:
super().__init__()
self.matcher = matcher
self.vocab_size = vocab_size
self.ctx = ctx
self.override_stop_tokens = override_stop_tokens
self.finished = False
from xgrammar.kernels.apply_token_bitmask_inplace_cpu import (
apply_token_bitmask_inplace_cpu,
)
self.apply_vocab_mask_cpu = apply_token_bitmask_inplace_cpu
self.accepted_tokens = []
self.key_string = key_string
def accept_token(self, token: int):
assert self.matcher.accept_token(token)
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
if not self.is_terminated():
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:
break
self.accepted_tokens.append(token)
# rollback to the last token that is the same
if k < len(old_output_ids):
self.matcher.rollback(len(old_output_ids) - k)
def rollback(self, k: int):
self.matcher.rollback(k)
self.accepted_tokens = self.accepted_tokens[:-k]
for i in range(k, len(new_output_ids)):
assert self.matcher.accept_token(new_output_ids[i])
def is_terminated(self):
return self.matcher.is_terminated()
def allocate_vocab_mask(
self, vocab_size: int, batch_size: int, device
@@ -122,9 +106,43 @@ class XGrammarGrammar(BaseGrammarObject):
override_stop_tokens=self.override_stop_tokens,
)
return XGrammarGrammar(
matcher, self.vocab_size, self.ctx, self.override_stop_tokens
matcher,
self.vocab_size,
self.ctx,
self.override_stop_tokens,
self.key_string,
)
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)):
assert self.matcher.accept_token(new_output_ids[i])
def __repr__(self):
return f"XGrammarGrammar({self.key_string=}, {self.accepted_tokens=})"
class XGrammarGrammarBackend(BaseGrammarBackend):
def __init__(
@@ -143,9 +161,15 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
self.vocab_size = vocab_size
self.override_stop_tokens = override_stop_tokens
def _from_context(self, ctx: CompiledGrammar) -> XGrammarGrammar:
matcher = GrammarMatcher(ctx, max_rollback_tokens=MAX_ROLLBACK_TOKENS)
return XGrammarGrammar(matcher, self.vocab_size, ctx, self.override_stop_tokens)
def _from_context(self, ctx: CompiledGrammar, key_string: str) -> 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
)
def dispatch_json(self, key_string: str) -> Optional[XGrammarGrammar]:
try:
@@ -157,7 +181,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
except RuntimeError as e:
logging.warning(f"Skip invalid json_schema: json_schema={key_string}, {e=}")
return None
return self._from_context(ctx)
return self._from_context(ctx, key_string)
def dispatch_ebnf(self, key_string: str) -> Optional[XGrammarGrammar]:
try:
@@ -165,7 +189,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
except RuntimeError as e:
logging.warning(f"Skip invalid ebnf: ebnf={key_string}, {e=}")
return None
return self._from_context(ctx)
return self._from_context(ctx, key_string)
def dispatch_regex(self, key_string: str) -> Optional[XGrammarGrammar]:
try:
@@ -173,7 +197,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
except RuntimeError as e:
logging.warning(f"Skip invalid regex: regex={key_string}, {e=}")
return None
return self._from_context(ctx)
return self._from_context(ctx, key_string)
def dispatch_structural_tag(self, key_string: str) -> Optional[XGrammarGrammar]:
try:
@@ -190,9 +214,11 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
tags, structural_tag["triggers"]
)
except RuntimeError as e:
logging.warning(f"Skip invalid regex: regex={key_string}, {e=}")
logging.warning(
f"Skip invalid structural_tag: structural_tag={key_string}, {e=}"
)
return None
return self._from_context(ctx)
return self._from_context(ctx, key_string)
def reset(self):
if self.grammar_compiler: