Skip to content

[UMBP] Enable partial hash hits by default for Mamba hybrid models - #2

Open
isytwu wants to merge 1 commit into
mainfrom
umbp-hybrid-partial-hash-default
Open

isytwu wants to merge 1 commit into
mainfrom
umbp-hybrid-partial-hash-default

Conversation

@isytwu

@isytwu isytwu commented Nov 4, 2026

Copy link
Copy Markdown
Owner

Internal review draft. Remove this note before submitting upstream. The Kimi K3 c70 runs are still in progress, so their results are pending. Links to sibling PRs point at this fork and need the upstream numbers later. Confirm the AI-assistance statement and checklist after review.

Overview

Default enable_partial_hash_hits to true when the model has Mamba "align" groups and prefix_match_unit is finer than the scheduler block. An explicit setting still wins.

Claims

  • Mamba hybrids configured with a fine prefix_match_unit restore 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_unit 128), external hits were close to zero without the option.
  • Other models keep the previous default.

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_*.py gives 234 passed. The new test covers the default for an align hybrid with a fine unit, an explicit false, 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_unit is 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_async keeps 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-checklist skill. (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.

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>
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.

1 participant