[Bugfix][Core] Skip Mamba partial-tail stops after mid-block hits - #51763
Closed
ivanium wants to merge 1 commit into
Closed
[Bugfix][Core] Skip Mamba partial-tail stops after mid-block hits#51763ivanium wants to merge 1 commit into
ivanium wants to merge 1 commit into
Conversation
In "align" mamba cache mode `_mamba_block_aligned_split` adds one extra stop at the prompt's last hash boundary so the partial-tail entry can be registered. That boundary is off the block grid by construction whenever `hash_block_size < block_size`, which is fine when the chunk *starts* on the grid -- the state at the start was materialized and the tail write is the only unaligned one. It is not fine when the request resumes mid-block. Then the block boundary already crossed was never materialized, and the extra stop cuts the chunk at a second unaligned position instead of re-aligning, leaving the SSM state wrong for the rest of the prefill. The production config that hit this is Kimi-K3: block_size 1536, `--prefix-match-unit 128`, external prefix hit at 24960 = 16*1536 + 384, prompt 25297. The re-align stop 26112 is past `last_cache_position` (24576) so it is suppressed; 24576 is behind the resume point; the tail boundary 25216 is left as the only stop and prefill is cut 640 tokens off the grid. Observed end-to-end as a silent HTTP 200 with `finish_reason: stop`, empty content, and the model resuming from the middle of a structural marker it never opened. Require a block-aligned start for that stop. Aligned starts keep registering the partial tail, so the existing behaviour and its test are unchanged; a mid-block resume now runs the fresh tail in one chunk, which is what the measured-clean configurations already did. Signed-off-by: Roger Wang <hey@rogerw.io> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
ivanium
marked this pull request as ready for review
August 11, 2026 05:17
ivanium
requested review from
ApostaC,
WoosukKwon,
alexm-redhat,
heheda12345,
njhill,
orozery,
robertgshaw2-redhat and
ywang96
as code owners
August 11, 2026 05:17
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Collaborator
Author
|
close as superseded by #51766 which is a better fix |
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
2 participants
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
Credit: @ywang96 found this issue and authored the original fix. This PR ports his change onto current upstream
main.In Mamba
"align"cache mode, a request resuming from a fine-grained prefix hit mid-block could still stop at the prompt's partial-tail hash boundary. That second off-grid cut leaves the SSM state misaligned and can corrupt the continuation.The fix only uses the partial-tail stop when the chunk starts on a Mamba block boundary. The tradeoff is that, after a mid-block hit, vLLM no longer forces a stop at the prompt's final hash boundary. It therefore will not always materialize and cache the precise end of the prompt, which can reduce prefix-cache granularity and reuse. This prioritizes correct SSM state; aligned starts retain the existing cache behavior.
This is not duplicate work: no open PR addresses this condition. #51113 fixes the adjacent block-realignment stop, not this distinct partial-tail stop.
Test Plan
Test Result
.venv/bin/python -m pytest tests/v1/core/prefix_cache/ -q— 33 passed..venv/bin/pre-commit run --files vllm/v1/core/sched/scheduler.py tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py— passed.OpenAI Codex was used to port and validate this change. The human submitter must review every changed line and the test results before marking the PR ready.
Essential Elements of an Effective PR Description Checklist
BEFORE SUBMITTING, PLEASE READ https://docs.vllm.ai/en/latest/contributing (anything written below this line will be removed by GitHub Actions)