[Bugfix][V1] Register accepted spec-decode tokens before terminal cleanup

tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py

@@ -487,7 +487,13 @@ def test_hybrid_mamba_partial_tail_owner_uses_cow_on_continue():
     req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
     computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
     assert num_computed == 0
-    assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
+    assert manager.allocate_slots(req0, 4, num_computed, computed_blocks) is not None
+
+    # Async chunked prefill may publish a new boundary while the preceding
+    # chunk is still in flight. This first registration must remain enabled.
+    req0.num_computed_tokens = 4
+    req0.num_in_flight_tokens = 4
+    assert manager.allocate_slots(req0, 2) is not None
 
     partial_mamba_hash = req0.block_hashes[6 // hash_block_size - 1]
     partial_mamba_block = manager.block_pool.get_cached_block(
@@ -497,8 +503,11 @@ def test_hybrid_mamba_partial_tail_owner_uses_cow_on_continue():
     partial_mamba_block_id = partial_mamba_block[0].block_id
     assert manager.get_blocks("0").get_block_ids()[1][1] == partial_mamba_block_id
 
+    # The next run-ahead allocation moves the boundary to a durable CoW block.
+    # Its end-of-allocation cache pass must not reattach the same hash to the
+    # mutable request-table source.
     req0.num_computed_tokens = 6
-    req0.append_output_token_ids([3])
+    req0.num_in_flight_tokens = 6
     new_blocks = manager.allocate_slots(req0, 1)
     assert new_blocks is not None
 
@@ -522,6 +531,17 @@ def test_hybrid_mamba_partial_tail_owner_uses_cow_on_continue():
     assert get_group_id(moved[0].block_hash) == 1
     assert moved[0].block_hash_num_tokens == 6
 
+    # The output-time cache path can retry the same settled boundary while a
+    # later step remains in flight; that must remain idempotent too.
+    req0.num_in_flight_tokens = 1
+    manager.cache_blocks(req0, 6)
+    grouped_hash = moved[0].block_hash
+    assert grouped_hash is not None
+    cache = manager.block_pool.cached_block_hash_to_block
+    assert cache.contain(grouped_hash, cow_copy.dst_block_id)
+    assert not cache.contain(grouped_hash, partial_mamba_block_id)
+    assert partial_mamba_block[0].block_hash is None
+
 
 def test_partial_hit_then_internal_checkpoint_uses_distinct_mamba_blocks():
     hash_block_size = 2

tests/v1/core/test_async_scheduler.py

@@ -13,7 +13,7 @@
 from vllm.v1.structured_output import StructuredOutputGrammar
 from vllm.v1.utils import ConstantList
 
-from .utils import create_requests, create_scheduler, mock_kv
+from .utils import EOS_TOKEN_ID, create_requests, create_scheduler, mock_kv
 
 pytestmark = pytest.mark.cpu_test
 
@@ -66,6 +66,42 @@ def test_stop_by_max_tokens(max_tokens: int):
     assert total_num_scheduled_tokens == expected_total_num_scheduled_tokens
 
 
+def test_finished_request_does_not_cache_in_flight_tokens():
+    """A terminal output must not publish KV from a later in-flight step."""
+    scheduler = create_scheduler(
+        async_scheduling=True,
+        enable_prefix_caching=True,
+        block_size=16,
+    )
+    (request,) = create_requests(
+        num_requests=1,
+        num_tokens=31,
+        max_tokens=10,
+        same_prompt=True,
+        block_size=16,
+    )
+    scheduler.add_request(request)
+
+    first_step = scheduler.schedule()
+    second_step = scheduler.schedule()
+    assert second_step.num_scheduled_tokens[request.request_id] == 1
+
+    model_output = _make_model_runner_output(first_step)
+    model_output.sampled_token_ids = [[EOS_TOKEN_ID]]
+    scheduler.update_from_output(first_step, model_output)
+    assert request.status == RequestStatus.FINISHED_STOPPED
+    assert request.num_in_flight_tokens == 1
+
+    (repeat,) = create_requests(
+        num_requests=1,
+        num_tokens=32,
+        same_prompt=True,
+        block_size=16,
+    )
+    _, num_cached_tokens, _ = scheduler.kv_cache_manager.get_computed_blocks(repeat)
+    assert num_cached_tokens == 16
+
+
 def test_no_spec_decode_padding_up_to_max_model_len():
     """Uniform spec-decode padding must leave room for the sampled token.
 

tests/v1/core/test_scheduler.py

@@ -1610,6 +1610,87 @@ def test_schedule_spec_decoding_stats(
         assert "per_step_accepted" not in payload  # summary level
 
 
+@pytest.mark.parametrize(
+    (
+        "async_scheduling",
+        "prompt_tokens",
+        "num_reprefillable_tokens",
+        "expected_cached_tokens",
+        "run_spec_verify",
+    ),
+    [
+        (False, 44, 0, 48, True),
+        (False, 30, 2, 16, True),
+        (False, 28, 2, 0, True),
+        (True, 30, 2, 16, True),
+        (False, 47, 0, 32, False),
+    ],
+    ids=[
+        "ordinary-spec",
+        "mtp-publish",
+        "mtp-fence",
+        "async-mtp-publish",
+        "sampled-output-fence",
+    ],
+)
+def test_finished_spec_decode_registers_accepted_tokens(
+    async_scheduling: bool,
+    prompt_tokens: int,
+    num_reprefillable_tokens: int,
+    expected_cached_tokens: int,
+    run_spec_verify: bool,
+):
+    """Terminal registration publishes all and only executed KV."""
+    scheduler = create_scheduler(
+        enable_prefix_caching=True,
+        block_size=16,
+        num_speculative_tokens=3,
+        speculative_method="ngram_gpu" if async_scheduling else None,
+        async_scheduling=async_scheduling,
+    )
+    if num_reprefillable_tokens:
+        # Mirror the cache-registration and lookup geometry of K=3
+        # multi-module MTP without requiring model weights or a GPU.
+        scheduler.use_eagle = True
+        scheduler.use_eagle_block_drop = True
+        scheduler.num_prefill_lookahead = num_reprefillable_tokens + 1
+        coordinator = scheduler.kv_cache_manager.coordinator
+        coordinator.num_reprefillable_tokens = num_reprefillable_tokens
+        coordinator.eagle_group_ids = {0}
+        coordinator.single_type_managers[0].use_eagle = True
+
+    (request,) = create_requests(
+        num_requests=1,
+        num_tokens=prompt_tokens,
+        max_tokens=4 if run_spec_verify else 1,
+        block_size=16,
+        same_prompt=True,
+    )
+    scheduler.add_request(request)
+
+    # Prefill samples one token. Most cases then accept all three drafts on a
+    # terminal verify; the verifier's bonus token is trimmed at max_tokens=4.
+    _model_output(scheduler, scheduler.schedule(), [[0]])
+    if run_spec_verify:
+        scheduler.update_draft_token_ids(
+            DraftTokenIds([request.request_id], [[0, 0, 0]])
+        )
+        _model_output(scheduler, scheduler.schedule(), [[0, 0, 0, 0]])
+
+    # One extra token keeps the lookup's required uncomputed tail outside the
+    # finalized boundary whose reuse this test measures.
+    (continuation,) = create_requests(
+        num_requests=1,
+        num_tokens=prompt_tokens + 5,
+        block_size=16,
+        same_prompt=True,
+    )
+    _, num_cached_tokens, _ = scheduler.kv_cache_manager.get_computed_blocks(
+        continuation
+    )
+    assert num_cached_tokens == expected_cached_tokens
+
+
 def _run_spec_verify_steps(scheduler, rounds, num_invalid_per_round=None):
     """Drive prefill + one draft/verify step per round for a single request.
 

vllm/v1/core/sched/scheduler.py

@@ -2137,6 +2137,21 @@ def update_from_output(
                 finish_reason = request.get_finished_reason()
                 finished = self._handle_stopped_request(request)
                 if finished:
+                    if status_before_stop == RequestStatus.RUNNING:
+                        # Register finalized output tokens before releasing
+                        # their blocks. A terminal request has no later
+                        # allocation pass to do so.
+                        finalized_num_computed_tokens = min(
+                            max(
+                                0,
+                                request.num_computed_tokens
+                                - request.num_in_flight_tokens,
+                            ),
+                            request.num_tokens,
+                        )
+                        self.kv_cache_manager.cache_blocks(
+                            request, finalized_num_computed_tokens
+                        )
                     kv_transfer_params, ec_transfer_params = self._free_request(request)
 
                 if status_before_stop == RequestStatus.RUNNING:

vllm/v1/core/single_type_kv_cache_manager.py

@@ -2039,6 +2039,17 @@ def _cache_partial_tail_block(
         request: Request,
         num_tokens: int,
     ) -> BlockHashWithGroupId | None:
+        # With async run-ahead, a boundary at or behind the optimistic
+        # computed frontier may already have been moved to a durable CoW
+        # block. Do not reattach its hash to the request-table state block,
+        # which a later in-flight step can overwrite. A boundary beyond the
+        # frontier is a first publication by the current allocation and must
+        # remain eligible.
+        if (
+            request.num_in_flight_tokens > 0
+            and num_tokens <= request.num_computed_tokens
+        ):
+            return None
         hash_block_size = self.block_pool.hash_block_size
         # Re-key the reserved block at its exported checkpoint boundary.
         checkpoint = self._checkpoints.get(request.request_id)