Skip to content

Support topk > 1 tree drafting for mamba/hybrid-linear models on spec v2#27463

Merged
hnyls2002 merged 6 commits into
mainfrom
lsyin/spec-v2-mamba-topk
Jun 8, 2026
Merged

Support topk > 1 tree drafting for mamba/hybrid-linear models on spec v2#27463
hnyls2002 merged 6 commits into
mainfrom
lsyin/spec-v2-mamba-topk

Conversation

@hnyls2002

@hnyls2002 hnyls2002 commented Jun 6, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Mamba / hybrid-linear models (GDN — Qwen3-Next; Mamba2 — Nemotron-H) with --speculative-eagle-topk > 1 previously 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-aware last_correct_step_indices by gathering each request's last accepted node from accept_index, replacing the chain-only accept_lens - 1; drop the topk != 1 guard (and the now-unused verify_input arg).
  • speculative_hook.py: remove the mamba + topk > 1 -> spec v1 fallback.

Hybrid KV pool

  • HybridLinearKVPool now sets enable_kv_cache_copy (matching the MHA / SWA / FP4 pools), so the topk > 1 tree path's move_kv_cache (accepted-branch KV compaction) is initialized for hybrid models. Without it the tree path hit KV copy not initialized.

Guards & docs

  • LightningAttentionBackend (seg_la): raise NotImplementedError on topk > 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.
  • Correct the stale "Mamba2 tree verify is structurally unsupported" comments in the attention unit-test suite — the production SSM kernel does consume the parent-indices plumbing; the unit test skips topk > 1 only for lack of a tree-aware reference.

Validation

  • TestQwen3NextMTPTopk (Qwen3-Next-80B, NEXTN, topk=4) now runs on spec v2.
  • H200, topk=4, accept length ~4: Qwen3-Next-80B gsm8k 0.96, Nemotron-H-120B gsm8k 0.965.

CI States

Latest PR Test (Base): ❌ Run #27088568107
Latest PR Test (Extra): ✅ Run #27088568046

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

@hnyls2002

Copy link
Copy Markdown
Collaborator Author

/rerun-test test/registered/models_e2e/test_qwen3_next_models_mtp.py

@github-actions

github-actions Bot commented Jun 6, 2026

Copy link
Copy Markdown
Contributor

Results for /rerun-test test/registered/models_e2e/test_qwen3_next_models_mtp.py:

🚀 4-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/models_e2e/test_qwen3_next_models_mtp.py

@hnyls2002

Copy link
Copy Markdown
Collaborator Author

/tag-and-rerun-ci

@github-actions github-actions Bot added the run-ci label Jun 7, 2026
@hnyls2002 hnyls2002 changed the title Support topk > 1 tree drafting for mamba/hybrid-linear models on spec v2 Support topk > 1 tree drafting for mamba/hybrid-linear models on spec v2 Jun 7, 2026

@hanming-lu hanming-lu left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

LGTM as long as test_qwen3_next_models_mtp.py‎ passes

@hnyls2002
hnyls2002 merged commit f68c796 into main Jun 8, 2026
146 of 162 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants