[hnyls2002]

Summary

• Run the EAGLE draft-extend forward as a captured CUDA graph for FlashInfer (add FlashInferAttnBackend to the draft-extend graph whitelist), replacing eager per-layer launches.
• ~10% online median E2E latency improvement (request-rate 8, Llama-2-7B EAGLE topk=4); accept length unchanged.

Changes

• Graph capture: add is_draft_extend_v2() branches to _prepare_cuda_graph_metadata / init_forward_metadata_out_graph so the draft-extend forward replays as a causal paged-prefill graph.
• Fixed qo layout: EagleDraftExtendInput.generate_attn_arg_prefill uses a constant num_draft_tokens qo layout (required for graph capture/replay), matching the padded draft-extend input.
• Drop a per-step D2H: pass host-known extend_prefix_lens_cpu into FlashInferIndicesUpdaterPrefill.update, replacing paged_kernel_lens.sum().item() with a host-side sum.

Notes

• FR-Spec (speculative_token_map) falls back to eager: the draft-extend graph buffers are sized to the full draft vocab and are incompatible with the reduced (hot) vocab head.

---

CI States

Latest PR Test (Base): ❌ Run #27916144655
Latest PR Test (Extra): 🚫 Run #27916144559

[gemini-code-assist]

Code Review

This pull request adds support for the FlashInfer attention backend in the Eagle draft-extend speculative execution path, enabling CUDA graph capture and replay by using a fixed query-output (qo) layout. It also optimizes performance by passing host-side prefix lengths (extend_prefix_lens_cpu) to avoid device-to-host synchronization. However, the review highlights a potential runtime crash: the update_sliding_window and update_cross_attention methods in FlashInferIndicesUpdaterPrefill need to be updated to accept the new extend_prefix_lens_cpu parameter to prevent TypeError exceptions when those features are active.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[hnyls2002]

/rerun-test test_spec_eagle_topk.py test_spec_eagle.py test_bench_serving_1gpu_large.py

[github-actions]

Results for /rerun-test test_spec_eagle_topk.py test_spec_eagle.py test_bench_serving_1gpu_large.py:

🚀 1-gpu-5090 (2 tests): ❌ View workflow run

cd test/ && python3 registered/spec/eagle/test_spec_eagle_topk.py
cd test/ && python3 registered/spec/eagle/test_spec_eagle.py

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

cd test/ && python3 registered/perf/test_bench_serving_1gpu_large.py

[github-actions]

Results for /rerun-test test_spec_eagle_topk.py test_spec_eagle.py test_bench_serving_1gpu_large.py:

🚀 1-gpu-5090 (2 tests): ❌ View workflow run

cd test/ && python3 registered/spec/eagle/test_spec_eagle_topk.py
cd test/ && python3 registered/spec/eagle/test_spec_eagle.py

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

cd test/ && python3 registered/perf/test_bench_serving_1gpu_large.py

[hnyls2002]

/rerun-test test_spec_eagle_topk.py test_spec_eagle.py

[github-actions]

Results for /rerun-test test_spec_eagle_topk.py test_spec_eagle.py:

🚀 1-gpu-5090 (2 tests): ✅ View workflow run

cd test/ && python3 registered/spec/eagle/test_spec_eagle_topk.py
cd test/ && python3 registered/spec/eagle/test_spec_eagle.py