[Spec] Consolidate the grammar sync decision into ScheduleBatch.grammar_needs_sync (#32353)
This commit is contained in:
@@ -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]
|
||||
):
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user