Support spec v2 tree drafting (eagle topk>1) with page_size==1#26866
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request updates the Eagle speculative decoding implementation (v2) to support tree drafting (topk > 1). It restricts tree drafting to page_size == 1, fixes an illegal memory access in move_accepted_tokens_to_target_kvcache by iterating over ai_size instead of size, and introduces a compaction mechanism (_compact_accepted_to_front) to align the accepted path to the front of each per-request block. The review feedback highlights an optimization opportunity to reduce VRAM allocation by sizing temporary cache location tensors to ai_size instead of size, and a potential correctness issue where routed_experts_output and indexer_topk_output are not compacted when topk > 1 is active.
|
/rerun-test test_eagle_infer_a.py TestEAGLE3EngineSpecV2Tree |
|
Results for 🚀 ⛔ |
# Conflicts: # test/registered/spec/eagle/test_eagle_infer_a.py
|
/rerun-test test/registered/spec/eagle/test_spec_eagle.py test/registered/spec/eagle/test_spec_eagle_topk.py test/registered/spec/eagle/test_spec_eagle_page.py test/registered/spec/eagle/test_spec_eagle_triton.py test/registered/spec/eagle/test_spec_eagle_fa3.py test/registered/spec/eagle/test_spec_eagle_parity.py test/registered/spec/eagle/test_spec_eagle_stress.py test/registered/spec/eagle/test_deepseek_v3_fp4_mtp_small.py |
|
Results for 🚀 |
|
Results for 🚀 🚀 |
… num_draft_tokens alloc
|
Results for 🚀 🚀 |
|
Results for 🚀 🚀 🚀 |
|
/rerun-test test_eagle_dp_attention.py |
#26997) Co-authored-by: Alison Shao <[email protected]>
1 participant
Routes eagle/eagle3 topk>1 to the spec-v2 (overlap) worker when page_size==1 instead of forcing v1, and fixes the v2 tree path: compact the accepted tree path to the front of each per-req block, wire up the (previously stubbed) accepted-KV move, fix its accept_index OOB, and fix the bonus-token stride for bs>1. Validated on EAGLE3 (Llama-3.1-8B) topk=4: greedy output is token-identical to topk=1 across short/long/concurrent bs>1. topk>1 with page_size>1 still falls back to v1.
CI States
Latest PR Test (Base): ❌ Run #26784184010
Latest PR Test (Extra): ❌ Run #26784183914