[Bugfix] Fix duplicate block load when a prefix-cache hit run spans a blend chunk boundary - #1181
Open
AmirF194 wants to merge 2 commits into
Open
Conversation
… blend chunk boundary UCMBlendConnector._get_req_chunk_hit trims only req_chunks_meta[0] for a partial prefix-cache overlap, using the raw global pc_hit_blocks count with no clamp to that chunk's own length. When the prefix-cache hit run covers more blocks than the first chunk (content already cached as a plain prefix before this request re-chunked it), chunk_blks_len goes negative and the chunk is never popped (the pop guard checks == 0), while every following chunk overlapping the run is left untouched with store_hits still True for blocks the prefix path already covers. _generate_blend_dispatch_meta then dispatches those blocks for LOAD twice: once via the prefix-path slice, once via the untouched chunk's own hit list, two different store keys racing to fill the same destination vLLM block. Fixes ModelEngine-Group#867 Fixes ModelEngine-Group#906
AmirF194
requested review from
Infinite666,
Wwwzff,
harrisonyhq,
mag1c-h,
qyh111 and
ygwpz
as code owners
August 2, 2026 18:50
Author
|
No rush, but checking in since it's been about a week. The branch is behind main now (not conflicting, so no rebase needed on my end) and CI hasn't run yet, it's sitting behind this repo's fork-PR approval gate. Happy to rebase or narrow the scope if that's useful. |
…rtial-pc-overlap-multi-chunk
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
Fixes #867 and #906:
UCMBlendConnector._get_req_chunk_hitdispatches thesame physical vLLM block for a KV load twice when a prefix-cache hit run
spans a blend chunk boundary, two different store keys racing to fill one
destination block.
Modifications
Root cause: when a prefix-cache hit run (
pc_hit_blocks) covers moreblocks than the first blend chunk,
_get_req_chunk_hitonly trimmedreq_chunks_meta[0], with the raw, unclampedpc_hit_blockscount:If
pc_hit_blocksexceeds that chunk's ownchunk_blks_len,update_meta_partial_pcdriveschunk_blks_len/chunk_tokens_lennegativeinstead of to
0, so the pop guard never fires and the corrupted chunkstays in
req_chunks_meta. Worse, any following chunk whose blocks alsofall inside the hit run is never trimmed at all, and its
store_hits(setearlier in the function from the same hit run) still reports
Trueforthose blocks.
_generate_blend_dispatch_metathen adds them to the loadlist a second time, on top of the prefix path's own
vllm_block_ids[:pc_hit_block_num]slice.Fix: walk
req_chunks_metafrom the front, clamping each chunk's trimto
min(remaining_pc_hit_blocks, chunk_blks_len)and consuming the runacross as many leading chunks as it actually covers, popping each one
fully absorbed and stopping at the first chunk only partially covered. For
a hit run confined to the first chunk (the common case) this is unchanged
from today's behavior.
This is a different, smaller fix than the earlier #910 (self-closed, no
stated reason): #910 rewrote the chunk-hash lookup to always query every
chunk hash independently and switched
store_hitsassignment tochunk-relative offsets. I traced the actual defect and found the
store_hitsassignment (global-index slicing into the concatenatedprefix+chunk lookup results) is index-correct on its own; the bug is
isolated to the single-chunk trim step, so the fix stays there.
Test
test/suites/Unit/test_blend_connector_chunk_hit.py(new): extracts_get_req_chunk_hit/_generate_blend_dispatch_metafrom this file viaastat test time (same technique astest_kv_cache_layout.py, no vLLMimport needed) and drives them with synthetic two-chunk fixtures.
test_prefix_hit_run_crossing_chunk_boundary_does_not_corrupt_chunkandtest_prefix_hit_run_crossing_chunk_boundary_does_not_double_load_blockfail on unmodified
develop(negativechunk_blks_len, block dispatchedfor load twice) and pass on this branch;
test_prefix_hit_within_first_chunk_still_trims_itguards the ordinarycase and passes on both.
black,isort --profile=black,codespellclean on both changed files(
python:3.12-slim, matchingpy-linter's pre-commit hooks)._post_process_chunk_cache,delta-rope application) that [Bug] Incorrect token generation after prefix cache hit (UCMBlendConnector, possible KV cache misalignment) #867 originally reported. This VPS has no GPU
and no vLLM install; the fix and test are scoped to the indexing/dispatch
defect I could reproduce and verify directly.