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

[Spec] Unify speculative grammar token-accept path in decode processing (#28682)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-06-19 02:49:40 -07:00
committed by GitHub
parent 941a11ada3
commit 9bb9d17e1a
5 changed files with 207 additions and 306 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler_components/batch_result_processor.py
+75 -101
View File
@@ -556,34 +556,71 @@ class SchedulerBatchResultProcessor:
assert stride is not None, "spec-v2 result missing speculative_num_draft_tokens"
for i, req in enumerate(batch.reqs):
predict_tokens.append(
next_token_ids[i * stride : i * stride + accept_lens[i]]
)
accept_tokens = next_token_ids[i * stride : i * stride + accept_lens[i]]
if req.is_retracted:
# reset_for_retract() already zeroes committed/allocated KV.
continue
if req.finished():
pass
elif req.finished():
if not batch.spec_algorithm.is_dflash():
# EAGLE prepare_for_decode pre-claimed the bonus slot.
req.kv_committed_len -= 1
continue
if batch.spec_algorithm.is_dflash():
# DFLASH materialized accepted draft tokens plus the bonus token.
req.kv_committed_len += accept_lens[i]
else:
# EAGLE prepare_for_decode pre-claimed the bonus slot.
req.kv_committed_len += accept_lens[i] - 1
req.spec_verify_ct += 1
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)
num_correct_drafts = result.num_correct_drafts_per_req_cpu[i]
req.spec_num_correct_drafts += num_correct_drafts
req.update_spec_correct_drafts_histogram(num_correct_drafts)
num_accept_tokens = len(accept_tokens)
if batch.spec_algorithm.is_dflash():
# DFLASH materialized accepted draft tokens plus the bonus token.
req.kv_committed_len += num_accept_tokens
else:
# EAGLE prepare_for_decode pre-claimed the bonus slot.
req.kv_committed_len += num_accept_tokens - 1
req.spec_verify_ct += 1
num_correct_drafts = result.num_correct_drafts_per_req_cpu[i]
req.spec_num_correct_drafts += num_correct_drafts
req.update_spec_correct_drafts_histogram(num_correct_drafts)
predict_tokens.append(accept_tokens)
return predict_tokens
def _accept_grammar_tokens(
self, req: Req, tokens: Union[int, List[int]]
) -> List[int]:
"""Advance the grammar over the accepted token(s), stopping at the token
that terminates it.
``tokens`` is a single sampled token (normal decode) or the whole
verified run (spec decode). Returns the retained prefix; for spec the
suffix past grammar completion is dropped so it is never committed to KV
nor emitted. Advances the grammar FSM only -- ``grammar.finished`` is
synced by the caller once the finish state is updated.
"""
if isinstance(tokens, int):
tokens = [tokens]
retained = []
try:
for token_id in tokens:
req.grammar.accept_token(token_id)
retained.append(token_id)
if req.grammar.is_terminated():
break
except ValueError as e:
# accept_token raises ValueError if the token is not in the grammar
# (misconfigured grammar or invalid token); abort the request.
logger.error(
f"Grammar accept_token failed for req {req.rid} with token "
f"{tokens}: {e}"
)
self.abort_request(AbortReq(rid=req.rid))
return retained
def process_batch_result_idle(
self,
batch: ScheduleBatch,
@@ -645,28 +682,24 @@ class SchedulerBatchResultProcessor:
# And all the over-allocated tokens will be freed in `release_kv_cache`.
continue
# Non-spec and V2: full post-processing
# Non-spec and Spec V2: full post-processing.
next_token_id = next_token_ids[i]
new_accepted_len = 1
needs_tokenwise_grammar_accept = (
not batch.spec_algorithm.is_none() and req.grammar is not None
)
if batch.spec_algorithm.is_none():
is_spec = not batch.spec_algorithm.is_none()
if not is_spec:
# Normal decode: a single sampled token.
req.output_ids.append(next_token_id)
elif needs_tokenwise_grammar_accept:
next_token_id = self._accept_spec_v2_grammar_tokens(req, next_token_id)
next_token_ids[i] = next_token_id
new_accepted_len = len(next_token_id)
new_accept_len = 1
else:
# Spec: accept the whole verified run. For grammar requests the
# run was already truncated at the grammar-terminating token in
# _resolve_spec_v2_tokens, so nothing is emitted past completion.
req.output_ids.extend(next_token_id)
new_accepted_len = len(next_token_id)
if not needs_tokenwise_grammar_accept:
self._maybe_update_reasoning_tokens(req, next_token_id)
new_accept_len = len(next_token_id)
self._maybe_update_reasoning_tokens(req, next_token_id)
req.time_stats.set_last_decode_finish_time()
if not needs_tokenwise_grammar_accept:
req.update_finish_state(new_accepted_len)
req.update_finish_state(new_accept_len)
self._handle_finish_state_updated_req(req, batch, result, i, logits_output)
@@ -681,14 +714,16 @@ class SchedulerBatchResultProcessor:
)
if req.return_hidden_states and logits_output.hidden_states is not None:
if batch.spec_algorithm.is_none():
if not is_spec:
req.hidden_states.append(
logits_output.hidden_states[i].cpu().clone().tolist()
)
else:
# Spec V2: hidden_states is [bs * speculative_num_draft_tokens, hidden_dim].
# One row per emitted token; next_token_id is already truncated
# at grammar termination, so this stays aligned with output_ids.
stride = result.speculative_num_draft_tokens
accept_len = result.num_correct_drafts_per_req_cpu[i] + 1
accept_len = len(next_token_id)
start = i * stride
req.hidden_states.extend(
logits_output.hidden_states[start : start + accept_len]
@@ -698,13 +733,11 @@ class SchedulerBatchResultProcessor:
)
if req.grammar is not None:
if needs_tokenwise_grammar_accept:
# Already advanced token-by-token above; just sync terminal flag.
req.grammar.finished = req.finished()
else:
self._apply_decode_grammar(
req=req, next_token_id=next_token_id, batch=batch
)
if not is_spec:
# Normal decode advances the grammar for its single token
# here; spec already advanced it in _resolve_spec_v2_tokens.
self._accept_grammar_tokens(req, next_token_id)
req.grammar.finished = req.finished()
self.output_streamer.stream_output(batch.reqs, batch.return_logprob)
self.token_to_kv_pool_allocator.free_group_end()
@@ -790,65 +823,6 @@ class SchedulerBatchResultProcessor:
logits_output.next_token_token_ids_logprobs_idx[flat_idx]
)
def _accept_spec_v2_grammar_tokens(
self, req: Req, proposed: List[int]
) -> List[int]:
"""Accept speculative grammar tokens until the request finishes.
Returns the retained prefix and rolls back KV commits for dropped suffix
tokens.
"""
accept_tokens = []
try:
for token_id in proposed:
req.grammar.accept_token(token_id)
req.output_ids.append(token_id)
accept_tokens.append(token_id)
self._maybe_update_reasoning_tokens(req, token_id)
req.update_finish_state()
if req.finished():
break
except ValueError as e:
# accept_token raises ValueError if the token is not in the grammar
# (misconfigured grammar or invalid token); abort the request.
logger.error(
f"Grammar accept_token failed for req {req.rid} with token {proposed}: {e}"
)
self.abort_request(AbortReq(rid=req.rid))
req.update_finish_state()
# _resolve_spec_v2_tokens committed the full proposed list; rollback the
# suffix that grammar termination dropped.
dropped = len(proposed) - len(accept_tokens)
if dropped > 0:
req.kv_committed_len -= dropped
return accept_tokens
def _apply_decode_grammar(
self,
*,
req: Req,
next_token_id: Union[int, List[int]],
batch: ScheduleBatch,
) -> None:
# FIXME: this try-except block is for handling unexpected xgrammar issue.
try:
if batch.spec_algorithm.is_none():
# Normal decode: single token
req.grammar.accept_token(next_token_id)
else:
# 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.
logger.error(
f"Grammar accept_token failed for req {req.rid} with token {next_token_id}: {e}"
)
self.abort_request(AbortReq(rid=req.rid))
req.grammar.finished = req.finished()
def _handle_finish_state_updated_req(
self,
req: Req,
python/sglang/test/kits/spec_server_kits.py
+71
View File
@@ -620,3 +620,74 @@ class SpecHiddenStatesKit:
for row in decode_rows:
self.assertIsInstance(row, list)
self.assertEqual(len(row), hidden_dim)
class SpecGrammarKit:
"""Grammar-constrained structured output under spec decoding.
Regression for spec verify accepting tokens past grammar termination: the
output must be valid JSON with nothing emitted after completion, and the
logprob count must match the (truncated) completion-token count.
"""
# Override per config if a different schema is desired.
grammar_json_schema = json.dumps(
{
"type": "object",
"properties": {
"name": {"type": "string", "pattern": "^[\\w]+$"},
"population": {"type": "integer"},
"country": {"type": "string", "pattern": "^[\\w ]+$"},
"capital": {"type": "string", "pattern": "^[\\w ]+$"},
},
"required": ["name", "population", "country", "capital"],
}
)
def _generate_grammar(self, return_logprob: bool):
response = requests.post(
self.base_url + "/generate",
json={
"text": "Here is the information of the capital of France in the JSON format.\n",
"sampling_params": {
"temperature": 0,
"max_new_tokens": 256,
"json_schema": self.grammar_json_schema,
},
"return_logprob": return_logprob,
"logprob_start_len": 0,
},
)
self.assertEqual(response.status_code, 200, response.text)
out = response.json()
self.assertGreater(
out["meta_info"]["spec_verify_ct"],
0,
"expected spec decoding to run (spec_verify_ct > 0)",
)
return out
def test_grammar_structured_output_no_trailing_tokens(self):
"""Output is valid JSON with nothing emitted past grammar completion."""
out = self._generate_grammar(return_logprob=False)
text = out["text"]
parsed = json.loads(text)
for key in ("name", "population", "country", "capital"):
self.assertIn(key, parsed)
self.assertTrue(
text.strip().endswith("}"), f"unexpected trailing tokens: {text!r}"
)
def test_grammar_logprob_count_matches_completion_tokens(self):
"""Trimmed spec tokens keep logprob count == completion token count."""
out = self._generate_grammar(return_logprob=True)
meta = out["meta_info"]
completion_tokens = meta["completion_tokens"]
output_logprobs = meta["output_token_logprobs"]
self.assertEqual(
len(output_logprobs),
completion_tokens,
"output logprobs must align with retained (trimmed) tokens: "
f"got {len(output_logprobs)} logprobs vs {completion_tokens} completion tokens",
)
json.loads(out["text"])