diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index 77221cd37..a9d5f0c28 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -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, diff --git a/python/sglang/test/kits/spec_server_kits.py b/python/sglang/test/kits/spec_server_kits.py index 89997d178..e7c792ae2 100644 --- a/python/sglang/test/kits/spec_server_kits.py +++ b/python/sglang/test/kits/spec_server_kits.py @@ -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"]) diff --git a/test/registered/disaggregation/test_disaggregation_basic.py b/test/registered/disaggregation/test_disaggregation_basic.py index ca9e250c3..2c6d50415 100644 --- a/test/registered/disaggregation/test_disaggregation_basic.py +++ b/test/registered/disaggregation/test_disaggregation_basic.py @@ -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): diff --git a/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py b/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py index ac1949c95..131d33b0e 100644 --- a/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py +++ b/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py @@ -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__": diff --git a/test/registered/unit/spec/test_decode_bookkeeping_ownership.py b/test/registered/unit/spec/test_decode_bookkeeping_ownership.py index a0f83dd15..b17b191ad 100644 --- a/test/registered/unit/spec/test_decode_bookkeeping_ownership.py +++ b/test/registered/unit/spec/test_decode_bookkeeping_ownership.py @@ -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",