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:
Yixin Dong
2025-11-27 04:31:46 -08:00
committed by GitHub
parent 25758647b1
commit 6350042696
8 changed files with 149 additions and 8 deletions

View File

@@ -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

View File

@@ -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

View File

@@ -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.