minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

Overlap grammar (constrained decoding) with speculative decode verify (#31488)

Browse Source 
Co-authored-by: Jason Park <jasonjk@fb.com>
This commit is contained in:
Lianmin Zheng
2026-07-20 17:36:05 -07:00
committed by GitHub
co-authored by Jason Park
parent 8905cbd42f
commit e7e8aaa73c
12 changed files with 214 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/schedule_batch.py
+1
View File
@@ -3008,6 +3008,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
forward_mode=self.forward_mode,
out_cache_loc=self.out_cache_loc,
return_logprob=self.return_logprob,
has_grammar=self.has_grammar,
decoding_reqs=self.decoding_reqs,
spec_algorithm=self.spec_algorithm,
spec_info=self.spec_info,
python/sglang/srt/managers/scheduler.py
+30 -11
View File
@@ -1646,19 +1646,34 @@ class Scheduler(
and last_batch_is_extend
)
# We do not support overlap + spec + grammar yet,
# so we need to turn off overlap for this batch.
# TODO(lsyin): support overlap + spec + grammar
# 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.
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.forward_mode.is_decode()
and len(self.result_queue) > 0
)
# Algorithms that support grammar overlap advance the FSM inside verify()
# via the grammar barrier (overlapping the target forward), which resolves
# whatever result is still pending in the queue — including the
# extend->decode boundary — so no grammar-specific overlap disable is needed.
return disable_overlap_for_batch or need_grammar_sync
def _advance_pending_grammar(self):
"""Grammar barrier (spec-v2 overlap): advance the FSM over any not-yet
-processed decode result still in the queue, so a following verify()'s
bitmask sees the previous batch's committed tokens. Invoked mid-worker
(before generate_token_bitmask) so the CPU advance overlaps the target
verify forward. Idempotent; no-op when the queue is empty or has no grammar.
"""
for prev_batch, prev_result in self.result_queue:
self.batch_result_processor.advance_grammar_fsm(prev_result, prev_batch)
@scheduler_nvtx_method("scheduler.process_input_requests")
def process_input_requests(self, recv_reqs: List):
now = time.monotonic()
@@ -3337,15 +3352,19 @@ class Scheduler(
# Spec_v2 fires on_publish mid-worker (between verify and
# draft_extend) so schedule prep can overlap with draft_extend.
# Non-spec has no later work — scheduler publishes after return.
fwd_kwargs = (
{
"on_publish": partial(
self.future_map.publish, future_indices
fwd_kwargs = {}
if not batch.spec_algorithm.is_none():
fwd_kwargs["on_publish"] = partial(
self.future_map.publish, future_indices
)
# Grammar-overlap-capable workers advance the grammar FSM
# inside verify() before building the bitmask; hand them the
# barrier that resolves the previous batch's committed
# tokens (overlapping the target forward).
if batch.spec_algorithm.supports_grammar_overlap():
fwd_kwargs["grammar_barrier"] = (
self._advance_pending_grammar
)
}
if not batch.spec_algorithm.is_none()
else {}
)
# FIXME: pp is not compatible with overlap
batch_result = self.model_worker.forward_batch_generation(
python/sglang/srt/managers/scheduler_components/batch_result_processor.py
+89 -17
View File
@@ -273,7 +273,9 @@ class SchedulerBatchResultProcessor:
if req.grammar is not None:
self._apply_prefill_grammar(
req=req, next_token_id=next_token_id
req=req,
next_token_id=next_token_id,
already_advanced=result.grammar_advanced,
)
else:
@@ -487,17 +489,24 @@ class SchedulerBatchResultProcessor:
)
return hidden_state_offset
def _apply_prefill_grammar(self, *, req: Req, next_token_id: int) -> None:
# FIXME: this try-except block is for handling unexpected xgrammar issue.
try:
req.grammar.accept_token(next_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.
logger.error(
f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}"
)
req.to_finish = FINISH_ABORT()
def _apply_prefill_grammar(
self, *, req: Req, next_token_id: int, already_advanced: bool = False
) -> None:
# The grammar barrier may have already advanced the FSM over this prefilled
# token (spec overlap path); only advance if not, but always sync
# grammar.finished.
if not already_advanced:
# FIXME: this try-except block is for handling unexpected xgrammar issue.
try:
req.grammar.accept_token(next_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.
logger.error(
f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}"
)
req.to_finish = FINISH_ABORT()
req.grammar.finished = req.finished()
def _apply_chunked_prefill_logprobs(
@@ -565,6 +574,13 @@ class SchedulerBatchResultProcessor:
result.num_correct_drafts_per_req_cpu, batch_size=len(batch.reqs)
)
# Advance the grammar FSM over this batch's committed tokens (idempotent):
# the EAGLE overlap path already did this inside verify() via the grammar
# barrier; otherwise advance now. advance_grammar_fsm self-gates on per-req
# grammar (the queued batch.copy() does not carry has_grammar) and consumes
# result.grammar_retained_tokens below instead of re-advancing.
self.advance_grammar_fsm(result, batch)
predict_tokens = []
# In adaptive spec-v2, the worker state may already have switched when this
# delayed result is processed. Use the draft token count recorded on result.
@@ -580,11 +596,9 @@ class SchedulerBatchResultProcessor:
pass
else:
if req.grammar is not None:
# Stop accepting once the grammar terminates, so the
# over-drafted suffix is never committed to KV nor emitted.
# This advances the grammar FSM; the result loop only syncs
# grammar.finished.
accept_tokens = self._accept_grammar_tokens(req, accept_tokens)
# FSM already advanced + truncated by advance_grammar_fsm; reuse
# the retained (grammar-legal) run instead of advancing again.
accept_tokens = result.grammar_retained_tokens[i]
# Commit the full accepted run (drafts + bonus).
num_accept_tokens = len(accept_tokens)
@@ -636,6 +650,64 @@ class SchedulerBatchResultProcessor:
req.to_finish = FINISH_ABORT()
return retained
def advance_grammar_fsm(
self, result: GenerationBatchResult, batch: ScheduleBatch
) -> None:
"""Advance each req's grammar FSM over the tokens THIS batch committed, and
(for decode) memoize the grammar-truncated run on ``result``.
This is the single place the spec-v2 FSM advances. It is idempotent
(``result.grammar_advanced``) so it runs either eagerly — inside ``verify(N)``
via the scheduler's grammar barrier, so the advance overlaps the target-verify
forward — or lazily from the result processors on the non-overlap / non-EAGLE
paths. It handles both decode (the accepted spec run) and extend (the single
prefilled token) results, so the barrier can resolve whatever the previous
batch was — e.g. the extend->decode boundary.
"""
if result.grammar_advanced or not batch.has_grammar:
return
is_decode = batch.forward_mode.is_decode()
if not (is_decode or batch.forward_mode.is_extend()):
return
if result.copy_done is not None:
result.copy_done.synchronize()
next_token_ids = result.next_token_ids.tolist()
if not is_decode:
# Extend: advance over the single token each completed-prefill req emitted
# (mirrors process_batch_result_prefill's per-req token indexing).
for i, req in enumerate(batch.reqs):
if (
req.grammar is None
or req.is_retracted
or req.finished()
or req.inflight_middle_chunks > 0
):
continue
self._accept_grammar_tokens(req, next_token_ids[i])
result.grammar_advanced = True
return
# Decode: only the spec-v2 path reaches here (the grammar barrier for
# spec-overlap workers and _resolve_spec_v2_tokens). Non-spec grammar decode
# advances its FSM in process_batch_result_decode and has no accept_lens, so
# bail out defensively.
if result.accept_lens is None:
return
accept_lens = result.accept_lens.tolist()
stride = result.speculative_num_draft_tokens
assert stride is not None, "spec-v2 result missing speculative_num_draft_tokens"
retained = [None] * len(batch.reqs)
for i, req in enumerate(batch.reqs):
if req.grammar is None or req.is_retracted or req.finished():
continue
accept_tokens = next_token_ids[i * stride : i * stride + accept_lens[i]]
# Stop accepting once the grammar terminates so the over-drafted suffix
# is never committed to KV nor emitted; this advances the FSM.
retained[i] = self._accept_grammar_tokens(req, accept_tokens)
result.grammar_retained_tokens = retained
result.grammar_advanced = True
def process_batch_result_idle(
self,
batch: ScheduleBatch,
python/sglang/srt/managers/utils.py
+7
View File
@@ -66,6 +66,13 @@ class GenerationBatchResult:
future_indices: Optional[torch.Tensor] = None
speculative_num_draft_tokens: Optional[int] = None
# Grammar FSM advance memoization (spec-v2 overlap). advance_grammar_fsm sets
# these once — eagerly via the scheduler's grammar barrier inside verify(), or
# lazily in _resolve_spec_v2_tokens — and the latter consumes
# grammar_retained_tokens instead of re-advancing the FSM.
grammar_advanced: bool = False
grammar_retained_tokens: Optional[list] = None
# FIXME(lsyin): maybe move to a better place?
# sync path: forward stream -> output processor
accept_lens: Optional[torch.Tensor] = None