[Bugfix] Flag every KV cache group holding a separately prefixed drafter's layers as a draft group - #56026
[Bugfix] Flag every KV cache group holding a separately prefixed drafter's layers as a draft group#56026Windless84 wants to merge 1 commit into
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
|
👋 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. 🚀 |
…l-groups draft fallback only when the block drop is on Qwen4ExpMTP registers its attention layers under "mtp." after the target's "model." layers, but no group carried a spec marker and the DeepseekV4 positional fallback is model-gated, so _annotate_eagle_groups flagged nothing, the coordinator treated every KV group as a draft group and Mamba groups could never serve a prefix hit (kv_cache_utils.py:1893 warning; zero hits across a full replay). Flag every group holding a layer under the drafter's distinct top-level prefix; drafters sharing the target's prefix keep the old conservative fallback. The warning now keys on use_eagle_block_drop(), which is what KVCacheManager actually receives. Claude-Session: https://claude.ai/code/session_01LqPLAsuNFXRF7yoojyna8j Signed-off-by: Windless84 <15976122+Windless84@users.noreply.github.com>
506a4bd to
4a365de
Compare
|
Tested exact head The patch applied cleanly. Its annotation behavior check passed: the Mamba group remained unflagged, the separately prefixed drafter group was flagged, same-prefix layouts remained unflagged, and block-drop-off left the groups untouched. The live Fresh-server 50,011-token cold/replay results:
This verifies #56026 fixes the DFlash group-annotation failure. The default-drop miss is a separate write-side issue; disabling the drop recovers reuse immediately. |
2 participants
Purpose
Follow-up to #52047, complementary to #55390.
On Qwen3.8-Flash-Next (
Qwen4ExpForConditionalGeneration: GDN + QSA + MTP, packed grouping path) no KV cache group is ever identified as the drafter's, so the coordinator falls back to flagging every group, the four Mamba groups included, as a draft group. Effects on a live server:kv_cache_utils.py:1893warning on every start.Why the existing rules miss it:
non_causal_multi_token_decode) is DSpark-only.model_type == "deepseek_v4". [Bugfix] Annotate MTP draft KV cache groups positionally on the hybrid grouping path #55390 generalizes the gate tomethod == "mtp"but, as its docstring notes, flags only the group holding the very last layer. The Qwen4Exp MTP block registers three caches undermtp.(QSA attention KV, compressed-keyMLAAttentionSpec, ringCircularBufferSpec) that land in two packed groups, so the group holding the drafter's attention KV can still be left unflagged.Fix
Rule 3 in
_annotate_eagle_groups: when no group was flagged, and the last registered layer's top-level module prefix differs from the first layer's (a separately built drafter registered after the target), flag every group holding a layer under that prefix. Drafters that reuse the target's prefix (EAGLE heads undermodel.) are not detected and keep the conservative all-groups fallback, as before._warn_if_unannotated_eagle_mambanow keys onuse_eagle_block_drop(), which is whatKVCacheManageractually receives; withdisable_eagle_block_dropno group is a draft group and the warning was spurious.If #55390 lands first this rebases on top of it cleanly (rule 3 only runs when rules 1 and 2 flagged nothing).
Test Plan
New:
test_annotate_eagle_groups_flags_separately_prefixed_drafter_layers(draft layer undermtp.flags its group and leaves the Mamba group alone; same-prefix layout flags nothing; no block drop configured leaves everything untouched).Test Result
106 passedonmain(d290609), CPU.Live, Qwen3.8-Flash-Next NVFP4 on one RTX PRO 6000, MTP 3, 262,144 context, 240 blocks of 1,600 tokens (3x12 GDN, 1 PLE, 13 QSA rings, 13 QSA main+compressed):
mtp.layers)Footprint unchanged (140K prompt = 113 of 240 blocks before and after). Note that the first request extending a prefix still misses with the trailing-block drop on, because the align-mode Mamba checkpoint sits on the last full block while the attention hit is offered one block lower; that is #52244's write-side fix, and the numbers above were taken with
disable_eagle_block_dropas a stand-in.Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.https://claude.ai/code/session_01LqPLAsuNFXRF7yoojyna8j