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

[Performance] Support both xgrammar and outlines for constrained decoding (#1752)

Browse Source
This commit is contained in:
DarkSharpness
2024-10-25 21:47:02 +00:00
committed by GitHub
parent 30643fed7f
commit b77a02cdfd
7 changed files with 325 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/schedule_batch.py
+20 -41
View File
@@ -37,8 +37,7 @@ import torch
from sglang.global_config import global_config
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.constrained import RegexGuide
from sglang.srt.constrained.jump_forward import JumpForwardMap
from sglang.srt.constrained.grammar import Grammar
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
from sglang.srt.mem_cache.chunk_cache import ChunkCache
from sglang.srt.mem_cache.memory_pool import BaseTokenToKVPool, ReqToTokenPool
@@ -247,9 +246,7 @@ class Req:
self.embedding = None
# Constrained decoding
self.regex_fsm: RegexGuide = None
self.regex_fsm_state: int = 0
self.jump_forward_map: JumpForwardMap = None
self.grammar: Optional[Grammar] = None
# For Qwen2-VL
self.mrope_position_delta = [] # use mutable object
@@ -359,6 +356,8 @@ class Req:
return
def jump_forward_and_retokenize(self, jump_forward_str, next_state):
assert self.grammar is not None and self.tokenizer is not None
if self.origin_input_text is None:
# Recovering text can only use unpadded ids
self.origin_input_text = self.tokenizer.decode(
@@ -398,7 +397,8 @@ class Req:
self.surr_offset = self.read_offset - i
break
self.regex_fsm_state = next_state
# update the inner state of the grammar
self.grammar.jump_and_retokenize(old_output_ids, self.output_ids, next_state)
if self.return_logprob:
# For fast-forward part's logprobs
@@ -468,8 +468,8 @@ class ScheduleBatch:
# Stream
has_stream: bool = False
# Has regex
has_regex: bool = False
# Has grammar
has_grammar: bool = False
# device
device: str = "cuda"
@@ -477,7 +477,7 @@ class ScheduleBatch:
@classmethod
def init_new(
cls,
reqs,
reqs: List[Req],
req_to_token_pool,
token_to_kv_pool,
tree_cache,
@@ -491,7 +491,7 @@ class ScheduleBatch:
model_config=model_config,
return_logprob=any(req.return_logprob for req in reqs),
has_stream=any(req.stream for req in reqs),
has_regex=any(req.regex_fsm for req in reqs),
has_grammar=any(req.grammar for req in reqs),
device=req_to_token_pool.device,
)
@@ -803,26 +803,10 @@ class ScheduleBatch:
keep_indices = set(i for i in range(len(self.reqs)))
for i, req in enumerate(self.reqs):
if req.jump_forward_map is not None:
jump_forward_bytes = req.jump_forward_map.jump_forward_byte(
req.regex_fsm_state
)
if jump_forward_bytes is not None and len(jump_forward_bytes) > 1:
suffix_bytes = []
continuation_range = range(0x80, 0xC0)
cur_state = req.regex_fsm_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 = req.tokenizer.convert_tokens_to_ids(suffix_tokens)
if req.grammar is not None:
jump_helper = req.grammar.try_jump(req.tokenizer)
if jump_helper.can_jump():
suffix_ids = jump_helper.suffix_ids
# Current ids, for cache and revert
cur_all_ids = tuple(req.origin_input_ids + req.output_ids)[:-1]
cur_output_ids = req.output_ids
@@ -836,10 +820,8 @@ class ScheduleBatch:
(
jump_forward_str,
next_state,
) = req.jump_forward_map.jump_forward_symbol(cur_state)
) = req.grammar.jump_forward_str_state(jump_helper)
# Make the incrementally decoded text part of jump_forward_str
# so that the UTF-8 will not corrupt
jump_forward_str = new_text + jump_forward_str
if not req.jump_forward_and_retokenize(
jump_forward_str, next_state
@@ -946,7 +928,7 @@ class ScheduleBatch:
self.top_logprobs_nums = None
self.has_stream = any(req.stream for req in self.reqs)
self.has_regex = any(req.regex_fsm for req in self.reqs)
self.has_grammar = any(req.grammar for req in self.reqs)
self.sampling_info.filter_batch(keep_indices, new_indices)
@@ -979,7 +961,7 @@ class ScheduleBatch:
self.return_logprob = self.return_logprob or other.return_logprob
self.has_stream = self.has_stream or other.has_stream
self.has_regex = self.has_regex or other.has_regex
self.has_grammar = self.has_grammar or other.has_grammar
def get_model_worker_batch(self):
if self.forward_mode.is_decode():
@@ -989,13 +971,10 @@ class ScheduleBatch:
extend_prefix_lens = self.prefix_lens
extend_logprob_start_lens = self.extend_logprob_start_lens
if self.has_regex:
self.sampling_info.regex_fsms = [req.regex_fsm for req in self.reqs]
self.sampling_info.regex_fsm_states = [
req.regex_fsm_state for req in self.reqs
]
if self.has_grammar:
self.sampling_info.grammars = [req.grammar for req in self.reqs]
else:
self.sampling_info.regex_fsms = None
self.sampling_info.grammars = None
global bid
bid += 1
python/sglang/srt/managers/scheduler.py
+21 -23
View File
@@ -29,8 +29,7 @@ import zmq
from sglang.global_config import global_config
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.constrained.fsm_cache import FSMCache
from sglang.srt.constrained.jump_forward import JumpForwardCache
from sglang.srt.constrained.grammar import GrammarCache
from sglang.srt.hf_transformers_utils import get_processor, get_tokenizer
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.managers.io_struct import (
@@ -225,17 +224,20 @@ class Scheduler:
)
# Init the FSM cache for constrained generation
self.grammar_cache = None
if not server_args.skip_tokenizer_init:
self.regex_fsm_cache = FSMCache(
self.grammar_cache = GrammarCache(
server_args.tokenizer_path,
{
"tokenizer_mode": server_args.tokenizer_mode,
"trust_remote_code": server_args.trust_remote_code,
},
skip_tokenizer_init=server_args.skip_tokenizer_init,
constrained_json_whitespace_pattern=server_args.constrained_json_whitespace_pattern,
whitespace_patterns=server_args.constrained_json_whitespace_pattern,
backend=server_args.grammar_backend,
allow_jump=not server_args.disable_regex_jump_forward,
)
self.jump_forward_cache = JumpForwardCache()
# Init new token estimation
assert (
@@ -402,22 +404,20 @@ class Scheduler:
# By default, only return the logprobs for output tokens
req.logprob_start_len = len(recv_req.input_ids) - 1
# Init regex FSM
# Init regex FSM or BNF
if (
req.sampling_params.json_schema is not None
or req.sampling_params.regex is not None
):
assert self.grammar_cache is not None
if req.sampling_params.json_schema is not None:
req.regex_fsm, computed_regex_string = self.regex_fsm_cache.query(
("json", req.sampling_params.json_schema)
req.grammar = self.grammar_cache.query(
("json", req.sampling_params.json_schema),
self.model_config.vocab_size,
)
elif req.sampling_params.regex is not None:
req.regex_fsm, computed_regex_string = self.regex_fsm_cache.query(
("regex", req.sampling_params.regex)
)
if not self.disable_regex_jump_forward:
req.jump_forward_map = self.jump_forward_cache.query(
computed_regex_string
req.grammar = self.grammar_cache.query(
("regex", req.sampling_params.regex), self.model_config.vocab_size
)
# Truncate prompts that are too long
@@ -796,10 +796,8 @@ class Scheduler:
elif not batch.decoding_reqs or req not in batch.decoding_reqs:
self.tree_cache.cache_unfinished_req(req)
if req.regex_fsm is not None:
req.regex_fsm_state = req.regex_fsm.get_next_state(
req.regex_fsm_state, next_token_ids[i]
)
if req.grammar is not None:
req.grammar.accept_token(next_token_ids[i])
if req.return_logprob:
logprob_pt += self.add_logprob_return_values(
@@ -855,10 +853,8 @@ class Scheduler:
req.output_ids.append(next_token_id)
req.check_finished()
if req.regex_fsm is not None:
req.regex_fsm_state = req.regex_fsm.get_next_state(
req.regex_fsm_state, next_token_id
)
if req.grammar is not None:
req.grammar.accept_token(next_token_id)
if req.finished():
self.tree_cache.cache_finished_req(req)
@@ -1056,7 +1052,9 @@ class Scheduler:
):
self.tree_cache.reset()
self.tree_cache_metrics = {"total": 0, "hit": 0}
self.regex_fsm_cache.reset()
if self.grammar_cache is not None:
self.grammar_cache.reset()
# TODO(dark): reset the bnf cache
self.req_to_token_pool.clear()
self.token_to_kv_pool.clear()
torch.cuda.empty_cache()