[Spec] Unify speculative grammar token-accept path in decode processing (#28682)
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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"])
|
||||
|
||||
@@ -16,6 +16,7 @@ from transformers import AutoTokenizer
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.kits.json_constrained_kit import JSONConstrainedMixin
|
||||
from sglang.test.kits.pause_generation_kit import PauseResumeInPlaceMixin
|
||||
from sglang.test.kits.spec_server_kits import SpecGrammarKit
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.server_fixtures.disaggregation_fixture import (
|
||||
PDDisaggregationServerBase,
|
||||
@@ -27,7 +28,7 @@ from sglang.test.test_utils import (
|
||||
DEFAULT_TARGET_MODEL_EAGLE3,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=890, stage="base-b", runner_config="2-gpu-large")
|
||||
register_cuda_ci(est_time=730, stage="base-b", runner_config="2-gpu-large")
|
||||
|
||||
|
||||
class TestDisaggregationAccuracy(PauseResumeInPlaceMixin, PDDisaggregationServerBase):
|
||||
@@ -221,7 +222,9 @@ class TestDisaggregationMooncakeFailure(PDDisaggregationServerBase):
|
||||
raise e from health_check_error
|
||||
|
||||
|
||||
class TestDisaggregationMooncakeSpec(JSONConstrainedMixin, PDDisaggregationServerBase):
|
||||
class TestDisaggregationMooncakeSpec(
|
||||
JSONConstrainedMixin, SpecGrammarKit, PDDisaggregationServerBase
|
||||
):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
@@ -260,104 +263,6 @@ class TestDisaggregationMooncakeSpec(JSONConstrainedMixin, PDDisaggregationServe
|
||||
self.assertGreater(metrics["score"], 0.74)
|
||||
|
||||
|
||||
class TestDisaggregationSpecV2Grammar(PDDisaggregationServerBase):
|
||||
"""Regression for PD disagg + EAGLE Spec V2 + grammar structured output."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super().setUpClass()
|
||||
cls.model = DEFAULT_TARGET_MODEL_EAGLE3
|
||||
spec_args = [
|
||||
"--speculative-algorithm",
|
||||
"EAGLE",
|
||||
"--speculative-draft-model-path",
|
||||
DEFAULT_DRAFT_MODEL_EAGLE3,
|
||||
"--speculative-num-steps",
|
||||
"3",
|
||||
"--speculative-eagle-topk",
|
||||
"1", # Spec V2 only supports topk=1
|
||||
"--speculative-num-draft-tokens",
|
||||
"4",
|
||||
"--grammar-backend",
|
||||
"xgrammar",
|
||||
"--cuda-graph-max-bs",
|
||||
"8",
|
||||
"--dtype=float16",
|
||||
# Cap context to the EAGLE3 draft's native length so the target and
|
||||
# draft ModelConfigs agree (the draft derives 2048 vs the Llama-3.1
|
||||
# target's 131072, which the Spec V2 draft worker otherwise rejects).
|
||||
# 2048 is far above this test's output length.
|
||||
"--context-length",
|
||||
"2048",
|
||||
]
|
||||
cls.extra_prefill_args = spec_args
|
||||
cls.extra_decode_args = spec_args
|
||||
cls.launch_all()
|
||||
|
||||
@staticmethod
|
||||
def _json_schema() -> str:
|
||||
return 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(self, return_logprob: bool):
|
||||
response = requests.post(
|
||||
f"{self.lb_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._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 V2 to run (spec_verify_ct > 0)",
|
||||
)
|
||||
return out
|
||||
|
||||
def test_structured_output_no_trailing_tokens(self):
|
||||
"""Output is valid JSON with nothing emitted past grammar completion."""
|
||||
out = self._generate(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_logprob_count_matches_completion_tokens(self):
|
||||
"""Trimmed Spec V2 tokens keep logprob count == completion token count."""
|
||||
out = self._generate(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"])
|
||||
|
||||
|
||||
class TestDisaggregationSimulatedRetract(PDDisaggregationServerBase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
|
||||
@@ -1,8 +1,14 @@
|
||||
"""Unit tests for Spec V2 grammar trimming in process_batch_result_decode."""
|
||||
"""Unit tests for Spec V2 grammar truncation in _resolve_spec_v2_tokens.
|
||||
|
||||
The grammar-constrained spec path stops accepting at the grammar-terminating
|
||||
token, so the over-drafted suffix is never committed to KV nor emitted.
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.managers.schedule_batch import Req
|
||||
from sglang.srt.managers.scheduler_components.batch_result_processor import (
|
||||
SchedulerBatchResultProcessor,
|
||||
@@ -15,7 +21,7 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class _FakeGrammar:
|
||||
"""Grammar stub that reports termination after `terminate_after` tokens."""
|
||||
"""Grammar stub that terminates after `terminate_after` accepted tokens."""
|
||||
|
||||
def __init__(self, terminate_after: int):
|
||||
self.accepted = []
|
||||
@@ -33,136 +39,86 @@ class _FakeSpecAlgorithm:
|
||||
def is_none(self) -> bool:
|
||||
return False
|
||||
|
||||
def is_dflash(self) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
class _FakeBatch:
|
||||
def __init__(self, reqs, return_logprob: bool):
|
||||
def __init__(self, reqs):
|
||||
self.reqs = reqs
|
||||
self.return_logprob = return_logprob
|
||||
self.spec_algorithm = _FakeSpecAlgorithm()
|
||||
|
||||
def batch_size(self) -> int:
|
||||
return len(self.reqs)
|
||||
|
||||
|
||||
class _TrimmingProcessor(SchedulerBatchResultProcessor):
|
||||
"""Runs the real result processor with GPU/IO-bound helpers stubbed."""
|
||||
|
||||
def _normalize_decode_outputs(
|
||||
self, *, batch, result, logits_output, next_token_ids
|
||||
):
|
||||
return result.test_next_token_ids, result.test_next_token_logprobs
|
||||
|
||||
def _mamba_prefix_cache_update(self, req, batch, result, i):
|
||||
pass
|
||||
|
||||
def _handle_finish_state_updated_req(self, req, batch, result, i, logits_output):
|
||||
pass
|
||||
|
||||
|
||||
def _make_processor() -> _TrimmingProcessor:
|
||||
metrics_reporter = SimpleNamespace(
|
||||
num_generated_tokens=0,
|
||||
forward_ct_decode=0,
|
||||
update_spec_metrics=lambda *a, **k: None,
|
||||
report_decode_stats=lambda *a, **k: None,
|
||||
)
|
||||
allocator = SimpleNamespace(
|
||||
free_group_begin=lambda: None,
|
||||
free_group_end=lambda: None,
|
||||
)
|
||||
output_streamer = SimpleNamespace(stream_output=lambda *a, **k: None)
|
||||
return _TrimmingProcessor(
|
||||
def _make_processor() -> SchedulerBatchResultProcessor:
|
||||
return SchedulerBatchResultProcessor(
|
||||
is_generation=True,
|
||||
disaggregation_mode=None,
|
||||
enable_overlap=False,
|
||||
enable_overlap_mlx=False,
|
||||
server_args=SimpleNamespace(enable_metrics=False),
|
||||
model_config=SimpleNamespace(think_end_id=None),
|
||||
token_to_kv_pool_allocator=allocator,
|
||||
token_to_kv_pool_allocator=None,
|
||||
tree_cache=None,
|
||||
hisparse_coordinator=None,
|
||||
req_to_token_pool=None,
|
||||
decode_offload_manager=None,
|
||||
metrics_collector=None,
|
||||
metrics_reporter=metrics_reporter,
|
||||
metrics_reporter=SimpleNamespace(),
|
||||
draft_worker=None,
|
||||
model_worker=None,
|
||||
model_worker=SimpleNamespace(on_verify_complete_cpu=lambda *a, **k: None),
|
||||
logprob_result_processor=None,
|
||||
output_streamer=output_streamer,
|
||||
output_streamer=SimpleNamespace(),
|
||||
abort_request=lambda *a, **k: None,
|
||||
)
|
||||
|
||||
|
||||
def _make_result(accept_tokens, logprobs):
|
||||
def _make_req(terminate_after: int) -> Req:
|
||||
sp = SamplingParams(max_new_tokens=256, temperature=0)
|
||||
sp.normalize(None)
|
||||
req = Req(
|
||||
rid="r0",
|
||||
origin_input_text="",
|
||||
origin_input_ids=[1, 2, 3],
|
||||
sampling_params=sp,
|
||||
)
|
||||
req.grammar = _FakeGrammar(terminate_after=terminate_after)
|
||||
req.kv_committed_len = 0
|
||||
return req
|
||||
|
||||
|
||||
def _make_result(num_draft_tokens, accept_lens, flat_tokens):
|
||||
return SimpleNamespace(
|
||||
copy_done=None,
|
||||
routed_experts_output=None,
|
||||
indexer_topk_output=None,
|
||||
logits_output=SimpleNamespace(hidden_states=None, customized_info=None),
|
||||
next_token_ids=None,
|
||||
can_run_cuda_graph=False,
|
||||
num_correct_drafts=len(accept_tokens),
|
||||
test_next_token_ids=[list(accept_tokens)],
|
||||
test_next_token_logprobs=[list(logprobs)],
|
||||
next_token_ids=torch.tensor(flat_tokens, dtype=torch.long),
|
||||
accept_lens=torch.tensor(accept_lens, dtype=torch.long),
|
||||
speculative_num_draft_tokens=num_draft_tokens,
|
||||
num_correct_drafts=None,
|
||||
num_correct_drafts_per_req_cpu=None,
|
||||
)
|
||||
|
||||
|
||||
class TestSpecV2GrammarTrimming(CustomTestCase):
|
||||
def _make_req(self, terminate_after: int) -> Req:
|
||||
sp = SamplingParams(max_new_tokens=256, temperature=0)
|
||||
sp.normalize(None)
|
||||
req = Req(
|
||||
rid="r0",
|
||||
origin_input_text="",
|
||||
origin_input_ids=[1, 2, 3],
|
||||
sampling_params=sp,
|
||||
)
|
||||
req.vocab_size = 32000
|
||||
req.return_logprob = True
|
||||
req.logprob.output_token_logprobs_val = []
|
||||
req.logprob.output_token_logprobs_idx = []
|
||||
req.grammar = _FakeGrammar(terminate_after=terminate_after)
|
||||
return req
|
||||
|
||||
def test_trims_tokens_after_grammar_completion(self):
|
||||
req = self._make_req(terminate_after=2)
|
||||
class TestSpecV2GrammarTruncation(CustomTestCase):
|
||||
def test_resolve_truncates_after_grammar_completion(self):
|
||||
req = _make_req(terminate_after=2)
|
||||
proc = _make_processor()
|
||||
result = _make_result([101, 102, 103], [-0.1, -0.2, -0.3])
|
||||
req.kv_committed_len = len(req.origin_input_ids) + len(
|
||||
result.test_next_token_ids[0]
|
||||
)
|
||||
batch = _FakeBatch([req], return_logprob=True)
|
||||
# stride=4, accept_len=3 -> proposed [101, 102, 103]; grammar finishes at 102.
|
||||
result = _make_result(4, [3], [101, 102, 103, 0])
|
||||
|
||||
proc.process_batch_result_decode(batch, result)
|
||||
predict_tokens = proc._resolve_spec_v2_tokens(result, _FakeBatch([req]))
|
||||
|
||||
self.assertTrue(req.finished())
|
||||
self.assertEqual(list(req.output_ids), [101, 102])
|
||||
self.assertEqual(result.test_next_token_ids[0], [101, 102])
|
||||
self.assertEqual(req.grammar.accepted, [101, 102])
|
||||
self.assertTrue(req.grammar.finished)
|
||||
self.assertEqual(req.logprob.output_token_logprobs_val, [-0.1, -0.2])
|
||||
self.assertEqual(req.logprob.output_token_logprobs_idx, [101, 102])
|
||||
self.assertEqual(req.kv_committed_len, len(req.origin_input_ids) + 2)
|
||||
self.assertEqual(predict_tokens, [[101, 102]])
|
||||
# EAGLE commits (retained - 1): prepare_for_decode pre-claimed the bonus
|
||||
# slot, and the dropped suffix is never committed.
|
||||
self.assertEqual(req.kv_committed_len, 2 - 1)
|
||||
|
||||
def test_keeps_all_tokens_when_grammar_not_terminated(self):
|
||||
req = self._make_req(terminate_after=99)
|
||||
def test_resolve_keeps_all_when_grammar_not_terminated(self):
|
||||
req = _make_req(terminate_after=99)
|
||||
proc = _make_processor()
|
||||
result = _make_result([201, 202, 203], [-0.5, -0.6, -0.7])
|
||||
req.kv_committed_len = len(req.origin_input_ids) + len(
|
||||
result.test_next_token_ids[0]
|
||||
)
|
||||
batch = _FakeBatch([req], return_logprob=True)
|
||||
result = _make_result(4, [3], [201, 202, 203, 0])
|
||||
|
||||
proc.process_batch_result_decode(batch, result)
|
||||
predict_tokens = proc._resolve_spec_v2_tokens(result, _FakeBatch([req]))
|
||||
|
||||
self.assertFalse(req.finished())
|
||||
self.assertEqual(list(req.output_ids), [201, 202, 203])
|
||||
self.assertEqual(result.test_next_token_ids[0], [201, 202, 203])
|
||||
self.assertEqual(req.grammar.accepted, [201, 202, 203])
|
||||
self.assertFalse(req.grammar.finished)
|
||||
self.assertEqual(req.logprob.output_token_logprobs_val, [-0.5, -0.6, -0.7])
|
||||
self.assertEqual(req.logprob.output_token_logprobs_idx, [201, 202, 203])
|
||||
self.assertEqual(req.kv_committed_len, len(req.origin_input_ids) + 3)
|
||||
self.assertEqual(predict_tokens, [[201, 202, 203]])
|
||||
self.assertEqual(req.kv_committed_len, 3 - 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -44,10 +44,6 @@ _RESOLVE = (
|
||||
"managers/scheduler_components/batch_result_processor.py",
|
||||
"SchedulerBatchResultProcessor._resolve_spec_v2_tokens",
|
||||
)
|
||||
_GRAMMAR_ACCEPT = (
|
||||
"managers/scheduler_components/batch_result_processor.py",
|
||||
"SchedulerBatchResultProcessor._accept_spec_v2_grammar_tokens",
|
||||
)
|
||||
_SS = "session/streaming_session.py"
|
||||
_OWNER_SITES = {
|
||||
# non-spec scheduler
|
||||
@@ -66,11 +62,10 @@ _OWNER_SITES = {
|
||||
(*_MIXIN, "kv_allocated_len"): 1,
|
||||
# 3rd resolve mutation: DFLASH settles its full commit_lens here (no
|
||||
# pre-claim in prepare_for_decode, unlike the EAGLE mixin).
|
||||
# Spec grammar truncation commits only the retained (pre-termination) length
|
||||
# here, so the dropped suffix is never over-committed (no later rollback).
|
||||
(*_RESOLVE, "kv_committed_len"): 3,
|
||||
(*_RESOLVE, "spec_verify_ct"): 1,
|
||||
# Spec grammar trim rolls back KV slots for tokens dropped after grammar
|
||||
# completion; resolve already committed the full accepted list.
|
||||
(*_GRAMMAR_ACCEPT, "kv_committed_len"): 1,
|
||||
(
|
||||
"speculative/dflash_info_v2.py",
|
||||
"DFlashDraftInputV2.prepare_for_decode",
|
||||
|
||||
Reference in New Issue
Block a user