Repository navigation
Conversation
A Mamba "align" model checkpoints its state only on prompt-tail hash boundaries, which rarely fall on a scheduler block edge, so with the default block-granular lookup a hybrid model almost never gets an external hit. Default enable_partial_hash_hits to true when the KV cache has Mamba "align" groups and the prefix-match unit is finer than the scheduler block. An explicit setting still wins. Signed-off-by: yutongwu <yutongwu@amd.com>
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.
Overview
Default
enable_partial_hash_hitsto true when the model has Mamba "align" groups andprefix_match_unitis finer than the scheduler block. An explicit setting still wins.Claims
prefix_match_unitrestore from UMBP without extra configuration. Their checkpoints sit on prompt-tail hash boundaries, which block-granular lookups almost never reach. On Kimi K3 (12288-token attention blocks,prefix_match_unit128), external hits were close to zero without the option.Validation
MI355X,
vllm/vllm-openai-rocm:nightly-40941f2c(MORI 1.2.3.post1 with UMBP).Unit tests:
pytest -q tests/v1/kv_connector/unit/test_umbp_*.pygives 234 passed. The new test covers the default for an align hybrid with a fine unit, an explicitfalse, a hybrid without a fine unit, and an attention-only model.End to end, together with #1 (Granite 4.0-H 350M,
prefix_match_unit=16, 400-token blocks, option unset): main restores none of the 832 tokens that a local GPU prefix hit reuses. With both changes it restores all 832, the same as setting the option explicitly.Details
prefix_match_unitis documented as the finest boundary a prefix-cache hit can land on, so a user who sets it below the block size asks for hits inside blocks. Hybrids are the case where UMBP can serve those hits, since it stores Mamba checkpoints and (with #1) the attention page at each boundary. Attention-only models store only full blocks, so the default stays off for them.Tradeoff: partial hits look up one key per hash unit instead of one per block. Kimi K3 sends up to about 37K keys per lookup across ranks.
lookup_asynckeeps that off the scheduler thread.Not a duplicate: no open PR changes this default.
AI assistance was used for this change and description. The submitter reviewed every line and ran the tests above.
Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.