Skip to content

[Bugfix][Core] Fix Gemma-4 MTP prefix-cache hit loss when lookahead block hashes are off - #65263

Open
xuanzic wants to merge 1 commit into
vllm-project:mainfrom
xuanzic:gemma4-mtp-prefix-hit
Open

xuanzic wants to merge 1 commit into
vllm-project:mainfrom
xuanzic:gemma4-mtp-prefix-hit

Conversation

@xuanzic

@xuanzic xuanzic commented Nov 3, 2026 •

Copy link
Copy Markdown

Overview

Follow-up to #64808, which reported that Gemma-4 MTP gets about half the prefix-cache hits of MTP0.

On latest main the default setup no longer shows it. Since #61589 (lookahead block hashes), prefix caching without a KV connector turns the EAGLE trailing-block drop off. However, the two underlying bugs still exist:

  • any KV connector (--kv-offloading-size, LMCache, NIXL/P-D, --kv-transfer-config);
  • "disable_lookahead_block_hashes": true.

This PR fixes those edge cases.

Claims

  • Gemma-4 MTP no longer drops the last block of a prefix-cache hit. EAGLE drafters drop it because their draft KV for that block depends on the next token. The Gemma-4 MTP drafter writes no KV of its own (it reuses the target model's), so it has nothing to recompute. needs_eagle_block_drop() now returns False for it, as it already does for DFlash/DSpark.
  • Keep the blocks that MTP requests actually hit at a retention checkpoint. With EAGLE-style drafting, a request that shares a prefix up to a checkpoint can only hit one alignment unit before it (e.g. 16352 instead of 16384 with Triton, where the alignment is 32 tokens). The sliding-window layers only kept the blocks ending at the checkpoint, so that hit missed and fell back to the previous checkpoint (8192). SlidingWindowManager.reachable_block_mask now keeps alignment_tokens // block_size more blocks before each checkpoint (2 for Gemma-4), so the request hits at 16352. This applies to any EAGLE-style drafter on a sliding-window model.
  • No change to the default setup. Without a KV connector, lookahead block hashes already turn the drop off.

Validation

Same setup as #64808. 4 sequential requests sharing a 16384-token prefix, cached tokens per request. Connector rows show local GPU hit + CPU hit.

Setup:

  • B200, TP2
  • nvidia/Gemma-4-31B-IT-NVFP4 with google/gemma-4-31B-it-assistant as MTP3 drafter
  • TRITON_ATTN backend, FP8 KV
  • --prefix-cache-retention-interval 8192
Config main 3fbc8d8b4 This PR
MTP0 0 / 16384 / 16384 / 16384 –
MTP3 (default) 0 / 16384 / 16384 / 16384 0 / 16384 / 16384 / 16384
MTP3 + disable_lookahead_block_hashes 0 / 8192 / 16352 / 16352 0 / 16384 / 16384 / 16384
MTP3 + OffloadingConnector 0 / 8192+8192 / 8192+8192 / 8192+8192 0 / 16384+0 / 16384+0 / 16384+0

To see what each change does on its own, I also applied them separately (MTP3 + disable_lookahead_block_hashes):

What was applied Cached tokens
Nothing (main) 0 / 8192 / 16352 / 16352
Only the second change: keep 2 more blocks at each checkpoint (Gemma-4 still drops the last block) 0 / 16352 / 16352 / 16352
Only the second change, but keeping 1 more block instead of 2 0 / 8192 / 16352 / 16352
Both changes (this PR) 0 / 16384 / 16384 / 16384

Results:

  • Hits. With this PR, Gemma-4 MTP3 matches MTP0 in every row. With OffloadingConnector, all hits are now served from GPU. The periodic-tail change alone removes the 8192 fallback; the remaining 32 tokens are the block drop, which the first change removes.
  • Acceptance. Greedy outputs and acceptance length (3.46) are identical to main's default config with and without disable_lookahead_block_hashes. The check used 8 GSM8K prompts and 6 long-prefix requests.
  • Unit tests. On B200, tests/v1/core/test_prefix_caching.py, tests/v1/core/prefix_cache/, tests/v1/core/test_single_type_kv_cache_manager.py and the eagle-block-drop tests in tests/test_config.py all pass, with no failures before or after this change.
  • Lint. pre-commit run --files on the changed files passes.

New tests:

  • tests/test_config.py::test_gemma4_mtp_needs_no_eagle_block_drop
  • tests/v1/core/test_prefix_caching.py::test_swa_reachable_block_mask_eagle_periodic_tail_reaches_sibling: covers SWA block < alignment.

Both new tests fail on main and pass with this change.

Details

Root cause (see #64808):

  1. Gemma-4 MTP falls into the coordinator fallback that flags every KV group as EAGLE, because its KV-shared draft layers own no group.
  2. The periodic tail uses shift = 1, so for the checkpoint at 16384 it caches the run whose peek block sits past the boundary.
  3. A sibling diverging in [16384, 16384 + alignment) needs the run one alignment unit lower. That run is never cached, so the lookup falls back a full retention interval, to 8192.

Why exclude Gemma-4 MTP in needs_eagle_block_drop() rather than change the fallback. The fallback is a conservative default for any drafter whose layers are not annotated. Changing it would affect all such drafters. This change touches only the one drafter known to write no KV.

Why alignment_tokens // block_size and not one block. One block is enough only when the SWA block equals the hit alignment. Gemma-4 has SWA block 64 vs alignment 128 (FlashInfer) and 16 vs 32 (Triton), so widening by one block still leaves the 8192 fallback (last row above).

Cost. Under EAGLE with sparse retention, each periodic checkpoint keeps alignment_tokens // block_size more blocks per SWA group: 17 → 19 for Gemma-4. Dense retention (retention_interval unset or equal to the alignment) is unchanged.


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.

…lock hashes are off

Gemma-4 MTP writes no draft KV, so skip the EAGLE trailing-block drop for
it. Under EAGLE with sparse retention, keep each periodic sliding-window
tail one hit-alignment unit wider, so a sibling that diverges right after
a boundary still finds its run. Fixes vllm-project#64808 with KV connectors or
disable_lookahead_block_hashes.

AI assistance was used.

Signed-off-by: Vivian Chen <140748220+xuanzic@users.noreply.github.com>

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@github-actions

github-actions Bot commented Nov 3, 2026

Copy link
Copy Markdown

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

This branch has not been deployed

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

Labels

bug Something isn't working kv-cache-manager

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant