From b5c7dff954ac353a8acbbc97e2fc5519d036fea4 Mon Sep 17 00:00:00 2001
From: zjy0516 <riverclouds.zhu@qq.com>
Date: Mon, 25 May 2026 05:51:39 +0000
Subject: [PATCH 1/4] init

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
---
 tests/v1/core/test_eager_cache_zombie.py | 332 +++++++++++++++++++++++
 tests/v1/core/test_prefix_caching.py     |  11 +
 vllm/v1/core/block_pool.py               |  31 +++
 vllm/v1/core/kv_cache_manager.py         |   9 +
 vllm/v1/core/sched/scheduler.py          |   3 +
 5 files changed, 386 insertions(+)
 create mode 100644 tests/v1/core/test_eager_cache_zombie.py

diff --git a/tests/v1/core/test_eager_cache_zombie.py b/tests/v1/core/test_eager_cache_zombie.py
new file mode 100644
index 000000000000..3674b4cd1909
--- /dev/null
+++ b/tests/v1/core/test_eager_cache_zombie.py
@@ -0,0 +1,332 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Tests for the eager-cache-registration zombie problem.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from tests.v1.core.test_prefix_caching import (
+    make_kv_cache_config,
+    make_kv_cache_config_hybrid_model,
+    make_request,
+)
+from tests.v1.core.utils import create_scheduler
+from vllm.utils.hashing import sha256
+from vllm.v1.core.kv_cache_manager import KVCacheManager
+from vllm.v1.core.kv_cache_utils import init_none_hash
+from vllm.v1.core.sched.request_queue import (
+    SchedulingPolicy,
+    create_request_queue,
+)
+from vllm.v1.request import RequestStatus
+
+pytestmark = pytest.mark.cpu_test
+
+
+@pytest.fixture(autouse=True)
+def _init_hash():
+    init_none_hash(sha256)
+
+
+# ---------------------------------------------------------------------------
+# Test 1: full-attention single-group case
+# ---------------------------------------------------------------------------
+
+
+def test_rollback_on_preempt_before_worker_write():
+    """
+    Schedule req_A then ``free`` it immediately (simulating preempt/abort
+    before the worker has executed). The eager-registered hash entries
+    must be evicted from the cache map and the blocks' ``block_hash``
+    must be reset, so that a later req_B with the same prefix does NOT
+    cache-hit on never-written blocks.
+    """
+    block_size = 16
+    manager = KVCacheManager(
+        make_kv_cache_config(block_size, 11),
+        max_model_len=8192,
+        enable_caching=True,
+        hash_block_size=block_size,
+    )
+
+    # --- req_A: 32 tokens = 2 full blocks --------------------------------
+    common_token_ids = [i for i in range(2) for _ in range(block_size)]
+    req_a = make_request("a", common_token_ids, block_size, sha256)
+    computed_a, num_computed_a = manager.get_computed_blocks(req_a)
+    assert num_computed_a == 0  # cold cache
+    blocks_a = manager.allocate_slots(req_a, 32, 0, computed_a)
+    assert blocks_a is not None
+
+    # After allocate_slots: 2 hashes registered eagerly and both tracked
+    # in ``_uncommitted[req_a]``.
+    cache_map = manager.block_pool.cached_block_hash_to_block._cache
+    assert len(cache_map) == 2, (
+        "Sanity: 2 full blocks should each register a hash on allocate_slots"
+    )
+    assert len(manager.block_pool._uncommitted["a"]) == 2, (
+        "Sanity: both eager registrations should be tracked as uncommitted "
+        "until commit_step or rollback runs."
+    )
+
+    block_ids_a = blocks_a.get_block_ids()[0]
+    eager_blocks = [manager.block_pool.blocks[bid] for bid in block_ids_a]
+
+    # --- Simulate preempt/abort BEFORE worker writes K/V bytes -----------
+    # In real life the worker would normally run between allocate_slots and
+    # the next free(). Here we go directly from allocate_slots → free
+    # *without* an intervening ``commit_step``, mirroring the race window
+    # where the request is preempted while the worker has not yet executed
+    # for this step.
+    manager.free(req_a)
+
+    # --- Rollback evicts the uncommitted entries -------------------------
+    assert len(cache_map) == 0, (
+        "free(uncommitted req) must evict both eager hash entries; got "
+        f"{len(cache_map)} left in cache map."
+    )
+    assert "a" not in manager.block_pool._uncommitted, (
+        "_uncommitted[req_a] must be cleared by rollback."
+    )
+    for blk in eager_blocks:
+        assert blk.ref_cnt == 0, "free() did decrement ref_cnt"
+        assert blk.block_hash is None, (
+            "block.block_hash must be reset by rollback path."
+        )
+
+    # --- req_B: same prefix → cache MISS ---------------------------------
+    # 48 tokens: first 32 would match req_a's now-rolled-back hashes, last
+    # 16 are new. Without zombies, this is a full cold miss.
+    req_b_tokens = common_token_ids + [99] * block_size
+    req_b = make_request("b", req_b_tokens, block_size, sha256)
+    _, num_computed_b = manager.get_computed_blocks(req_b)
+
+    assert num_computed_b == 0, (
+        f"req_b must not cache-hit on rolled-back entries; got "
+        f"{num_computed_b} cached tokens, expected 0."
+    )
+
+
+# ---------------------------------------------------------------------------
+# Test 2: hybrid (full-attention + Mamba) case
+# ---------------------------------------------------------------------------
+
+
+def test_rollback_on_preempt_for_mamba_hybrid():
+    """
+    Same race as test 1, but for full-attention + 2 Mamba groups. Confirms
+    the rollback path covers every manager that calls ``cache_blocks``
+    during ``allocate_slots``, not just full-attention.
+    """
+    block_size = 16
+    # 1 full-attention group + 2 Mamba groups (slice 0/1).
+    manager = KVCacheManager(
+        make_kv_cache_config_hybrid_model(
+            block_size,
+            num_blocks=16,
+            sliding_window_blocks=0,
+            second_spec_type="mamba",
+        ),
+        max_model_len=8192,
+        enable_caching=True,
+        hash_block_size=block_size,
+    )
+
+    common_token_ids = [i for i in range(2) for _ in range(block_size)]
+    req_a = make_request("a", common_token_ids, block_size, sha256)
+    computed_a, _ = manager.get_computed_blocks(req_a)
+    blocks_a = manager.allocate_slots(req_a, 32, 0, computed_a)
+    assert blocks_a is not None
+
+    cache_map = manager.block_pool.cached_block_hash_to_block._cache
+    eager_entry_count = len(cache_map)
+    assert eager_entry_count > 0, (
+        "Sanity: hybrid manager registers at least some eager entries on alloc"
+    )
+    assert len(manager.block_pool._uncommitted["a"]) == eager_entry_count, (
+        "Sanity: every eager registration across all groups should be tracked "
+        "as uncommitted."
+    )
+
+    # Preempt before worker write.
+    manager.free(req_a)
+
+    assert len(cache_map) == 0, (
+        f"Rollback must clear all groups' eager entries; got "
+        f"{len(cache_map)} left in the cache map."
+    )
+    assert "a" not in manager.block_pool._uncommitted
+
+    # Future request cache MISS — no entries left to match.
+    req_b = make_request("b", common_token_ids, block_size, sha256)
+    _, num_computed_b = manager.get_computed_blocks(req_b)
+    assert num_computed_b == 0, (
+        f"req_b must not cache-hit on rolled-back entries; got "
+        f"{num_computed_b} cached tokens, expected 0."
+    )
+
+
+# ---------------------------------------------------------------------------
+# Test 3: control case — successful step + free should keep cache (no bug)
+# ---------------------------------------------------------------------------
+
+
+def test_commit_step_keeps_cache_across_normal_free():
+    """
+    Control: when a request runs the worker to completion (modelled here
+    by an explicit ``commit_step()`` call) and is then freed normally,
+    the cache entries SHOULD remain. ``commit_step`` clears
+    ``_uncommitted`` so the subsequent ``free`` does not roll anything
+    back.
+
+    Without the ``commit_step`` call, ``free`` would correctly roll the
+    entries back — that's exactly the eager-rollback behavior exercised
+    by tests 1 and 2 above. This test pins down the other side: the
+    commit hook is what preserves cache hits across normal completion.
+    """
+    block_size = 16
+    manager = KVCacheManager(
+        make_kv_cache_config(block_size, 11),
+        max_model_len=8192,
+        enable_caching=True,
+        hash_block_size=block_size,
+    )
+
+    common_token_ids = [i for i in range(2) for _ in range(block_size)]
+    req_a = make_request("a", common_token_ids, block_size, sha256)
+    computed_a, _ = manager.get_computed_blocks(req_a)
+    manager.allocate_slots(req_a, 32, 0, computed_a)
+    assert len(manager.block_pool._uncommitted["a"]) == 2
+
+    # Simulate successful worker execution by committing the step.
+    # Scheduler.schedule() calls this at the start of every step; here we
+    # call it manually to model "previous step's worker has confirmed".
+    manager.commit_step()
+    assert manager.block_pool._uncommitted == {}, (
+        "commit_step should clear all pending uncommitted registrations."
+    )
+
+    manager.free(req_a)
+
+    cache_map = manager.block_pool.cached_block_hash_to_block._cache
+    assert len(cache_map) == 2, (
+        "Normal completion (commit_step + free) should leave the cache "
+        f"entries intact, got {len(cache_map)} entries."
+    )
+
+
+# ---------------------------------------------------------------------------
+# Test 4: scheduler-level preemption test
+# ---------------------------------------------------------------------------
+
+
+def test_scheduler_preempt_rolls_back_target_step_eager_cache():
+    """
+    Exercise the rollback through a real ``Scheduler.schedule()`` call.
+
+    Setup seeds two RUNNING chunked-prefill requests with one already-committed
+    block each. In the target schedule call, req_A is scheduled first and
+    eagerly caches two new full blocks. req_B then cannot allocate, so priority
+    preemption removes req_A from the same scheduler output before any worker
+    can see or execute that work.
+
+    The two target-step entries must be evicted on preempt; the prior-step
+    seeded block remains in the cache map (it was promoted to committed when
+    ``schedule`` called ``commit_step`` at the start of this step). A future
+    req_C with req_A's first two blocks as prefix therefore hits exactly the
+    one seeded block.
+    """
+    block_size = 16
+    # Resource budget is intentionally tight to force the exact race:
+    #   max_num_batched_tokens=48 = 32 (req_a's 2 remaining blocks)
+    #                             + 16 (req_b's 1 remaining block).
+    #     Both requests *want* to advance fully in the target step.
+    #   num_blocks=5 = 1 null
+    #                + 1 seeded committed block for req_a
+    #                + 1 seeded committed block for req_b
+    #                + 2 free blocks.
+    #     The 2 free blocks are exactly enough for req_a's target-step
+    #     allocation (2 new full blocks); req_b then has 0 free → triggers
+    #     preempt of the lowest-priority running request, which is req_a.
+    scheduler = create_scheduler(
+        max_num_seqs=3,
+        max_num_batched_tokens=48,
+        max_model_len=8192,
+        enable_prefix_caching=True,
+        num_blocks=5,
+        block_size=block_size,
+    )
+    # PRIORITY policy is required: under FCFS req_a would not be preempted
+    # for req_b, and the race window we want to test would not open.
+    # ``create_scheduler`` does not expose ``policy`` directly, so swap it
+    # post-construction and rebuild the queues that depend on the policy.
+    scheduler.policy = SchedulingPolicy.PRIORITY
+    scheduler.waiting = create_request_queue(scheduler.policy)
+    scheduler.skipped_waiting = create_request_queue(scheduler.policy)
+
+    # 48-token prompt = 3 full blocks; 1 seeded as committed, 2 to be
+    # eagerly cached in the target step (these become the zombies).
+    tokens_a = [0] * block_size + [1] * block_size + [2] * block_size
+    # 32-token prompt = 2 full blocks; 1 seeded as committed, 1 needed in
+    # the target step (this is what fails to allocate and triggers preempt).
+    tokens_b = [10] * block_size + [11] * block_size
+    req_a = make_request("a", tokens_a, block_size, sha256)
+    req_b = make_request("b", tokens_b, block_size, sha256)
+    # In vLLM priority semantics, smaller value = higher priority. req_b
+    # outranks req_a, so when free blocks run out, req_a is the one
+    # preempted to make room for req_b.
+    req_a.priority = 5
+    req_b.priority = 0
+    req_a.arrival_time = 1.0
+    req_b.arrival_time = 2.0
+
+    # Plant both requests directly in RUNNING with their first block already
+    # computed. This bypasses the normal admission path so we start the test
+    # mid-chunked-prefill (the only state where this race opens) without
+    # having to drive multiple scheduler steps to set it up.
+    manager = scheduler.kv_cache_manager
+    for req in (req_a, req_b):
+        seeded_blocks = manager.allocate_slots(req, block_size)
+        assert seeded_blocks is not None
+        req.num_computed_tokens = block_size
+        req.status = RequestStatus.RUNNING
+        scheduler.requests[req.request_id] = req
+    scheduler.running = [req_a, req_b]
+
+    # Sanity: post-seeding state matches the budget plan above.
+    cache_map = manager.block_pool.cached_block_hash_to_block._cache
+    assert len(cache_map) == 2  # one committed hash per seeded request.
+    assert manager.block_pool.get_num_free_blocks() == 2  # exactly the race window.
+
+    output = scheduler.schedule()
+
+    assert output.preempted_req_ids == {"a"}
+    assert "a" not in output.num_scheduled_tokens, (
+        "Sanity: req_a was removed from the scheduler output before worker "
+        "execution, so its target-step KV writes cannot happen."
+    )
+    assert output.num_scheduled_tokens == {"b": block_size}
+    assert req_a.status == RequestStatus.PREEMPTED
+
+    # A future request with A's first two blocks as a prefix should hit
+    # exactly one block: the prior-step seeded block (committed via
+    # commit_step at the start of this scheduler step). The target-step
+    # eager entry for the would-be second block was rolled back on preempt.
+    req_c = make_request(
+        "c",
+        tokens_a[: 2 * block_size] + [99] * block_size,
+        block_size,
+        sha256,
+    )
+    computed_c, num_computed_c = manager.get_computed_blocks(req_c)
+    assert num_computed_c == block_size, (
+        "Scheduler-level preemption must roll back the target-step eager "
+        "cache entry; only the prior-step committed block should hit. Got "
+        f"{num_computed_c} cached tokens, expected {block_size}."
+    )
+    hit_blocks = computed_c.blocks[0]
+    assert len(hit_blocks) == 1, (
+        f"Expected 1 hit (seeded committed block), got {len(hit_blocks)}."
+    )
diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py
index 546412b1d2f8..32d4ab08bbf9 100644
--- a/tests/v1/core/test_prefix_caching.py
+++ b/tests/v1/core/test_prefix_caching.py
@@ -269,6 +269,9 @@ def test_prefill(hash_fn):
     free_block_queue = manager.block_pool.free_block_queue
     assert free_block_queue.num_free_blocks == 5
 
+    # Model successful worker execution before freeing so the eager-rollback
+    # path does not evict the committed cache entries.
+    manager.commit_step()
     manager.free(req0)
     manager.free(req1)
 
@@ -400,6 +403,7 @@ def test_prefill_hybrid_model():
                 assert block.ref_cnt == 2
 
     block_hashes = req1.block_hashes
+    manager.commit_step()
     manager.free(req0)
     manager.free(req1)
 
@@ -579,6 +583,7 @@ def test_prefill_hybrid_model_eagle():
                 assert block.ref_cnt == 2
 
     block_hashes = req1.block_hashes
+    manager.commit_step()
     manager.free(req0)
     manager.free(req1)
 
@@ -1222,6 +1227,7 @@ def test_evict():
     # 10 - (6 + 3) == 1
     assert manager.block_pool.free_block_queue.num_free_blocks == 1
 
+    manager.commit_step()
     manager.free(req0)
     manager.free(req1)
     assert manager.block_pool.free_block_queue.num_free_blocks == 10
@@ -1321,6 +1327,7 @@ def test_computed_blocks_not_evicted():
     assert blocks.blocks[0][0].block_id == 2
 
     # Free the blocks.
+    manager.commit_step()
     manager.free(req0)
     manager.free(req1)
 
@@ -1754,6 +1761,7 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():
     ]
     # | Common-0 | Common-1 | Common-2 | Req1-3 (F) | Req1-4 (F) |
     # | Req1-5(F)| ... |
+    manager.commit_step()
     manager.free(req1)
     assert {block.ref_cnt for block in block_part1[:3]} == {1}
     assert {block.ref_cnt for block in block_part1[3:]} == {0}
@@ -2345,6 +2353,7 @@ def test_eagle_enabled_removes_last_block():
     manager.allocate_slots(
         req, len(token_ids), len(computed_blocks.blocks[0]) * 16, computed_blocks
     )
+    manager.commit_step()
     manager.free(req)
 
     # New request with same tokens + Eagle enabled
@@ -2377,6 +2386,7 @@ def test_eagle_with_partial_blocks():
     manager.allocate_slots(
         req, len(token_ids), len(computed_blocks.blocks[0]) * 16, computed_blocks
     )
+    manager.commit_step()
     manager.free(req)
 
     # New request with Eagle enabled
@@ -2421,6 +2431,7 @@ def test_eagle_with_sliding_window():
     # record the block hash of the first block in the request for later use
     block_hash_first_block = req.block_hashes[0]
     assert block_hash_first_block is not None
+    manager.commit_step()
     manager.free(req)
 
     # New request with Eagle enabled
diff --git a/vllm/v1/core/block_pool.py b/vllm/v1/core/block_pool.py
index 513e4bf380b9..7c56bcbea49c 100644
--- a/vllm/v1/core/block_pool.py
+++ b/vllm/v1/core/block_pool.py
@@ -181,6 +181,11 @@ def __init__(
 
         self.metrics_collector = metrics_collector
 
+        # Blocks whose hashes were eagerly registered this scheduler step but
+        # whose K/V bytes are not yet worker-confirmed. ``commit_step`` clears
+        # at step boundaries; ``rollback_uncommitted`` evicts on early free.
+        self._uncommitted: dict[str, list[KVCacheBlock]] = {}
+
     def get_cached_block(
         self, block_hash: BlockHash, kv_cache_group_ids: list[int]
     ) -> list[KVCacheBlock] | None:
@@ -264,6 +269,7 @@ def cache_full_blocks(
         new_hashes: list[ExternalBlockHash] | None = (
             [] if self.enable_kv_cache_events else None
         )
+        uncommitted_for_req: list[KVCacheBlock] | None = None
         for i, blk in enumerate(new_full_blocks):
             # Some blocks may be null or masked out when enabling sparse attention
             # like sliding window attention, or Mamba models with prefix-caching
@@ -282,6 +288,13 @@ def cache_full_blocks(
             if new_hashes is not None:
                 new_hashes.append(maybe_convert_block_hash(block_hash))
 
+            # Track for commit_step / rollback_uncommitted. Lazy list alloc.
+            if uncommitted_for_req is None:
+                uncommitted_for_req = self._uncommitted.setdefault(
+                    request.request_id, []
+                )
+            uncommitted_for_req.append(blk)
+
         if self.enable_kv_cache_events:
             if num_cached_blocks == 0:
                 parent_block_hash: ExternalBlockHash | None = None
@@ -432,6 +445,22 @@ def free_blocks(self, ordered_blocks: Iterable[KVCacheBlock]) -> None:
             [block for block in blocks_list if block.ref_cnt == 0 and not block.is_null]
         )
 
+    def rollback_uncommitted(self, request_id: str) -> int:
+        """Evict ``request_id``'s eager-registered-but-not-yet-worker-confirmed
+        cache entries. Returns the number evicted. Idempotent.
+        """
+        blocks = self._uncommitted.pop(request_id, None)
+        if not blocks:
+            return 0
+        return sum(self._maybe_evict_cached_block(b) for b in blocks)
+
+    def commit_step(self) -> None:
+        """Promote all pending eager registrations to committed. Called at
+        scheduler step boundaries. Idempotent.
+        """
+        if self._uncommitted:
+            self._uncommitted.clear()
+
     def evict_blocks(self, block_ids: set[int]) -> None:
         """evict blocks from the prefix cache by their block IDs.
 
@@ -472,6 +501,8 @@ def reset_prefix_cache(self) -> bool:
         # Remove all hashes so that no new blocks will hit.
         self.cached_block_hash_to_block = BlockHashToBlockMap()
 
+        self._uncommitted.clear()
+
         # Remove all hashes from all blocks.
         for block in self.blocks:
             block.reset_hash()
diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py
index 9359d8843a91..50e14fd551dd 100644
--- a/vllm/v1/core/kv_cache_manager.py
+++ b/vllm/v1/core/kv_cache_manager.py
@@ -434,6 +434,9 @@ def free(self, request: Request) -> None:
         Args:
             request: The request to free the blocks.
         """
+        # Roll back uncommitted eager cache entries (no-op for cross-step
+        # free; cleans up zombies for preempt/abort mid-step).
+        self.block_pool.rollback_uncommitted(request.request_id)
         self.coordinator.free(request.request_id)
 
     def remove_skipped_blocks(
@@ -568,3 +571,9 @@ def take_new_block_ids(self) -> list[int]:
     def new_step_starts(self) -> None:
         """Called when a new step is started."""
         self.coordinator.new_step_starts()
+
+    def commit_step(self) -> None:
+        """Promote eager cache registrations from the previous scheduler step
+        to committed. Called at step boundaries. Idempotent.
+        """
+        self.block_pool.commit_step()
diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
index c69c9a8119ab..b11f16512803 100644
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -359,6 +359,9 @@ def schedule(self) -> SchedulerOutput:
         # For logging.
         scheduled_timestamp = time.monotonic()
 
+        # Previous step's worker has run by now; promote its eager cache
+        # registrations from uncommitted to committed.
+        self.kv_cache_manager.commit_step()
         self.kv_cache_manager.new_step_starts()
 
         # First, schedule the RUNNING requests.

From 904629a6340b037bc132a5d4558e9922b59186fc Mon Sep 17 00:00:00 2001
From: zjy0516 <riverclouds.zhu@qq.com>
Date: Mon, 25 May 2026 08:07:47 +0000
Subject: [PATCH 2/4] update

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
---
 tests/v1/core/test_eager_cache_zombie.py | 187 +++++++++++------------
 tests/v1/core/test_prefix_caching.py     |  11 --
 vllm/v1/core/block_pool.py               |  36 +++--
 vllm/v1/core/kv_cache_manager.py         |  17 ++-
 vllm/v1/core/sched/scheduler.py          |  16 +-
 5 files changed, 135 insertions(+), 132 deletions(-)

diff --git a/tests/v1/core/test_eager_cache_zombie.py b/tests/v1/core/test_eager_cache_zombie.py
index 3674b4cd1909..b80d3f6ee159 100644
--- a/tests/v1/core/test_eager_cache_zombie.py
+++ b/tests/v1/core/test_eager_cache_zombie.py
@@ -1,7 +1,24 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """
-Tests for the eager-cache-registration zombie problem.
+Tests for the eager-cache-registration zombie protection.
+
+``KVCacheManager.allocate_slots`` registers block hashes into the cache map
+before the worker writes the K/V bytes. If a request is preempted/aborted
+before that write happens, those hash entries would otherwise become zombies
+that future requests could cache-hit, reading uninitialized memory.
+
+BlockPool tracks each in-flight scheduler step's eager registrations in a
+FIFO bucket queue (``_uncommitted``). The scheduler hooks them at two points:
+
+- ``schedule()`` top calls ``begin_step()`` to push a new bucket.
+- ``update_from_output()`` top calls ``commit_step()`` to pop the oldest
+  bucket once the matching worker batch has been confirmed.
+
+The preempt and abort paths call ``rollback_uncommitted(req_id)`` *before*
+``free(req)`` to evict that request's pending entries from every bucket.
+The normal-finish path calls only ``free(req)`` (no rollback), so worker-
+written entries that survived ``commit_step`` stay in the cache map.
 """
 
 from __future__ import annotations
@@ -31,18 +48,26 @@ def _init_hash():
     init_none_hash(sha256)
 
 
+def _uncommitted_blocks_for(manager: KVCacheManager, req_id: str) -> list:
+    """Flatten all uncommitted buckets and return req_id's tracked blocks."""
+    out: list = []
+    for bucket in manager.block_pool._uncommitted:
+        out.extend(bucket.get(req_id, []))
+    return out
+
+
 # ---------------------------------------------------------------------------
 # Test 1: full-attention single-group case
 # ---------------------------------------------------------------------------
 
 
-def test_rollback_on_preempt_before_worker_write():
+def test_rollback_uncommitted_evicts_preempt_zombies():
     """
-    Schedule req_A then ``free`` it immediately (simulating preempt/abort
-    before the worker has executed). The eager-registered hash entries
-    must be evicted from the cache map and the blocks' ``block_hash``
-    must be reset, so that a later req_B with the same prefix does NOT
-    cache-hit on never-written blocks.
+    Schedule req_A, then explicitly invoke the preempt cleanup
+    (``rollback_uncommitted`` + ``free``). The eager-registered hash entries
+    must be evicted from the cache map and the blocks' ``block_hash`` must be
+    reset, so that a later req_B with the same prefix does NOT cache-hit on
+    never-written blocks.
     """
     block_size = 16
     manager = KVCacheManager(
@@ -60,52 +85,39 @@ def test_rollback_on_preempt_before_worker_write():
     blocks_a = manager.allocate_slots(req_a, 32, 0, computed_a)
     assert blocks_a is not None
 
-    # After allocate_slots: 2 hashes registered eagerly and both tracked
-    # in ``_uncommitted[req_a]``.
+    # After allocate_slots: 2 hashes registered eagerly, both tracked in the
+    # current step's bucket.
     cache_map = manager.block_pool.cached_block_hash_to_block._cache
-    assert len(cache_map) == 2, (
-        "Sanity: 2 full blocks should each register a hash on allocate_slots"
-    )
-    assert len(manager.block_pool._uncommitted["a"]) == 2, (
-        "Sanity: both eager registrations should be tracked as uncommitted "
-        "until commit_step or rollback runs."
-    )
+    assert len(cache_map) == 2
+    assert len(_uncommitted_blocks_for(manager, "a")) == 2
 
     block_ids_a = blocks_a.get_block_ids()[0]
     eager_blocks = [manager.block_pool.blocks[bid] for bid in block_ids_a]
 
-    # --- Simulate preempt/abort BEFORE worker writes K/V bytes -----------
-    # In real life the worker would normally run between allocate_slots and
-    # the next free(). Here we go directly from allocate_slots → free
-    # *without* an intervening ``commit_step``, mirroring the race window
-    # where the request is preempted while the worker has not yet executed
-    # for this step.
+    # --- Simulate preempt before worker write ----------------------------
+    # The preempt/abort code paths invoke ``rollback_uncommitted`` *before*
+    # ``free`` so that any hash entries the worker hasn't yet backed with
+    # K/V bytes are evicted from the cache map.
+    manager.rollback_uncommitted(req_a.request_id)
     manager.free(req_a)
 
-    # --- Rollback evicts the uncommitted entries -------------------------
     assert len(cache_map) == 0, (
-        "free(uncommitted req) must evict both eager hash entries; got "
-        f"{len(cache_map)} left in cache map."
-    )
-    assert "a" not in manager.block_pool._uncommitted, (
-        "_uncommitted[req_a] must be cleared by rollback."
+        f"rollback must evict both eager hash entries, got {len(cache_map)}."
     )
+    assert _uncommitted_blocks_for(manager, "a") == []
     for blk in eager_blocks:
         assert blk.ref_cnt == 0, "free() did decrement ref_cnt"
         assert blk.block_hash is None, (
             "block.block_hash must be reset by rollback path."
         )
 
-    # --- req_B: same prefix → cache MISS ---------------------------------
-    # 48 tokens: first 32 would match req_a's now-rolled-back hashes, last
-    # 16 are new. Without zombies, this is a full cold miss.
+    # --- req_B: superset prefix → cache MISS -----------------------------
     req_b_tokens = common_token_ids + [99] * block_size
     req_b = make_request("b", req_b_tokens, block_size, sha256)
     _, num_computed_b = manager.get_computed_blocks(req_b)
-
     assert num_computed_b == 0, (
         f"req_b must not cache-hit on rolled-back entries; got "
-        f"{num_computed_b} cached tokens, expected 0."
+        f"{num_computed_b} cached tokens."
     )
 
 
@@ -114,14 +126,13 @@ def test_rollback_on_preempt_before_worker_write():
 # ---------------------------------------------------------------------------
 
 
-def test_rollback_on_preempt_for_mamba_hybrid():
+def test_rollback_uncommitted_covers_mamba_hybrid_groups():
     """
     Same race as test 1, but for full-attention + 2 Mamba groups. Confirms
     the rollback path covers every manager that calls ``cache_blocks``
     during ``allocate_slots``, not just full-attention.
     """
     block_size = 16
-    # 1 full-attention group + 2 Mamba groups (slice 0/1).
     manager = KVCacheManager(
         make_kv_cache_config_hybrid_model(
             block_size,
@@ -142,49 +153,33 @@ def test_rollback_on_preempt_for_mamba_hybrid():
 
     cache_map = manager.block_pool.cached_block_hash_to_block._cache
     eager_entry_count = len(cache_map)
-    assert eager_entry_count > 0, (
-        "Sanity: hybrid manager registers at least some eager entries on alloc"
-    )
-    assert len(manager.block_pool._uncommitted["a"]) == eager_entry_count, (
-        "Sanity: every eager registration across all groups should be tracked "
-        "as uncommitted."
-    )
+    assert eager_entry_count > 0
+    assert len(_uncommitted_blocks_for(manager, "a")) == eager_entry_count
 
     # Preempt before worker write.
+    manager.rollback_uncommitted(req_a.request_id)
     manager.free(req_a)
 
     assert len(cache_map) == 0, (
-        f"Rollback must clear all groups' eager entries; got "
-        f"{len(cache_map)} left in the cache map."
+        f"rollback must clear all groups' eager entries; got {len(cache_map)}."
     )
-    assert "a" not in manager.block_pool._uncommitted
+    assert _uncommitted_blocks_for(manager, "a") == []
 
-    # Future request cache MISS — no entries left to match.
     req_b = make_request("b", common_token_ids, block_size, sha256)
     _, num_computed_b = manager.get_computed_blocks(req_b)
-    assert num_computed_b == 0, (
-        f"req_b must not cache-hit on rolled-back entries; got "
-        f"{num_computed_b} cached tokens, expected 0."
-    )
+    assert num_computed_b == 0
 
 
 # ---------------------------------------------------------------------------
-# Test 3: control case — successful step + free should keep cache (no bug)
+# Test 3: normal-finish path keeps cache entries (no rollback called)
 # ---------------------------------------------------------------------------
 
 
-def test_commit_step_keeps_cache_across_normal_free():
+def test_finish_path_free_alone_preserves_cache_entries():
     """
-    Control: when a request runs the worker to completion (modelled here
-    by an explicit ``commit_step()`` call) and is then freed normally,
-    the cache entries SHOULD remain. ``commit_step`` clears
-    ``_uncommitted`` so the subsequent ``free`` does not roll anything
-    back.
-
-    Without the ``commit_step`` call, ``free`` would correctly roll the
-    entries back — that's exactly the eager-rollback behavior exercised
-    by tests 1 and 2 above. This test pins down the other side: the
-    commit hook is what preserves cache hits across normal completion.
+    Control: the normal-finish path calls only ``free()`` (no rollback).
+    Cache entries the worker has confirmed must remain hittable for future
+    requests, with or without an intervening ``commit_step``.
     """
     block_size = 16
     manager = KVCacheManager(
@@ -198,27 +193,27 @@ def test_commit_step_keeps_cache_across_normal_free():
     req_a = make_request("a", common_token_ids, block_size, sha256)
     computed_a, _ = manager.get_computed_blocks(req_a)
     manager.allocate_slots(req_a, 32, 0, computed_a)
-    assert len(manager.block_pool._uncommitted["a"]) == 2
-
-    # Simulate successful worker execution by committing the step.
-    # Scheduler.schedule() calls this at the start of every step; here we
-    # call it manually to model "previous step's worker has confirmed".
-    manager.commit_step()
-    assert manager.block_pool._uncommitted == {}, (
-        "commit_step should clear all pending uncommitted registrations."
-    )
+    assert len(_uncommitted_blocks_for(manager, "a")) == 2
 
+    # ``free`` on its own must not touch the cache map. The pending bucket
+    # entries get cleaned up later when the matching commit_step runs (or
+    # stay tracked but harmless until then).
     manager.free(req_a)
 
     cache_map = manager.block_pool.cached_block_hash_to_block._cache
     assert len(cache_map) == 2, (
-        "Normal completion (commit_step + free) should leave the cache "
-        f"entries intact, got {len(cache_map)} entries."
+        "Normal free should leave the cache entries intact for future hits."
     )
 
+    # commit_step pops the oldest bucket. After that, no rollback can affect
+    # these entries even if rollback_uncommitted were called.
+    manager.commit_step()
+    assert _uncommitted_blocks_for(manager, "a") == []
+    assert len(cache_map) == 2
+
 
 # ---------------------------------------------------------------------------
-# Test 4: scheduler-level preemption test
+# Test 4: scheduler-level preemption test (real Scheduler.schedule() path)
 # ---------------------------------------------------------------------------
 
 
@@ -226,27 +221,22 @@ def test_scheduler_preempt_rolls_back_target_step_eager_cache():
     """
     Exercise the rollback through a real ``Scheduler.schedule()`` call.
 
-    Setup seeds two RUNNING chunked-prefill requests with one already-committed
-    block each. In the target schedule call, req_A is scheduled first and
-    eagerly caches two new full blocks. req_B then cannot allocate, so priority
-    preemption removes req_A from the same scheduler output before any worker
-    can see or execute that work.
-
-    The two target-step entries must be evicted on preempt; the prior-step
-    seeded block remains in the cache map (it was promoted to committed when
-    ``schedule`` called ``commit_step`` at the start of this step). A future
-    req_C with req_A's first two blocks as prefix therefore hits exactly the
-    one seeded block.
+    Setup seeds two RUNNING chunked-prefill requests with one already-
+    committed block each. In the target schedule call, req_A is scheduled
+    first and eagerly caches two new full blocks. req_B then cannot
+    allocate, so priority preemption removes req_A from the same scheduler
+    output before any worker can see or execute that work.
+
+    The preempt path (`_preempt_request`) calls ``rollback_uncommitted``
+    before ``free``, evicting the two target-step entries. The prior-step
+    seeded block remains in the cache map. A future req_C with req_A's
+    first two blocks as prefix therefore hits exactly the one seeded block.
     """
     block_size = 16
     # Resource budget is intentionally tight to force the exact race:
     #   max_num_batched_tokens=48 = 32 (req_a's 2 remaining blocks)
     #                             + 16 (req_b's 1 remaining block).
-    #     Both requests *want* to advance fully in the target step.
-    #   num_blocks=5 = 1 null
-    #                + 1 seeded committed block for req_a
-    #                + 1 seeded committed block for req_b
-    #                + 2 free blocks.
+    #   num_blocks=5 = 1 null + 2 seeded committed + 2 free.
     #     The 2 free blocks are exactly enough for req_a's target-step
     #     allocation (2 new full blocks); req_b then has 0 free → triggers
     #     preempt of the lowest-priority running request, which is req_a.
@@ -283,9 +273,9 @@ def test_scheduler_preempt_rolls_back_target_step_eager_cache():
     req_b.arrival_time = 2.0
 
     # Plant both requests directly in RUNNING with their first block already
-    # computed. This bypasses the normal admission path so we start the test
-    # mid-chunked-prefill (the only state where this race opens) without
-    # having to drive multiple scheduler steps to set it up.
+    # computed. This bypasses the normal admission path so we start the
+    # test mid-chunked-prefill (the only state where this race opens)
+    # without having to drive multiple scheduler steps to set it up.
     manager = scheduler.kv_cache_manager
     for req in (req_a, req_b):
         seeded_blocks = manager.allocate_slots(req, block_size)
@@ -295,10 +285,15 @@ def test_scheduler_preempt_rolls_back_target_step_eager_cache():
         scheduler.requests[req.request_id] = req
     scheduler.running = [req_a, req_b]
 
+    # The seed allocations registered hashes in the BlockPool's first lazy
+    # bucket. Production puts seed entries in a *prior* committed step, not
+    # the current step. Commit here to model that.
+    manager.commit_step()
+
     # Sanity: post-seeding state matches the budget plan above.
     cache_map = manager.block_pool.cached_block_hash_to_block._cache
     assert len(cache_map) == 2  # one committed hash per seeded request.
-    assert manager.block_pool.get_num_free_blocks() == 2  # exactly the race window.
+    assert manager.block_pool.get_num_free_blocks() == 2  # the race window.
 
     output = scheduler.schedule()
 
@@ -312,8 +307,8 @@ def test_scheduler_preempt_rolls_back_target_step_eager_cache():
 
     # A future request with A's first two blocks as a prefix should hit
     # exactly one block: the prior-step seeded block (committed via
-    # commit_step at the start of this scheduler step). The target-step
-    # eager entry for the would-be second block was rolled back on preempt.
+    # commit_step before this scheduler step). The target-step eager entry
+    # for the would-be second block was rolled back on preempt.
     req_c = make_request(
         "c",
         tokens_a[: 2 * block_size] + [99] * block_size,
diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py
index 32d4ab08bbf9..546412b1d2f8 100644
--- a/tests/v1/core/test_prefix_caching.py
+++ b/tests/v1/core/test_prefix_caching.py
@@ -269,9 +269,6 @@ def test_prefill(hash_fn):
     free_block_queue = manager.block_pool.free_block_queue
     assert free_block_queue.num_free_blocks == 5
 
-    # Model successful worker execution before freeing so the eager-rollback
-    # path does not evict the committed cache entries.
-    manager.commit_step()
     manager.free(req0)
     manager.free(req1)
 
@@ -403,7 +400,6 @@ def test_prefill_hybrid_model():
                 assert block.ref_cnt == 2
 
     block_hashes = req1.block_hashes
-    manager.commit_step()
     manager.free(req0)
     manager.free(req1)
 
@@ -583,7 +579,6 @@ def test_prefill_hybrid_model_eagle():
                 assert block.ref_cnt == 2
 
     block_hashes = req1.block_hashes
-    manager.commit_step()
     manager.free(req0)
     manager.free(req1)
 
@@ -1227,7 +1222,6 @@ def test_evict():
     # 10 - (6 + 3) == 1
     assert manager.block_pool.free_block_queue.num_free_blocks == 1
 
-    manager.commit_step()
     manager.free(req0)
     manager.free(req1)
     assert manager.block_pool.free_block_queue.num_free_blocks == 10
@@ -1327,7 +1321,6 @@ def test_computed_blocks_not_evicted():
     assert blocks.blocks[0][0].block_id == 2
 
     # Free the blocks.
-    manager.commit_step()
     manager.free(req0)
     manager.free(req1)
 
@@ -1761,7 +1754,6 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():
     ]
     # | Common-0 | Common-1 | Common-2 | Req1-3 (F) | Req1-4 (F) |
     # | Req1-5(F)| ... |
-    manager.commit_step()
     manager.free(req1)
     assert {block.ref_cnt for block in block_part1[:3]} == {1}
     assert {block.ref_cnt for block in block_part1[3:]} == {0}
@@ -2353,7 +2345,6 @@ def test_eagle_enabled_removes_last_block():
     manager.allocate_slots(
         req, len(token_ids), len(computed_blocks.blocks[0]) * 16, computed_blocks
     )
-    manager.commit_step()
     manager.free(req)
 
     # New request with same tokens + Eagle enabled
@@ -2386,7 +2377,6 @@ def test_eagle_with_partial_blocks():
     manager.allocate_slots(
         req, len(token_ids), len(computed_blocks.blocks[0]) * 16, computed_blocks
     )
-    manager.commit_step()
     manager.free(req)
 
     # New request with Eagle enabled
@@ -2431,7 +2421,6 @@ def test_eagle_with_sliding_window():
     # record the block hash of the first block in the request for later use
     block_hash_first_block = req.block_hashes[0]
     assert block_hash_first_block is not None
-    manager.commit_step()
     manager.free(req)
 
     # New request with Eagle enabled
diff --git a/vllm/v1/core/block_pool.py b/vllm/v1/core/block_pool.py
index 7c56bcbea49c..0c000c00f5d3 100644
--- a/vllm/v1/core/block_pool.py
+++ b/vllm/v1/core/block_pool.py
@@ -181,10 +181,8 @@ def __init__(
 
         self.metrics_collector = metrics_collector
 
-        # Blocks whose hashes were eagerly registered this scheduler step but
-        # whose K/V bytes are not yet worker-confirmed. ``commit_step`` clears
-        # at step boundaries; ``rollback_uncommitted`` evicts on early free.
-        self._uncommitted: dict[str, list[KVCacheBlock]] = {}
+        # FIFO queue of eager-registration buckets, one per in-flight step.
+        self._uncommitted: list[dict[str, list[KVCacheBlock]]] = []
 
     def get_cached_block(
         self, block_hash: BlockHash, kv_cache_group_ids: list[int]
@@ -288,9 +286,11 @@ def cache_full_blocks(
             if new_hashes is not None:
                 new_hashes.append(maybe_convert_block_hash(block_hash))
 
-            # Track for commit_step / rollback_uncommitted. Lazy list alloc.
+            # Track in the current bucket; lazy-create if no step open.
             if uncommitted_for_req is None:
-                uncommitted_for_req = self._uncommitted.setdefault(
+                if not self._uncommitted:
+                    self._uncommitted.append({})
+                uncommitted_for_req = self._uncommitted[-1].setdefault(
                     request.request_id, []
                 )
             uncommitted_for_req.append(blk)
@@ -446,20 +446,24 @@ def free_blocks(self, ordered_blocks: Iterable[KVCacheBlock]) -> None:
         )
 
     def rollback_uncommitted(self, request_id: str) -> int:
-        """Evict ``request_id``'s eager-registered-but-not-yet-worker-confirmed
-        cache entries. Returns the number evicted. Idempotent.
+        """Evict ``request_id``'s uncommitted cache entries across all buckets.
+        Called by preempt/abort paths. Returns the number evicted.
         """
-        blocks = self._uncommitted.pop(request_id, None)
-        if not blocks:
-            return 0
-        return sum(self._maybe_evict_cached_block(b) for b in blocks)
+        evicted = 0
+        for bucket in self._uncommitted:
+            blocks = bucket.pop(request_id, None)
+            if blocks:
+                evicted += sum(self._maybe_evict_cached_block(b) for b in blocks)
+        return evicted
+
+    def begin_step(self) -> None:
+        """Open a new eager-registration bucket for the current step."""
+        self._uncommitted.append({})
 
     def commit_step(self) -> None:
-        """Promote all pending eager registrations to committed. Called at
-        scheduler step boundaries. Idempotent.
-        """
+        """Promote the oldest pending bucket to committed. Idempotent."""
         if self._uncommitted:
-            self._uncommitted.clear()
+            self._uncommitted.pop(0)
 
     def evict_blocks(self, block_ids: set[int]) -> None:
         """evict blocks from the prefix cache by their block IDs.
diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py
index 50e14fd551dd..375bb7cf8ad2 100644
--- a/vllm/v1/core/kv_cache_manager.py
+++ b/vllm/v1/core/kv_cache_manager.py
@@ -434,11 +434,14 @@ def free(self, request: Request) -> None:
         Args:
             request: The request to free the blocks.
         """
-        # Roll back uncommitted eager cache entries (no-op for cross-step
-        # free; cleans up zombies for preempt/abort mid-step).
-        self.block_pool.rollback_uncommitted(request.request_id)
         self.coordinator.free(request.request_id)
 
+    def rollback_uncommitted(self, request_id: str) -> int:
+        """Evict pending eager cache entries for ``request_id``. Call before
+        ``free`` on preempt/abort paths to avoid zombie hash entries.
+        """
+        return self.block_pool.rollback_uncommitted(request_id)
+
     def remove_skipped_blocks(
         self, request_id: str, total_computed_tokens: int
     ) -> None:
@@ -572,8 +575,12 @@ def new_step_starts(self) -> None:
         """Called when a new step is started."""
         self.coordinator.new_step_starts()
 
+    def begin_step(self) -> None:
+        """Open a new eager-registration bucket at the top of ``schedule()``."""
+        self.block_pool.begin_step()
+
     def commit_step(self) -> None:
-        """Promote eager cache registrations from the previous scheduler step
-        to committed. Called at step boundaries. Idempotent.
+        """Pop the oldest pending bucket at the top of ``update_from_output``.
+        Idempotent.
         """
         self.block_pool.commit_step()
diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
index b11f16512803..4ef2ab30f966 100644
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -359,9 +359,9 @@ def schedule(self) -> SchedulerOutput:
         # For logging.
         scheduled_timestamp = time.monotonic()
 
-        # Previous step's worker has run by now; promote its eager cache
-        # registrations from uncommitted to committed.
-        self.kv_cache_manager.commit_step()
+        # Open this step's eager-registration bucket; the matching
+        # commit_step runs at the top of update_from_output.
+        self.kv_cache_manager.begin_step()
         self.kv_cache_manager.new_step_starts()
 
         # First, schedule the RUNNING requests.
@@ -938,6 +938,8 @@ def _preempt_request(self, request: Request, timestamp: float) -> None:
         assert request.status == RequestStatus.RUNNING, (
             "Only running requests can be preempted"
         )
+        # Worker hasn't run for this step yet; evict zombies before free.
+        self.kv_cache_manager.rollback_uncommitted(request.request_id)
         self.kv_cache_manager.free(request)
         self.encoder_cache_manager.free(request)
         request.status = RequestStatus.PREEMPTED
@@ -1288,6 +1290,10 @@ def update_from_output(
         scheduler_output: SchedulerOutput,
         model_runner_output: ModelRunnerOutput,
     ) -> dict[int, EngineCoreOutputs]:
+        # Worker has confirmed writes for this step; commit its eager
+        # registrations before any finish-triggered free() runs below.
+        self.kv_cache_manager.commit_step()
+
         sampled_token_ids = model_runner_output.sampled_token_ids
         logprobs = model_runner_output.logprobs
         prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict
@@ -2087,7 +2093,9 @@ def _update_waiting_for_remote_kv(self, request: Request) -> None:
                 self.kv_cache_manager.cache_blocks(request, request.num_computed_tokens)
             else:
                 # No valid computed tokens, release allocated blocks.
-                # There may be a local cache hit on retry.
+                # There may be a local cache hit on retry. KV load failed,
+                # so evict any zombies before releasing.
+                self.kv_cache_manager.rollback_uncommitted(request.request_id)
                 self.kv_cache_manager.free(request)
 
             self.failed_recving_kv_req_ids.remove(request.request_id)

From 859678780c02b71a7dfe15653ca20292feeb47af Mon Sep 17 00:00:00 2001
From: zjy0516 <riverclouds.zhu@qq.com>
Date: Tue, 26 May 2026 08:24:48 +0000
Subject: [PATCH 3/4] update

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
---
 tests/v1/core/test_eager_cache_zombie.py | 238 +++++++++++++++++++++--
 vllm/v1/core/block_pool.py               |  48 +++--
 vllm/v1/core/kv_cache_manager.py         |  20 +-
 vllm/v1/core/sched/async_scheduler.py    |   7 +-
 vllm/v1/core/sched/scheduler.py          |  16 +-
 5 files changed, 291 insertions(+), 38 deletions(-)

diff --git a/tests/v1/core/test_eager_cache_zombie.py b/tests/v1/core/test_eager_cache_zombie.py
index b80d3f6ee159..1f665e597f23 100644
--- a/tests/v1/core/test_eager_cache_zombie.py
+++ b/tests/v1/core/test_eager_cache_zombie.py
@@ -2,23 +2,6 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """
 Tests for the eager-cache-registration zombie protection.
-
-``KVCacheManager.allocate_slots`` registers block hashes into the cache map
-before the worker writes the K/V bytes. If a request is preempted/aborted
-before that write happens, those hash entries would otherwise become zombies
-that future requests could cache-hit, reading uninitialized memory.
-
-BlockPool tracks each in-flight scheduler step's eager registrations in a
-FIFO bucket queue (``_uncommitted``). The scheduler hooks them at two points:
-
-- ``schedule()`` top calls ``begin_step()`` to push a new bucket.
-- ``update_from_output()`` top calls ``commit_step()`` to pop the oldest
-  bucket once the matching worker batch has been confirmed.
-
-The preempt and abort paths call ``rollback_uncommitted(req_id)`` *before*
-``free(req)`` to evict that request's pending entries from every bucket.
-The normal-finish path calls only ``free(req)`` (no rollback), so worker-
-written entries that survived ``commit_step`` stay in the cache map.
 """
 
 from __future__ import annotations
@@ -30,7 +13,7 @@
     make_kv_cache_config_hybrid_model,
     make_request,
 )
-from tests.v1.core.utils import create_scheduler
+from tests.v1.core.utils import create_requests, create_scheduler
 from vllm.utils.hashing import sha256
 from vllm.v1.core.kv_cache_manager import KVCacheManager
 from vllm.v1.core.kv_cache_utils import init_none_hash
@@ -38,6 +21,7 @@
     SchedulingPolicy,
     create_request_queue,
 )
+from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import RequestStatus
 
 pytestmark = pytest.mark.cpu_test
@@ -213,7 +197,223 @@ def test_finish_path_free_alone_preserves_cache_entries():
 
 
 # ---------------------------------------------------------------------------
-# Test 4: scheduler-level preemption test (real Scheduler.schedule() path)
+# Test 4: committed=True kwarg skips _uncommitted tracking
+# ---------------------------------------------------------------------------
+
+
+def test_cache_blocks_committed_kwarg_skips_uncommitted_tracking():
+    """
+    When ``cache_blocks(committed=True)`` is called from a path that knows the
+    worker has already confirmed the K/V writes (AsyncScheduler post-output,
+    KV connector load completion), the new cache entries must NOT be tracked
+    in ``_uncommitted``. A subsequent ``rollback_uncommitted`` for the same
+    request must leave those entries intact in ``cached_block_hash_to_block``.
+    """
+    block_size = 16
+    manager = KVCacheManager(
+        make_kv_cache_config(block_size, 11),
+        max_model_len=8192,
+        enable_caching=True,
+        hash_block_size=block_size,
+    )
+
+    tokens = [i for i in range(2) for _ in range(block_size)]
+    req = make_request("a", tokens, block_size, sha256)
+    computed, _ = manager.get_computed_blocks(req)
+    blocks = manager.allocate_slots(req, 32, 0, computed)
+    assert blocks is not None
+
+    cache_map = manager.block_pool.cached_block_hash_to_block._cache
+    assert len(cache_map) == 2
+    # The eager (default) path tracks the 2 new registrations.
+    assert len(_uncommitted_blocks_for(manager, "a")) == 2
+
+    # Now drop the eager tracking by committing the step, then invoke a
+    # committed-mode cache_blocks call (mimicking the AsyncScheduler/connector
+    # post-confirm paths). It must register into cache_map but NOT add to
+    # _uncommitted.
+    manager.commit_step()
+    assert _uncommitted_blocks_for(manager, "a") == []
+
+    # Construct a longer prompt request with the same first 32 tokens; calling
+    # cache_blocks(committed=True) for those tokens should be a no-op (already
+    # cached) but importantly must not poison _uncommitted.
+    tokens_longer = tokens + [99] * block_size
+    req_longer = make_request("b", tokens_longer, block_size, sha256)
+    computed_longer, num_computed_longer = manager.get_computed_blocks(req_longer)
+    assert num_computed_longer == 32  # hits the two committed blocks
+    manager.allocate_slots(req_longer, 16, 32, computed_longer)
+
+    # Now simulate a committed-path registration for req_longer (the bytes are
+    # already worker-confirmed for the suffix). This should NOT add to
+    # _uncommitted["b"].
+    before = len(_uncommitted_blocks_for(manager, "b"))
+    manager.cache_blocks(req_longer, 48, committed=True)
+    after = len(_uncommitted_blocks_for(manager, "b"))
+    assert after == before, (
+        f"committed=True must not grow _uncommitted; before={before}, after={after}"
+    )
+
+    # And subsequent rollback for req_longer must not evict any of the
+    # committed entries from cache_map.
+    cache_map_before_rollback = dict(cache_map)
+    manager.rollback_uncommitted(req_longer.request_id)
+    # The cache_map should still contain at least the two original committed
+    # entries (req_a's prefix). The eager block(s) for req_longer's suffix
+    # would be evicted by rollback, which is expected.
+    for key, blk in cache_map_before_rollback.items():
+        if blk in [manager.block_pool.blocks[bid] for bid in blocks.get_block_ids()[0]]:
+            assert key in cache_map, "rollback evicted a committed entry from cache_map"
+
+
+# ---------------------------------------------------------------------------
+# Test 5: reset_prefix_cache also clears _uncommitted
+# ---------------------------------------------------------------------------
+
+
+def test_reset_prefix_cache_clears_uncommitted():
+    """
+    ``reset_prefix_cache`` wipes ``cached_block_hash_to_block`` and resets every
+    block's ``block_hash``. ``_uncommitted`` would dangle if not cleared
+    alongside — pointing at blocks whose hashes have been reset to None — so
+    a subsequent ``rollback_uncommitted`` call would walk stale entries and
+    potentially evict blocks that have been re-allocated to other requests.
+
+    Confirm reset zeroes ``_uncommitted`` too.
+    """
+    block_size = 16
+    manager = KVCacheManager(
+        make_kv_cache_config(block_size, 11),
+        max_model_len=8192,
+        enable_caching=True,
+        hash_block_size=block_size,
+    )
+
+    tokens = [i for i in range(2) for _ in range(block_size)]
+    req = make_request("a", tokens, block_size, sha256)
+    computed, _ = manager.get_computed_blocks(req)
+    manager.allocate_slots(req, 32, 0, computed)
+
+    # Sanity: eager registration populated _uncommitted.
+    assert len(_uncommitted_blocks_for(manager, "a")) == 2
+    assert len(manager.block_pool._uncommitted) == 1
+
+    # Free the request so reset_prefix_cache succeeds (it requires zero
+    # outstanding blocks except the null block).
+    manager.free(req)
+    assert manager.reset_prefix_cache()
+
+    # Both the cache map and the uncommitted queue must be empty.
+    assert len(manager.block_pool.cached_block_hash_to_block) == 0
+    assert len(manager.block_pool._uncommitted) == 0, (
+        "reset_prefix_cache must clear _uncommitted; otherwise stale entries "
+        "point at blocks whose hashes have been reset and a subsequent "
+        "rollback_uncommitted could touch the wrong physical blocks."
+    )
+
+    # Sanity follow-through: a fresh request after reset goes through the
+    # normal eager-registration path without surprises.
+    req2 = make_request("b", tokens, block_size, sha256)
+    computed2, _ = manager.get_computed_blocks(req2)
+    assert _ == 0  # cache was cleared by reset
+    manager.allocate_slots(req2, 32, 0, computed2)
+    assert len(_uncommitted_blocks_for(manager, "b")) == 2
+
+
+# ---------------------------------------------------------------------------
+# Test 6: FIFO bucket discipline across multiple in-flight scheduler steps
+# ---------------------------------------------------------------------------
+
+
+def test_fifo_bucket_discipline_under_multi_in_flight_steps():
+    """
+    Drive ``Scheduler.schedule()`` twice without an intervening
+    ``update_from_output``, simulating the async batch-queue path. Each
+    schedule must push a bucket via ``begin_step``; ``update_from_output``
+    for the OLDEST step must pop only that bucket via ``commit_step``,
+    leaving the newer in-flight bucket intact and its entries available for
+    ``rollback_uncommitted`` if that newer step's request is later aborted.
+    """
+    # async_scheduling + PP>=2 lets the scheduler keep multiple in-flight
+    # schedule() calls outstanding before any update_from_output fires —
+    # the precise condition that exposes the FIFO discipline.
+    # enable_prefix_caching is needed for cache_blocks (and therefore the
+    # eager registrations we want to track) to fire at all.
+    scheduler = create_scheduler(
+        async_scheduling=True,
+        pipeline_parallel_size=2,
+        enable_prefix_caching=True,
+    )
+    manager = scheduler.kv_cache_manager
+    cache_map = manager.block_pool.cached_block_hash_to_block._cache
+
+    # Two requests with full-block prompts so each schedule() actually
+    # registers eager cache entries (block_size=16 default → 32 tokens =
+    # 2 full blocks per request).
+    reqs = create_requests(num_requests=2, num_tokens=32)
+
+    assert len(manager.block_pool._uncommitted) == 0
+
+    # --- Step N: admit and schedule req_a --------------------------------
+    scheduler.add_request(reqs[0])
+    output_n = scheduler.schedule()
+    assert len(manager.block_pool._uncommitted) == 1, (
+        "schedule() must push a bucket via begin_step"
+    )
+    assert reqs[0].request_id in manager.block_pool._uncommitted[0]
+    entries_after_n = len(cache_map)
+    assert entries_after_n > 0
+
+    # --- Step N+1: admit and schedule req_b WITHOUT update_from_output ---
+    scheduler.add_request(reqs[1])
+    scheduler.schedule()
+    assert len(manager.block_pool._uncommitted) == 2, (
+        "second schedule() must push another bucket while the first remains"
+    )
+    # Newer bucket carries req_b's eager registrations.
+    assert reqs[1].request_id in manager.block_pool._uncommitted[-1]
+    # Older bucket still carries req_a's, unchanged.
+    assert reqs[0].request_id in manager.block_pool._uncommitted[0]
+
+    # --- Worker N completes; update_from_output(N) commits oldest bucket -
+    req_ids_n = list(output_n.num_scheduled_tokens.keys())
+    scheduler.update_from_output(
+        output_n,
+        ModelRunnerOutput(
+            req_ids=req_ids_n,
+            req_id_to_index={r: i for i, r in enumerate(req_ids_n)},
+            sampled_token_ids=[[] for _ in req_ids_n],  # still prefilling
+            logprobs=None,
+            prompt_logprobs_dict={},
+            pooler_output=[],
+        ),
+    )
+    assert len(manager.block_pool._uncommitted) == 1, (
+        "update_from_output must commit only the oldest bucket via commit_step"
+    )
+    assert reqs[0].request_id not in manager.block_pool._uncommitted[0], (
+        "commit_step must remove req_a's tracking; req_a's entries stay in "
+        "cache_map as committed."
+    )
+    assert reqs[1].request_id in manager.block_pool._uncommitted[0]
+
+    # --- Abort req_b mid-flight (its worker batch has NOT confirmed) -----
+    # rollback must reach into the still-pending bucket and evict req_b's
+    # entries; req_a's committed entries are untouched.
+    manager.rollback_uncommitted(reqs[1].request_id)
+    manager.free(reqs[1])
+    assert _uncommitted_blocks_for(manager, reqs[1].request_id) == []
+    assert len(cache_map) < entries_after_n + 1, (
+        "rollback must evict req_b's eager entries from cache_map"
+    )
+    # req_a's entries (committed) still hittable.
+    assert _uncommitted_blocks_for(manager, reqs[0].request_id) == []
+    rollback_a = manager.rollback_uncommitted(reqs[0].request_id)
+    assert rollback_a == 0, "rollback after commit must be a no-op for req_a"
+
+
+# ---------------------------------------------------------------------------
+# Test 5: scheduler-level preemption test (real Scheduler.schedule() path)
 # ---------------------------------------------------------------------------
 
 
diff --git a/vllm/v1/core/block_pool.py b/vllm/v1/core/block_pool.py
index 0c000c00f5d3..cb461a7af5fc 100644
--- a/vllm/v1/core/block_pool.py
+++ b/vllm/v1/core/block_pool.py
@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections import deque
 from collections.abc import Iterable, Sequence
+from contextlib import contextmanager
 from typing import Any
 
 from vllm.distributed.kv_events import (
@@ -182,7 +184,12 @@ def __init__(
         self.metrics_collector = metrics_collector
 
         # FIFO queue of eager-registration buckets, one per in-flight step.
-        self._uncommitted: list[dict[str, list[KVCacheBlock]]] = []
+        self._uncommitted: deque[dict[str, list[KVCacheBlock]]] = deque()
+        # When True, cache_full_blocks skips _uncommitted tracking. Used by
+        # post-worker-confirmed paths (AsyncScheduler post-output, KV connector
+        # load completion) where the bytes are already valid and the entry
+        # must not be evictable by a later rollback_uncommitted call.
+        self._suppress_uncommitted_tracking: bool = False
 
     def get_cached_block(
         self, block_hash: BlockHash, kv_cache_group_ids: list[int]
@@ -267,7 +274,19 @@ def cache_full_blocks(
         new_hashes: list[ExternalBlockHash] | None = (
             [] if self.enable_kv_cache_events else None
         )
-        uncommitted_for_req: list[KVCacheBlock] | None = None
+
+        # Where to record this request's new registrations for later
+        # rollback_uncommitted. ``None`` means skip tracking (the caller
+        # passed committed=True, signalling worker-confirmed bytes).
+        if self._suppress_uncommitted_tracking:
+            uncommitted_bucket: list[KVCacheBlock] | None = None
+        else:
+            if not self._uncommitted:
+                self._uncommitted.append({})
+            uncommitted_bucket = self._uncommitted[-1].setdefault(
+                request.request_id, []
+            )
+
         for i, blk in enumerate(new_full_blocks):
             # Some blocks may be null or masked out when enabling sparse attention
             # like sliding window attention, or Mamba models with prefix-caching
@@ -286,14 +305,8 @@ def cache_full_blocks(
             if new_hashes is not None:
                 new_hashes.append(maybe_convert_block_hash(block_hash))
 
-            # Track in the current bucket; lazy-create if no step open.
-            if uncommitted_for_req is None:
-                if not self._uncommitted:
-                    self._uncommitted.append({})
-                uncommitted_for_req = self._uncommitted[-1].setdefault(
-                    request.request_id, []
-                )
-            uncommitted_for_req.append(blk)
+            if uncommitted_bucket is not None:
+                uncommitted_bucket.append(blk)
 
         if self.enable_kv_cache_events:
             if num_cached_blocks == 0:
@@ -460,10 +473,23 @@ def begin_step(self) -> None:
         """Open a new eager-registration bucket for the current step."""
         self._uncommitted.append({})
 
+    @contextmanager
+    def suppress_uncommitted_tracking(self):
+        """Skip ``_uncommitted`` tracking within this context. Used by paths
+        that register cache entries whose K/V bytes are already worker-
+        confirmed (AsyncScheduler post-output, KV connector load completion).
+        """
+        prev = self._suppress_uncommitted_tracking
+        self._suppress_uncommitted_tracking = True
+        try:
+            yield
+        finally:
+            self._suppress_uncommitted_tracking = prev
+
     def commit_step(self) -> None:
         """Promote the oldest pending bucket to committed. Idempotent."""
         if self._uncommitted:
-            self._uncommitted.pop(0)
+            self._uncommitted.popleft()
 
     def evict_blocks(self, block_ids: set[int]) -> None:
         """evict blocks from the prefix cache by their block IDs.
diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py
index 375bb7cf8ad2..961574f6a40a 100644
--- a/vllm/v1/core/kv_cache_manager.py
+++ b/vllm/v1/core/kv_cache_manager.py
@@ -547,15 +547,31 @@ def get_block_ids(self, request_id: str) -> tuple[list[int], ...]:
         """Get the block ids of a request."""
         return self.get_blocks(request_id).get_block_ids()
 
-    def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:
+    def cache_blocks(
+        self,
+        request: Request,
+        num_computed_tokens: int,
+        *,
+        committed: bool = False,
+    ) -> None:
         """Cache the blocks for the request, if enabled.
 
         Args:
             request: The request to cache the blocks.
             num_computed_tokens: The number of computed tokens, including tokens
                 that are already cached and tokens to be cached.
+            committed: Pass True when the worker has already confirmed writing
+                the bytes for these blocks (e.g. AsyncScheduler post-output
+                cache update, KV connector load completion). Skips
+                ``_uncommitted`` tracking so the entries are not evictable by
+                a subsequent ``rollback_uncommitted`` call.
         """
-        if self.enable_caching:
+        if not self.enable_caching:
+            return
+        if committed:
+            with self.block_pool.suppress_uncommitted_tracking():
+                self.coordinator.cache_blocks(request, num_computed_tokens)
+        else:
             self.coordinator.cache_blocks(request, num_computed_tokens)
 
     def create_kv_cache_blocks(
diff --git a/vllm/v1/core/sched/async_scheduler.py b/vllm/v1/core/sched/async_scheduler.py
index cb61bcabd3ee..43aa90d05a74 100644
--- a/vllm/v1/core/sched/async_scheduler.py
+++ b/vllm/v1/core/sched/async_scheduler.py
@@ -54,8 +54,13 @@ def _update_request_with_output(
         assert request.num_output_placeholders >= 0
 
         # Cache the new tokens. Preempted requests should be skipped.
+        # Worker has already confirmed these bytes (we're past future.result()
+        # for this step), so register as committed — must not be evictable by
+        # a later rollback_uncommitted call.
         if status_before_update == RequestStatus.RUNNING:
             self.kv_cache_manager.cache_blocks(
-                request, request.num_computed_tokens - request.num_output_placeholders
+                request,
+                request.num_computed_tokens - request.num_output_placeholders,
+                committed=True,
             )
         return new_token_ids, stopped
diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
index 4ef2ab30f966..d8cbc9db8007 100644
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -2089,8 +2089,12 @@ def _update_waiting_for_remote_kv(self, request: Request) -> None:
             # Request had KV load failures; num_computed_tokens was already
             # updated in _update_requests_with_invalid_blocks
             if request.num_computed_tokens:
-                # Cache any valid computed tokens.
-                self.kv_cache_manager.cache_blocks(request, request.num_computed_tokens)
+                # Cache any valid computed tokens. These bytes are
+                # connector-loaded (valid by construction), so register as
+                # committed to keep them safe from later rollback_uncommitted.
+                self.kv_cache_manager.cache_blocks(
+                    request, request.num_computed_tokens, committed=True
+                )
             else:
                 # No valid computed tokens, release allocated blocks.
                 # There may be a local cache hit on retry. KV load failed,
@@ -2100,9 +2104,11 @@ def _update_waiting_for_remote_kv(self, request: Request) -> None:
 
             self.failed_recving_kv_req_ids.remove(request.request_id)
         else:
-            # Now that the blocks are ready, actually cache them.
-            # This will cache the blocks iff caching is enabled.
-            self.kv_cache_manager.cache_blocks(request, request.num_computed_tokens)
+            # Now that the blocks are ready, actually cache them. Bytes are
+            # connector-loaded and worker-confirmed, so register as committed.
+            self.kv_cache_manager.cache_blocks(
+                request, request.num_computed_tokens, committed=True
+            )
 
             # on a full prompt hit, we need to re-compute the last token
             # in order to be able to sample the next token

From 660c9326c06aa2afd7170a5fb423cb705ad39d4f Mon Sep 17 00:00:00 2001
From: zjy0516 <riverclouds.zhu@qq.com>
Date: Tue, 26 May 2026 12:41:04 +0000
Subject: [PATCH 4/4] update

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
---
 tests/v1/core/test_eager_cache_zombie.py               |  6 ++++++
 .../unit/test_invalid_blocks_correctness.py            |  4 ++--
 vllm/v1/core/block_pool.py                             | 10 ++++------
 3 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/tests/v1/core/test_eager_cache_zombie.py b/tests/v1/core/test_eager_cache_zombie.py
index 1f665e597f23..096f36127169 100644
--- a/tests/v1/core/test_eager_cache_zombie.py
+++ b/tests/v1/core/test_eager_cache_zombie.py
@@ -66,6 +66,7 @@ def test_rollback_uncommitted_evicts_preempt_zombies():
     req_a = make_request("a", common_token_ids, block_size, sha256)
     computed_a, num_computed_a = manager.get_computed_blocks(req_a)
     assert num_computed_a == 0  # cold cache
+    manager.begin_step()
     blocks_a = manager.allocate_slots(req_a, 32, 0, computed_a)
     assert blocks_a is not None
 
@@ -132,6 +133,7 @@ def test_rollback_uncommitted_covers_mamba_hybrid_groups():
     common_token_ids = [i for i in range(2) for _ in range(block_size)]
     req_a = make_request("a", common_token_ids, block_size, sha256)
     computed_a, _ = manager.get_computed_blocks(req_a)
+    manager.begin_step()
     blocks_a = manager.allocate_slots(req_a, 32, 0, computed_a)
     assert blocks_a is not None
 
@@ -176,6 +178,7 @@ def test_finish_path_free_alone_preserves_cache_entries():
     common_token_ids = [i for i in range(2) for _ in range(block_size)]
     req_a = make_request("a", common_token_ids, block_size, sha256)
     computed_a, _ = manager.get_computed_blocks(req_a)
+    manager.begin_step()
     manager.allocate_slots(req_a, 32, 0, computed_a)
     assert len(_uncommitted_blocks_for(manager, "a")) == 2
 
@@ -220,6 +223,7 @@ def test_cache_blocks_committed_kwarg_skips_uncommitted_tracking():
     tokens = [i for i in range(2) for _ in range(block_size)]
     req = make_request("a", tokens, block_size, sha256)
     computed, _ = manager.get_computed_blocks(req)
+    manager.begin_step()
     blocks = manager.allocate_slots(req, 32, 0, computed)
     assert blocks is not None
 
@@ -292,6 +296,7 @@ def test_reset_prefix_cache_clears_uncommitted():
     tokens = [i for i in range(2) for _ in range(block_size)]
     req = make_request("a", tokens, block_size, sha256)
     computed, _ = manager.get_computed_blocks(req)
+    manager.begin_step()
     manager.allocate_slots(req, 32, 0, computed)
 
     # Sanity: eager registration populated _uncommitted.
@@ -316,6 +321,7 @@ def test_reset_prefix_cache_clears_uncommitted():
     req2 = make_request("b", tokens, block_size, sha256)
     computed2, _ = manager.get_computed_blocks(req2)
     assert _ == 0  # cache was cleared by reset
+    manager.begin_step()
     manager.allocate_slots(req2, 32, 0, computed2)
     assert len(_uncommitted_blocks_for(manager, "b")) == 2
 
diff --git a/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py b/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py
index 77d629729776..bbb25d82c772 100644
--- a/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py
+++ b/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py
@@ -438,9 +438,9 @@ def evict_blocks_spy(block_ids):
     original_cache_blocks = recompute_scheduler.kv_cache_manager.cache_blocks
     cache_blocks_calls = []
 
-    def cache_blocks_spy(req, num_tokens):
+    def cache_blocks_spy(req, num_tokens, **kwargs):
         cache_blocks_calls.append((req.request_id, num_tokens))
-        return original_cache_blocks(req, num_tokens)
+        return original_cache_blocks(req, num_tokens, **kwargs)
 
     with patch.object(
         recompute_scheduler.kv_cache_manager, "cache_blocks", cache_blocks_spy
diff --git a/vllm/v1/core/block_pool.py b/vllm/v1/core/block_pool.py
index cb461a7af5fc..57c706f64f4a 100644
--- a/vllm/v1/core/block_pool.py
+++ b/vllm/v1/core/block_pool.py
@@ -275,14 +275,12 @@ def cache_full_blocks(
             [] if self.enable_kv_cache_events else None
         )
 
-        # Where to record this request's new registrations for later
-        # rollback_uncommitted. ``None`` means skip tracking (the caller
-        # passed committed=True, signalling worker-confirmed bytes).
-        if self._suppress_uncommitted_tracking:
+        # Track new registrations in the current step's bucket so they can
+        # be rolled back on preempt. Skipped when committed=True (bytes are
+        # worker-confirmed) or no step is open (caller bypassed schedule()).
+        if self._suppress_uncommitted_tracking or not self._uncommitted:
             uncommitted_bucket: list[KVCacheBlock] | None = None
         else:
-            if not self._uncommitted:
-                self._uncommitted.append({})
             uncommitted_bucket = self._uncommitted[-1].setdefault(
                 request.request_id, []
             )