[Spec] Consolidate the grammar sync decision into ScheduleBatch.grammar_needs_sync (#32353)

This commit is contained in:
Liangsheng Yin
2026-07-24 20:42:38 -07:00
committed by GitHub
parent 95865de24f
commit cff20a2fbb
4 changed files with 12 additions and 7 deletions
@@ -2042,6 +2042,11 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
def is_dllm(self):
return self.dllm_config is not None
def grammar_needs_sync(self) -> bool:
"""Whether grammar forces this batch onto the synchronous path, i.e. the
previous batch's result is resolved before this forward."""
return self.has_grammar and not self.spec_algorithm.supports_grammar_overlap()
def prepare_encoder_info_extend(
self, input_ids: List[array[int]], seq_lens: List[int]
):
+3 -5
View File
@@ -1672,14 +1672,12 @@ class Scheduler(
and last_batch_is_extend
)
# Spec algorithms that don't advance the grammar FSM inside verify() (see
# supports_grammar_overlap) still need overlap forced off for grammar decode
# batches, so the FSM is advanced before the next batch's bitmask.
# Sync so the FSM advance lands before the next batch's bitmask. Permanent
# path for host-draft algorithms, not a pending migration.
need_grammar_sync = (
batch
and not batch.spec_algorithm.is_none()
and not batch.spec_algorithm.supports_grammar_overlap()
and batch.has_grammar
and batch.grammar_needs_sync()
and batch.forward_mode.is_decode()
and len(self.result_queue) > 0
)
@@ -230,7 +230,7 @@ class NGRAMWorker(BaseSpecWorker):
# spliced in from spec_info. Sync mode and grammar batches process
# results before the next draft prep, so output_ids is already
# complete and splicing would duplicate the tail.
use_prev_tokens = self.enable_overlap and not batch.has_grammar
use_prev_tokens = self.enable_overlap and not batch.grammar_needs_sync()
i = 0
for req in batch.reqs:
prev_tokens = (
@@ -344,7 +344,7 @@ class NGRAMWorker(BaseSpecWorker):
i, stride = 0, self.draft_token_num
# Same splice condition as _prepare_draft_tokens: only overlap mode
# has accepted tokens missing from req.output_ids.
use_prev_tokens = self.enable_overlap and not batch.has_grammar
use_prev_tokens = self.enable_overlap and not batch.grammar_needs_sync()
for req in batch.reqs:
# FIXME: Whether to insert 'extend' into the cache or not, after testing,
# there is not much difference, so we will not insert it for now.
@@ -136,6 +136,8 @@ class SpeculativeAlgorithm(Enum):
def supports_grammar_overlap(self) -> bool:
# Whether the worker advances the grammar FSM inside verify() (via the
# scheduler's grammar barrier), letting spec + grammar decode overlap.
# Needs a GPU draft phase to hide the grammar CPU work under: NGRAM drafts
# from a host corpus lookup, so it stays synchronous by design.
# STANDALONE inherits the EAGLE V2 worker's verify path, barrier included.
return self.is_eagle() or self.is_standalone()