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

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(
+            accept_tokens = next_token_ids[i * stride : i * stride + accept_lens[i]]
-                next_token_ids[i * stride : i * stride + accept_lens[i]]
-            )
 
             if req.is_retracted:
                 # reset_for_retract() already zeroes committed/allocated KV.
-                continue
+                pass
-
+            elif req.finished():
-            if req.finished():
                 if not batch.spec_algorithm.is_dflash():
                     # EAGLE prepare_for_decode pre-claimed the bonus slot.
                     req.kv_committed_len -= 1
-                continue
+            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)
 
+                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 += accept_lens[i]
+                    req.kv_committed_len += num_accept_tokens
                 else:
                     # EAGLE prepare_for_decode pre-claimed the bonus slot.
-                req.kv_committed_len += accept_lens[i] - 1
+                    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
+            is_spec = not batch.spec_algorithm.is_none()
-            needs_tokenwise_grammar_accept = (
+
-                not batch.spec_algorithm.is_none() and req.grammar is not None
+            if not is_spec:
-            )
+                # Normal decode: a single sampled token.
-            if batch.spec_algorithm.is_none():
                 req.output_ids.append(next_token_id)
-            elif needs_tokenwise_grammar_accept:
+                new_accept_len = 1
-                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)
             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)
+                new_accept_len = len(next_token_id)
 
-            if not needs_tokenwise_grammar_accept:
             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_accept_len)
-                req.update_finish_state(new_accepted_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:
+                if not is_spec:
-                    # Already advanced token-by-token above; just sync terminal flag.
+                    # 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()
-                else:
-                    self._apply_decode_grammar(
-                        req=req, next_token_id=next_token_id, batch=batch
-                    )
 
         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

@@ -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"])

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):

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,82 +39,40 @@ 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)
 
-
+def _make_processor() -> SchedulerBatchResultProcessor:
-class _TrimmingProcessor(SchedulerBatchResultProcessor):
+    return 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(
         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:
-    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)],
-    )
-
-
-class TestSpecV2GrammarTrimming(CustomTestCase):
-    def _make_req(self, terminate_after: int) -> Req:
     sp = SamplingParams(max_new_tokens=256, temperature=0)
     sp.normalize(None)
     req = Req(
@@ -117,52 +81,44 @@ class TestSpecV2GrammarTrimming(CustomTestCase):
         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)
+    req.kv_committed_len = 0
     return req
 
-    def test_trims_tokens_after_grammar_completion(self):
+
-        req = self._make_req(terminate_after=2)
+def _make_result(num_draft_tokens, accept_lens, flat_tokens):
-        proc = _make_processor()
+    return SimpleNamespace(
-        result = _make_result([101, 102, 103], [-0.1, -0.2, -0.3])
+        next_token_ids=torch.tensor(flat_tokens, dtype=torch.long),
-        req.kv_committed_len = len(req.origin_input_ids) + len(
+        accept_lens=torch.tensor(accept_lens, dtype=torch.long),
-            result.test_next_token_ids[0]
+        speculative_num_draft_tokens=num_draft_tokens,
+        num_correct_drafts=None,
+        num_correct_drafts_per_req_cpu=None,
     )
-        batch = _FakeBatch([req], return_logprob=True)
 
-        proc.process_batch_result_decode(batch, result)
 
-        self.assertTrue(req.finished())
+class TestSpecV2GrammarTruncation(CustomTestCase):
-        self.assertEqual(list(req.output_ids), [101, 102])
+    def test_resolve_truncates_after_grammar_completion(self):
-        self.assertEqual(result.test_next_token_ids[0], [101, 102])
+        req = _make_req(terminate_after=2)
-        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)
-
-    def test_keeps_all_tokens_when_grammar_not_terminated(self):
-        req = self._make_req(terminate_after=99)
         proc = _make_processor()
-        result = _make_result([201, 202, 203], [-0.5, -0.6, -0.7])
+        # stride=4, accept_len=3 -> proposed [101, 102, 103]; grammar finishes at 102.
-        req.kv_committed_len = len(req.origin_input_ids) + len(
+        result = _make_result(4, [3], [101, 102, 103, 0])
-            result.test_next_token_ids[0]
-        )
-        batch = _FakeBatch([req], return_logprob=True)
 
-        proc.process_batch_result_decode(batch, result)
+        predict_tokens = proc._resolve_spec_v2_tokens(result, _FakeBatch([req]))
 
-        self.assertFalse(req.finished())
+        self.assertEqual(predict_tokens, [[101, 102]])
-        self.assertEqual(list(req.output_ids), [201, 202, 203])
+        # EAGLE commits (retained - 1): prepare_for_decode pre-claimed the bonus
-        self.assertEqual(result.test_next_token_ids[0], [201, 202, 203])
+        # slot, and the dropped suffix is never committed.
-        self.assertEqual(req.grammar.accepted, [201, 202, 203])
+        self.assertEqual(req.kv_committed_len, 2 - 1)
-        self.assertFalse(req.grammar.finished)
+
-        self.assertEqual(req.logprob.output_token_logprobs_val, [-0.5, -0.6, -0.7])
+    def test_resolve_keeps_all_when_grammar_not_terminated(self):
-        self.assertEqual(req.logprob.output_token_logprobs_idx, [201, 202, 203])
+        req = _make_req(terminate_after=99)
-        self.assertEqual(req.kv_committed_len, len(req.origin_input_ids) + 3)
+        proc = _make_processor()
+        result = _make_result(4, [3], [201, 202, 203, 0])
+
+        predict_tokens = proc._resolve_spec_v2_tokens(result, _FakeBatch([req]))
+
+        self.assertEqual(predict_tokens, [[201, 202, 203]])
+        self.assertEqual(req.kv_committed_len, 3 - 1)
 
 
 if __name__ == "__main__":

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",