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

fast regex decode

Browse Source 
Auto-detect constant str path in regex FSM, then extend instead.
This commit is contained in:
Liangsheng Yin
2024-01-25 01:16:25 +08:00
committed by GitHub
parent 711d343530
commit 01ee0fbc05
16 changed files with 968 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/detokenizer_manager.py
+2
View File
@@ -60,6 +60,8 @@ class DetokenizerManager:
if first_token.startswith("▁"):
output_strs[i] = " " + output_strs[i]
output_strs[i] = recv_obj.output_and_fast_forward_strs[i] + output_strs[i]
self.send_to_tokenizer.send_pyobj(
BatchStrOut(
recv_obj.rids,
python/sglang/srt/managers/io_struct.py
+2
View File
@@ -59,6 +59,7 @@ class GenerateReqInput:
@dataclass
class TokenizedGenerateReqInput:
rid: str
input_text: str
input_ids: List[int]
pixel_values: List[float]
image_hash: int
@@ -73,6 +74,7 @@ class TokenizedGenerateReqInput:
class BatchTokenIDOut:
rids: List[str]
output_tokens: List[List[int]]
output_and_fast_forward_strs: List[str]
hit_stop_str: List[Optional[str]]
skip_special_tokens: List[bool]
meta_info: List[Dict]
python/sglang/srt/managers/router/infer_batch.py
+77
View File
@@ -23,6 +23,7 @@ class FinishReason(Enum):
class Req:
def __init__(self, rid):
self.rid = rid
self.input_text = None
self.input_ids = []
self.output_ids = []
self.pixel_values = None
@@ -48,10 +49,44 @@ class Req:
# for constrained decoding
self.regex_fsm = None
self.regex_fsm_state = 0
self.fast_forward_map = None
self.output_and_fast_forward_str = ""
def max_new_tokens(self):
return self.sampling_params.max_new_tokens
def tokenize_fast_forward(self, fast_forward_str, next_state):
old_output_str = self.tokenizer.decode(self.output_ids)
if self.tokenizer.convert_ids_to_tokens(self.output_ids[0]).startswith("▁"):
old_output_str = " " + old_output_str
new_input_string = (
self.input_text
+ self.output_and_fast_forward_str
+ old_output_str
+ fast_forward_str
)
new_input_ids = self.tokenizer.encode(new_input_string)
fast_forward_tokens_len = (
len(new_input_ids) - len(self.input_ids) - len(self.output_ids)
)
# print("=" * 100)
# print(f"Catch fast forward:\n{fast_forward_str}")
# print(self.tokenizer.convert_ids_to_tokens(self.input_ids))
# print(self.tokenizer.convert_ids_to_tokens(new_input_ids))
self.input_ids = new_input_ids
self.output_ids = []
self.sampling_params.max_new_tokens = max(
self.sampling_params.max_new_tokens - fast_forward_tokens_len, 0
)
self.regex_fsm_state = next_state
self.output_and_fast_forward_str = (
self.output_and_fast_forward_str + old_output_str + fast_forward_str
)
# print(f"Output and fast forward str:\n{self.output_and_fast_forward_str}")
# print("*" * 100)
def check_finished(self):
if self.finished:
return
@@ -263,6 +298,8 @@ class Batch:
req.last_node = None
req.extend_input_len = 0
req.output_ids = []
req.regex_fsm_state = 0
# TODO: apply more fine-grained retraction
token_indices = self.req_to_token_pool.req_to_token[
@@ -274,6 +311,46 @@ class Batch:
return retracted_reqs
def check_for_fast_forward(self):
fast_forward_reqs = []
filter_indices = [i for i in range(len(self.reqs))]
req_pool_indices_cpu = None
for i, req in enumerate(self.reqs):
if req.fast_forward_map is not None:
res = req.fast_forward_map.fast_forward(req.regex_fsm_state)
if res is not None:
fast_forward_str, next_state = res
if len(fast_forward_str) <= 1:
continue
# insert the old request into tree_cache
token_ids_in_memory = tuple(req.input_ids + req.output_ids)[:-1]
if req_pool_indices_cpu is None:
req_pool_indices_cpu = self.req_pool_indices.cpu().tolist()
req_pool_idx = req_pool_indices_cpu[i]
indices = self.req_to_token_pool.req_to_token[
req_pool_idx, : len(token_ids_in_memory)
]
prefix_len = self.tree_cache.insert(
token_ids_in_memory, indices.clone()
)
self.token_to_kv_pool.free(indices[:prefix_len])
self.req_to_token_pool.free(req_pool_idx)
self.tree_cache.dec_ref_counter(req.last_node)
# fast forward
req.tokenize_fast_forward(fast_forward_str, next_state)
fast_forward_reqs.append(req)
filter_indices.remove(i)
if len(filter_indices) < len(self.reqs):
self.filter_batch(filter_indices)
return fast_forward_reqs
def prepare_for_decode(self, input_ids=None):
if input_ids is None:
input_ids = [
python/sglang/srt/managers/router/model_rpc.py
+18 -5
View File
@@ -21,6 +21,7 @@ from sglang.srt.managers.router.radix_cache import RadixCache
from sglang.srt.managers.router.scheduler import Scheduler
from sglang.srt.model_config import ModelConfig
from sglang.srt.server_args import PortArgs, ServerArgs
from sglang.srt.constrained.fast_forward import FastForwardCache
from sglang.srt.utils import (
get_exception_traceback,
get_int_token_logit_bias,
@@ -45,6 +46,7 @@ class ModelRpcServer(rpyc.Service):
self.tp_rank = tp_rank
self.tp_size = server_args.tp_size
self.schedule_heuristic = server_args.schedule_heuristic
self.no_regex_fast_forward = server_args.no_regex_fast_forward
# Init model and tokenizer
self.model_config = ModelConfig(
@@ -118,6 +120,7 @@ class ModelRpcServer(rpyc.Service):
"trust_remote_code": server_args.trust_remote_code,
},
)
self.fast_forward_cache = FastForwardCache()
# Init new token estimation
self.new_token_ratio = min(0.4 * server_args.schedule_conservativeness, 1.0)
@@ -201,6 +204,7 @@ class ModelRpcServer(rpyc.Service):
recv_req: TokenizedGenerateReqInput,
):
req = Req(recv_req.rid)
req.input_text = recv_req.input_text
req.input_ids = recv_req.input_ids
req.pixel_values = recv_req.pixel_values
req.image_size = recv_req.image_size
@@ -223,6 +227,10 @@ class ModelRpcServer(rpyc.Service):
# Init regex fsm
if req.sampling_params.regex is not None:
req.regex_fsm = self.regex_fsm_cache.init_fsm(req.sampling_params.regex)
if not self.no_regex_fast_forward:
req.fast_forward_map = self.fast_forward_cache.init_fast_forward_map(
req.sampling_params.regex
)
# Truncate long prompts
req.input_ids = req.input_ids[: self.model_config.context_len - 1]
@@ -334,11 +342,6 @@ class ModelRpcServer(rpyc.Service):
self.model_config.vocab_size, self.int_token_logit_bias
)
# Reset regex fsm state before first sampling due to retractions
for req in batch.reqs:
if req.sampling_params.regex is not None:
req.regex_fsm_state = 0
if batch.extend_num_tokens != 0:
# Forward
logits, (logprobs, normalized_logprobs) = self.model_runner.forward(
@@ -388,6 +391,13 @@ class ModelRpcServer(rpyc.Service):
self.min_new_token_ratio,
)
if not self.no_regex_fast_forward:
# check for fast forward
fast_forward_reqs = batch.check_for_fast_forward()
self.forward_queue.extend(fast_forward_reqs)
if batch.is_empty():
return
# Update batch tensors
self.decode_forward_ct += 1
batch.prepare_for_decode()
@@ -408,6 +418,7 @@ class ModelRpcServer(rpyc.Service):
def handle_finished_requests(self, batch: Batch):
output_rids = []
output_tokens = []
output_and_fast_forward_strs = []
output_hit_stop_str = []
output_skip_special_tokens = []
output_meta_info = []
@@ -425,6 +436,7 @@ class ModelRpcServer(rpyc.Service):
):
output_rids.append(req.rid)
output_tokens.append(req.output_ids)
output_and_fast_forward_strs.append(req.output_and_fast_forward_str)
output_hit_stop_str.append(req.hit_stop_str)
output_skip_special_tokens.append(
req.sampling_params.skip_special_tokens
@@ -445,6 +457,7 @@ class ModelRpcServer(rpyc.Service):
BatchTokenIDOut(
output_rids,
output_tokens,
output_and_fast_forward_strs,
output_hit_stop_str,
output_skip_special_tokens,
output_meta_info,
python/sglang/srt/managers/tokenizer_manager.py
+1
View File
@@ -157,6 +157,7 @@ class TokenizerManager:
)
tokenized_obj = TokenizedGenerateReqInput(
rid=rid,
input_text=obj.text,
input_ids=input_ids,
pixel_values=pixel_values,
image_hash=image_hash,