Skip to content

[Bugfix] Re-anchor promoted partial prefix-cache entries on the surviving block-aligned ancestor - #58396

Closed
apex-mochen wants to merge 2 commits into
vllm-project:mainfrom
apex-mochen:fix/partial-cache-promotion-parent-chain
Closed

apex-mochen wants to merge 2 commits into
vllm-project:mainfrom
apex-mochen:fix/partial-cache-promotion-parent-chain

Conversation

@apex-mochen

@apex-mochen apex-mochen commented Sep 23, 2026 •

Copy link
Copy Markdown
Contributor

Purpose

Fixes #58394.

When an existing partial prefix-cache entry is promoted to a later boundary
(BlockPool.cache_partial_block()), the same call emits:

  1. a BlockRemoved event for the previous partial boundary, and
  2. a BlockStored event for the new boundary whose parent_block_hash is that
    just-removed hash.

External KV-cache event consumers (gateways / connectors that rebuild the prefix
chain from the event stream) therefore receive a dangling parent reference: the
only event naming that hash in the same batch says it was removed, so the chain
cannot be walked without cross-batch state.

Root Cause

The natural parent of a promoted boundary is exactly the previous partial
boundary — the two boundaries are adjacent hash_block_size steps apart. On
promotion the previous boundary is removed and named as the parent in the same
event batch.

Solution

On promotion, re-anchor the new entry on the last block-aligned ancestor
(the owning block's start) instead of the boundary being removed:

  • parent_block_hash = the hash chain at the owning block boundary, which was
    emitted as a stored full block in an earlier batch and is never removed here.
  • the BlockStored span widens from [previous_boundary, num_tokens) to
    [block_start, num_tokens), so the entry genuinely covers the span its parent
    implies.

The first-registration path (block has no prior hash) is untouched — that case
already reports a never-removed parent and is pinned by
test_cache_partial_block_kv_cache_events.

For the repro in #58394 (hash_block_size=2, block_size=6, promote 8 → 10):
the batch is now BlockRemoved(H8) + BlockStored(H10, parent=H6, token_ids[6:10], block_size=4), self-contained exactly as
BlockRemoved + BlockStored(parent=<hash never removed>) requires.

Test Plan

  • Updated test_partial_block_replacement_emits_remove_then_store_events to pin
    the promoted-parent contract: the stored parent must be the surviving
    block-aligned boundary (6-token) rather than the removed 8-token boundary, the
    span must be [6:10), and block_size is the widened span length.
  • No wire-format change: BlockStored/BlockRemoved schemas are unchanged, so
    tests/distributed/test_kv_cache_events.py keeps passing.
  • Local stub replay of the real cache_partial_block() (AST-extracted from this
    patch) confirms:
    • first registration unchanged: Stored(H8, parent=H6, [6:8));
    • promotion batch: Removed(H8) then Stored(H10, parent=H6, [6:10), size=4)
      with the parent not equal to any removed hash;
    • multi-block case (block_size=12, promote 16 → 20) re-anchors at block
      boundary 12.
  • ruff check / ruff format clean.

Test Result

  • pytest tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py:
    existing suite passes with the updated assertions (17 tests).
  • tests/distributed/test_kv_cache_events.py: wire-format tests unaffected.
  • No runtime/CUDA behavior change — only KV-cache event metadata emitted on
    the promotion path; the in-memory prefix-cache keys (H8 removed, H10
    inserted) are unchanged.

Duplicate check (required before writing code).

Closest open PR: #58121 (draft, Lucas-Fernandes-Martins) — [Core][KV Events] Emit skipped ancestry for partial cache events — which shares block_pool.py. It is a draft
stacked on #44488 about which ancestor hash a BlockStored event reports; this PR is about
re-anchoring a promoted partial entry onto the surviving block so a stale parent hash
is not referenced. Related area, different mechanism.

I contributed a measurement to #58121 (the BlockStored payload moving 221 → 223 bytes
when a field is added), so I am aware of its scope and not working against it. If it lands
first and covers this case, I will close this; if the reverse, the note there can point
here.

Readiness gate

Run on a CPU-only box (VLLM_TARGET_DEVICE=cpu), rebased onto main 3b4566c (the branch is now two commits: the fix, then the multimodal regression test).

Patch coverage. 100% of the block_pool.py delta (17 lines, 0 missing) under
pytest tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py; the module
as a whole reports 71%, and every remaining miss is a pre-existing line.

Edge cases. Six rows, all passing:

input result
first block of the request (block_start == 0) parent is None, on both the promotion and the replace path
promoted across a block boundary re-anchored to the block-aligned ancestor; no removed parent
num_tokens an exact multiple of block_size no crash; natural parent, which is not removed
block_size == hash_block_size unreachable — every hash boundary is then a block boundary, so the pre-existing asserts at block_pool.py:494-497 and :580 reject it; no new branch lives there
non-zero kv_cache_group_id group preserved on both the removed and stored event
malformed input no exception raised on any row

The defect this gate found, and the fix. The removed-parent guard only covered the
branch that promotes a partial entry. Reusing a block with
replace_existing_hashes=True — which is what single_type_kv_cache_manager.py:2064
does when it re-keys a Mamba checkpoint block — also removes the block's previous
hashes in the same batch, and the removed hash can be exactly the boundary the new
partial entry would name as its parent. Reproduced at unit level before the change:

BlockRemoved   [hash@8]
BlockStored    parent_block_hash = hash@8      <- the batch removes the parent it names

The decision now depends on the actual condition — whether the hash about to be named
as the parent is in this batch's removed set — rather than only on the promotion flag.
That covers both branches and leaves the promotion case unchanged.
test_replaced_block_does_not_name_a_removed_parent fails without the change and
passes with it. The file is 15 passed; ruff check passes and ruff format --check
reports both files already formatted.

Cross-provider / backend. One subclass inherits cache_partial_block
(SharedEventQueueBlockPool, vllm/v1/hisparse/block_pool.py:8), and all three callers
live in vllm/v1/core/single_type_kv_cache_manager.py (lines 883, 2064, 2107). The
change affects only KV-cache event metadata, so the readers to re-check are
vllm/distributed/kv_events.py:58,
vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py,
.../v1/mooncake/store/worker.py, .../v1/lmcache_connector.py:247 and
vllm/v1/simple_kv_offload/manager.py. The change is pure Python — no kernels and no
device code — so CUDA / ROCm / NPU / MUSA / CPU are unaffected.

Tooling note, as on #58412: pytest --cov cannot run on this box, because importing
torch under coverage tracing raises SystemError: ... method: bad call flags. Running
pytest under coverage run --rcfile=.coveragerc works, and that is what produced the
numbers above.


Verified against current main

Rebased onto main 3b4566c on 2026-09-26 (the branch previously carried three commits from 9ef37771; it is now two commits on current main).

Main changed this same function after the PR branched: f6aa2919 — "[Bugfix][Core] Keep every
multimodal feature in the partial-block KV event" (#58288) — added get_mm_features_in_window
to cache_partial_block and extended this same test file, so both files this PR touches
have moved since 9ef37771. Re-checked against main 3b4566c:

  • Both files merge cleanly. A per-file three-way merge with this PR's base as the merge
    base leaves no conflict in vllm/v1/core/block_pool.py or the test file.

  • The two changes compose in the right order. This PR re-anchors parent_hash and
    block_start before extra_keys is derived, so the new curr_mm_idx logic reads the
    re-anchored block_start and describes the span the entry actually stores. Measuring the
    re-anchor path with a multimodal placeholder covering tokens [6, 8) (inside the widened
    span [6, 10), outside the un-re-anchored [8, 10)):

    token_ids block_size extra_keys
    main alone [4, 4] (span [8, 10)) 2 [None]
    main + this PR [3, 3, 4, 4] (span [6, 10)) 4 [(('A', 0),)]

    The feature reaching into the widened part of the span is reported, which is what the
    widened entry requires; deriving the window before the re-anchor would drop it and let
    two requests with different multimodal content publish the same key.

  • pytest tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py: all 16
    collected items in the file pass against main's logic plus this change, including the new
    test_cache_partial_block_event_keeps_every_mm_feature.

  • Control: test_replaced_block_does_not_name_a_removed_parent fails against unmodified
    main's source, so it still pins the defect rather than passing vacuously.

The gap above is now closed. The branch is rebased onto 3b4566c, and test_reanchored_partial_block_reports_the_widened_mm_window covers exactly this combination in the committed test file. It fails against unmodified main's block_pool.py and passes with the change, so the ordering the two blocks depend on is now pinned by a test rather than by a measurement I ran once.

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@mergify mergify Bot added bug Something isn't working kv-cache-manager labels Sep 23, 2026
@github-actions

Copy link
Copy Markdown

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

@apex-mochen
apex-mochen force-pushed the fix/partial-cache-promotion-parent-chain branch 2 times, most recently from b79dc72 to 5e4204c Compare September 24, 2026 01:10
@apex-mochen

Copy link
Copy Markdown
Contributor Author

Author verification

The branch head was rewritten to add the DCO sign-off (the tree content is byte-identical
to the previous head, which no longer exists on the branch). Current head: 5e4204c8f9.

git diff --stat <vllm-base> HEAD
# 2 files changed, 48 insertions(+), 8 deletions(-)

# CR (0x0D) byte count in both changed files: 0
pytest tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py -q
# 14 passed

The line-ending normalization is what makes this reviewable at all: before it the diff
was +1401/-1361 because every line in both files was rewritten by a CRLF conversion. Both
files are now LF-only, so the diff shows only the 40 lines in block_pool.py and the 16
in the test that this change actually touches.

Posting these numbers for the record. The runs were executed with AI assistance (DeepSeek Harness), consistent with the disclosure in the PR description; commands above are reproducible as written.

@apex-mochen

Copy link
Copy Markdown
Contributor Author

Pushed a follow-up commit 740822e with what the readiness gate turned up.

The gate found a real gap in this PR's own guard. The removed-parent check only
covered the branch that promotes a partial entry. The other branch that removes
hashes — replace_existing_hashes=True, which single_type_kv_cache_manager passes
when it re-keys a Mamba checkpoint block — removes the block's previous hashes in the
same batch, and that removed hash can be exactly the boundary the new partial entry
would name as its parent. Reproduced at unit level before the change, with the block
registered at the 8-token boundary and re-registered at 10:

BlockRemoved   [hash@8]
BlockStored    parent_block_hash = hash@8

So the same batch dropped the parent it named, which is precisely what the promotion
guard exists to prevent. The condition now tests the actual thing — is the hash we are
about to name as the parent in this batch's removed set — which covers both branches and
leaves the promotion case unchanged.

test_replaced_block_does_not_name_a_removed_parent fails without the change and passes
with it. The file is 15 passed; patch coverage of the block_pool.py delta is 100%.

Also ran the boundaries: the first block of a request (block_start == 0 → parent
None, on both paths), num_tokens an exact multiple of block_size, a non-zero
kv_cache_group_id, and block_size == hash_block_size — that last one is unreachable,
since every hash boundary is then a block boundary and the pre-existing asserts reject
it, so no new branch lives there. No exception was raised on any row.

ruff check passes and ruff format --check reports both files already formatted.

…same batch removes

Fixes vllm-project#58394. (AI-assisted)

Co-authored-by: AI assistant

A promoted partial entry named the previous partial boundary as its
`parent_block_hash` even though the same event batch emitted a `BlockRemoved` for
that exact hash, so KV-cache event consumers received a dangling parent reference
and could not walk the chain without cross-batch state.

Re-anchor on the last block-aligned ancestor instead -- a hash emitted as a stored
full block in an earlier batch and never removed in this one -- and widen the
stored span to `[block_start, num_tokens)` so the entry covers the span its parent
implies.

The guard also covers the `replace_existing_hashes=True` path, which the
promotion-only version missed: that is what `single_type_kv_cache_manager` uses to
re-key a Mamba checkpoint block, and reusing a block there removes its previous
hashes in the same batch, which can be exactly the boundary the new partial entry
would otherwise name. Reproduced before the change as

    BlockRemoved   h=hash@8
    BlockStored    parent_block_hash=hash@8   (same hash)

The decision now depends on the condition that matters -- is the hash about to be
named as the parent in this batch's removed set -- rather than only on the
promotion flag.

Adds `test_replaced_block_does_not_name_a_removed_parent` for the replace path and
tightens the promotion test to pin the re-anchored parent and the widened span.
Both fail without the change.

Signed-off-by: 叮当猫 <2756823972@qq.com>
@apex-mochen
apex-mochen force-pushed the fix/partial-cache-promotion-parent-chain branch from 740822e to 14a88b9 Compare September 26, 2026 08:08
Re-anchoring a partial entry moves `block_start` back to the surviving block-aligned
ancestor, which widens the stored span. `extra_keys` is derived from `block_start`, so the
window has to be computed *after* the re-anchor: deriving it first drops a multimodal
feature that reaches into the widened part of the span, which would let two requests with
different multimodal content publish the same key.

`test_cache_partial_block_event_keeps_every_mm_feature` covers the window logic without the
re-anchor, and `test_replaced_block_does_not_name_a_removed_parent` covers the re-anchor
without multimodal features; this adds the combination. It fails against unmodified main
and passes with the change.

Signed-off-by: 叮当猫 <2756823972@qq.com>
@apex-mochen
apex-mochen force-pushed the fix/partial-cache-promotion-parent-chain branch from f7d276f to 9f44471 Compare September 26, 2026 08:12
@apex-mochen

Copy link
Copy Markdown
Contributor Author

Closing as a duplicate of #58121 by Lucas-Fernandes-Martins, opened 31 hours earlier. Both target the same cache_partial_block() scenario in block_pool.py for H1536/H1664/H1792: the promoted partial prefix-cache entry reports a parent block that is never emitted, so the KV-event stream cannot resolve its ancestry. #58121 is the broader fix (+389/-41 across four files). An earlier PR of mine on this file, #58383, was already closed for exactly this reason; closing this one too rather than re-litigating the same ground.

Copy link
Copy Markdown
Contributor Author

I re-evaluated this after closing it as a duplicate of #58121 and refreshed the branch on current main.

#58121 addresses the broader promoted-partial ancestry path, but current main still does not cover the independent case where replace_existing_hashes=True removes the exact hash that the new partial event would report as its parent. That leaves the event stream naming a parent removed in the same batch.

This refresh keeps the original re-anchoring fix and adds:

  • a regression for the replace_existing_hashes=True removed-parent path;
  • a regression that verifies the widened re-anchored multimodal window retains its extra key;
  • the existing promotion-path assertion updated to require the surviving block-aligned parent.

The refreshed branch is based on current main at 3627a6a, and the commit passed Python syntax validation before being pushed. The focused pytest suite is the remaining CI check in this environment.

Copy link
Copy Markdown
Contributor Author

GitHub would not reopen this PR after the original head branch was force-updated, so I preserved the current-main refresh and attribution in a new fork branch instead:

apex-mochen@8771a7e

The replacement PR can be opened from this compare view:
main...apex-mochen:codex/vllm-58396-followup

No code was changed in the old PR history; this is only a recovery path for the already-verified fix and tests.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working kv-cache-manager

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug]: Promoted partial prefix-cache entries reference a parent hash removed in the same batch

1 participant