Support topk > 1 tree drafting for mamba/hybrid-linear models on spec v2#27463
Merged
Conversation
hnyls2002
requested review from
Fridge003,
Qiaolin-Yu,
Ying1123,
alphabetc1,
hanming-lu,
hzh0425,
ispobock,
merrymercy,
xiezhq-hermann and
yizhang2077
as code owners
June 6, 2026 22:36
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
Collaborator
Author
|
/rerun-test test/registered/models_e2e/test_qwen3_next_models_mtp.py |
Contributor
|
Results for 🚀 |
Collaborator
Author
|
/tag-and-rerun-ci |
topk > 1 tree drafting for mamba/hybrid-linear models on spec v2
yizhang2077
approved these changes
Jun 7, 2026
hanming-lu
approved these changes
Jun 7, 2026
hanming-lu
left a comment
Collaborator
There was a problem hiding this comment.
LGTM as long as test_qwen3_next_models_mtp.py passes
5 tasks
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.
Summary
--speculative-eagle-topk > 1previously fell back to the spec v1 worker. This runs them on spec v2, reusing the tree-drafting machinery from Spec v2 tree drafting (topk>1) with page_size>1 #26972.Changes
Spec v2 mamba tree verify
eagle_worker_v2.py: compute the tree-awarelast_correct_step_indicesby gathering each request's last accepted node fromaccept_index, replacing the chain-onlyaccept_lens - 1; drop thetopk != 1guard (and the now-unusedverify_inputarg).speculative_hook.py: remove themamba + topk > 1 -> spec v1fallback.Hybrid KV pool
HybridLinearKVPoolnow setsenable_kv_cache_copy(matching the MHA / SWA / FP4 pools), so thetopk > 1tree path'smove_kv_cache(accepted-branch KV compaction) is initialized for hybrid models. Without it the tree path hitKV copy not initialized.Guards & docs
LightningAttentionBackend(seg_la): raiseNotImplementedErrorontopk > 1— seg_la has no parent-indices plumbing and verifies a draft tree as a chain, so tree verify is genuinely unsupported there; fail fast instead of silently mis-decoding.topk > 1only for lack of a tree-aware reference.Validation
TestQwen3NextMTPTopk(Qwen3-Next-80B, NEXTN, topk=4) now runs on spec v2.CI States
Latest PR Test (Base): ❌ Run #27088568107
Latest PR Test (Extra): ✅ Run #27088568046