Repository navigation
Conversation
…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>
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: 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. 🚀 |
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:
--kv-offloading-size, LMCache, NIXL/P-D,--kv-transfer-config);"disable_lookahead_block_hashes": true.This PR fixes those edge cases.
Claims
needs_eagle_block_drop()now returns False for it, as it already does for DFlash/DSpark.SlidingWindowManager.reachable_block_masknow keepsalignment_tokens // block_sizemore 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.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:
nvidia/Gemma-4-31B-IT-NVFP4withgoogle/gemma-4-31B-it-assistantas MTP3 drafterTRITON_ATTNbackend, FP8 KV--prefix-cache-retention-interval 81923fbc8d8b4disable_lookahead_block_hashesTo see what each change does on its own, I also applied them separately (MTP3 +
disable_lookahead_block_hashes):Results:
disable_lookahead_block_hashes. The check used 8 GSM8K prompts and 6 long-prefix requests.tests/v1/core/test_prefix_caching.py,tests/v1/core/prefix_cache/,tests/v1/core/test_single_type_kv_cache_manager.pyand the eagle-block-drop tests intests/test_config.pyall pass, with no failures before or after this change.pre-commit run --fileson the changed files passes.New tests:
tests/test_config.py::test_gemma4_mtp_needs_no_eagle_block_droptests/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):
shift = 1, so for the checkpoint at 16384 it caches the run whose peek block sits past the boundary.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_sizeand 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_sizemore blocks per SWA group: 17 → 19 for Gemma-4. Dense retention (retention_intervalunset or equal to the alignment) is unchanged.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.