feat: Naive support Spec V2 + Constrained Decoding (#13425)
Signed-off-by: Ubospica <ubospica@gmail.com> Co-authored-by: Liangsheng Yin <lsyincs@gmail.com>
This commit is contained in:
@@ -1924,6 +1924,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
||||
dimensions=self.dimensions,
|
||||
dllm_block_offsets=[req.dllm_block_offset for req in self.reqs],
|
||||
dllm_config=self.dllm_config,
|
||||
reqs=self.reqs,
|
||||
has_grammar=self.has_grammar,
|
||||
)
|
||||
|
||||
def copy(self):
|
||||
@@ -2041,3 +2043,8 @@ class ModelWorkerBatch:
|
||||
# Diffusion LLM
|
||||
dllm_block_offsets: Optional[List[int]] = None
|
||||
dllm_config: Optional[DllmConfig] = None
|
||||
|
||||
# For constrained decoding
|
||||
# FIXME(lsyin): remove this after fully overlap grammar
|
||||
reqs: Optional[List[Req]] = None
|
||||
has_grammar: bool = False
|
||||
|
||||
@@ -1016,7 +1016,16 @@ class Scheduler(
|
||||
and self.last_batch.forward_mode.is_extend()
|
||||
)
|
||||
|
||||
if disable_overlap_for_batch:
|
||||
# FIXME(lsyin): remove this grammar sync
|
||||
need_grammar_sync = (
|
||||
batch is not None
|
||||
and batch.forward_mode.is_decode()
|
||||
and batch.has_grammar
|
||||
and batch.is_v2_eagle
|
||||
and len(self.result_queue) > 0
|
||||
)
|
||||
|
||||
if disable_overlap_for_batch or need_grammar_sync:
|
||||
pop_and_process()
|
||||
|
||||
batch_result = None
|
||||
@@ -1025,7 +1034,7 @@ class Scheduler(
|
||||
self.result_queue.append((batch.copy(), batch_result))
|
||||
|
||||
if self.last_batch:
|
||||
if not disable_overlap_for_batch:
|
||||
if not disable_overlap_for_batch and not need_grammar_sync:
|
||||
pop_and_process()
|
||||
elif batch is None:
|
||||
# When the server is idle, do self-check and re-init some states
|
||||
|
||||
@@ -395,10 +395,16 @@ class SchedulerOutputProcessorMixin:
|
||||
logits_output.hidden_states[i].cpu().clone().tolist()
|
||||
)
|
||||
|
||||
if req.grammar is not None and batch.spec_algorithm.is_none():
|
||||
if req.grammar is not None:
|
||||
# FIXME: this try-except block is for handling unexpected xgrammar issue.
|
||||
try:
|
||||
req.grammar.accept_token(next_token_id)
|
||||
if batch.spec_algorithm.is_none():
|
||||
# Normal decode: single token
|
||||
req.grammar.accept_token(next_token_id)
|
||||
elif batch.is_v2_eagle:
|
||||
# Speculative decode: next_token_id is a list of accepted tokens
|
||||
for token_id in next_token_id:
|
||||
req.grammar.accept_token(token_id)
|
||||
except ValueError as e:
|
||||
# Grammar accept_token can raise ValueError if the token is not in the grammar.
|
||||
# This can happen if the grammar is not set correctly or the token is invalid.
|
||||
|
||||
Reference in New Issue
Block a user