Skip to content

[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
ModelEngine-Group:developfrom
AmirF194:fix/867-blend-partial-pc-overlap-multi-chunk
Open

[Bugfix] Fix duplicate block load when a prefix-cache hit run spans a blend chunk boundary#1181
AmirF194 wants to merge 2 commits into
ModelEngine-Group:developfrom
AmirF194:fix/867-blend-partial-pc-overlap-multi-chunk

Conversation

@AmirF194

@AmirF194 AmirF194 commented Aug 2, 2026

Copy link
Copy Markdown

Purpose

Fixes #867 and #906: UCMBlendConnector._get_req_chunk_hit dispatches the
same 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 more
blocks than the first blend chunk, _get_req_chunk_hit only trimmed
req_chunks_meta[0], with the raw, unclamped pc_hit_blocks count:

first_chunk_meta = req_chunks_meta[0]
first_chunk_meta.update_meta_partial_pc(pc_hit_blocks, self.block_size)
if first_chunk_meta.chunk_tokens_len == 0:
    req_chunks_meta.pop(0)

If pc_hit_blocks exceeds that chunk's own chunk_blks_len,
update_meta_partial_pc drives chunk_blks_len/chunk_tokens_len negative
instead of to 0, so the pop guard never fires and the corrupted chunk
stays in req_chunks_meta. Worse, any following chunk whose blocks also
fall inside the hit run is never trimmed at all, and its store_hits (set
earlier in the function from the same hit run) still reports True for
those blocks. _generate_blend_dispatch_meta then adds them to the load
list a second time, on top of the prefix path's own
vllm_block_ids[:pc_hit_block_num] slice.

Fix: walk req_chunks_meta from the front, clamping each chunk's trim
to min(remaining_pc_hit_blocks, chunk_blks_len) and consuming the run
across 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_hits assignment to
chunk-relative offsets. I traced the actual defect and found the
store_hits assignment (global-index slicing into the concatenated
prefix+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_meta from this file via
    ast at test time (same technique as test_kv_cache_layout.py, no vLLM
    import needed) and drives them with synthetic two-chunk fixtures.
    test_prefix_hit_run_crossing_chunk_boundary_does_not_corrupt_chunk and
    test_prefix_hit_run_crossing_chunk_boundary_does_not_double_load_block
    fail on unmodified develop (negative chunk_blks_len, block dispatched
    for load twice) and pass on this branch;
    test_prefix_hit_within_first_chunk_still_trims_it guards the ordinary
    case and passes on both.
  • black, isort --profile=black, codespell clean on both changed files
    (python:3.12-slim, matching py-linter's pre-commit hooks).
  • Not verified: the full vLLM/GPU inference path (_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.

… 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

Copy link
Copy Markdown
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.

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] Incorrect token generation after prefix cache hit (UCMBlendConnector, possible KV cache misalignment)

1 participant