[V1][Spec Decode] Apply torch.compile & cudagraph to EAGLE3

[zixi-qi]

Follow up on #17211 to add support for Eagle3 torch.compile & cudagraph by moving out the multiple auxiliary hidden states merge

Test plan:

• offline:

VLLM_TORCH_PROFILER_DIR=/data/users/qizixi/gitrepos/vllm/profiles VLLM_USE_V1=1 python examples/offline_inference/eagle.py --num_spec_tokens 3 --max_num_seqs 1 --num_prompts 10 --method eagle3

• trace before (propose -> ~16ms)

• trace after (propose -> ~6ms)

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[WoosukKwon]

LGTM. Thanks for doing this! It'd be nice if @luyuzhe111 can also take a look.

[WoosukKwon]

@zixi-qi Oh can you please fix the format error in the pre-commit test?

[luyuzhe111]

@zixi-qi Thanks for the PR! LGTM overall and only left some minor suggestions.

Factoring the fc layer out of the compiled model seems to be the right thing to do! In this case fc layer will have to run in eager mode, but it's not causing much overhead based on your profiling right?

[zixi-qi]

Address comments and apply pre-commit lint suggestions

[zixi-qi]

@zixi-qi Thanks for the PR! LGTM overall and only left some minor suggestions.

Factoring the fc layer out of the compiled model seems to be the right thing to do! In this case fc layer will have to run in eager mode, but it's not causing much overhead based on your profiling right?

RE moving fc layer into eager mode: this is mostly just an aten:mm so seems pretty cheap

[luyuzhe111]

@zixi-qi thanks for incorporating the suggestions! final thing, were you able to check the acceptance length with torch.compile & cuda graph matches that with eager mode as in #16937 (comment)?

[zixi-qi]

@zixi-qi thanks for incorporating the suggestions! final thing, were you able to check the acceptance length with torch.compile & cuda graph matches that with eager mode as in #16937 (comment)?

This is a great point, there does seem to be a difference in acceptance rate between with and without eager, below tests are run with mt_bench + patching #16367

• eagle3 eager

--------------------------------------------------
mean acceptance length:         3.52
--------------------------------------------------
acceptance at token 0:1.00
acceptance at token 1:0.74
acceptance at token 2:0.54
acceptance at token 3:0.41
acceptance at token 4:0.30
acceptance at token 5:0.23
acceptance at token 6:0.17
acceptance at token 7:0.13

• eagle3 cudagraph

--------------------------------------------------
mean acceptance length:         2.79
--------------------------------------------------
acceptance at token 0:1.00
acceptance at token 1:0.55
acceptance at token 2:0.38
acceptance at token 3:0.29
acceptance at token 4:0.21
acceptance at token 5:0.16
acceptance at token 6:0.12
acceptance at token 7:0.09

Eagle pass does not have the same issue, so perhaps moving the fc had a negative impact on numerics?

• eagle eager

--------------------------------------------------
mean acceptance length:         2.47
--------------------------------------------------
acceptance at token 0:1.00
acceptance at token 1:0.67
acceptance at token 2:0.39
acceptance at token 3:0.21
acceptance at token 4:0.11
acceptance at token 5:0.05
acceptance at token 6:0.03
acceptance at token 7:0.01

• eagle cudagraph

--------------------------------------------------
mean acceptance length:         2.47
--------------------------------------------------
acceptance at token 0:1.00
acceptance at token 1:0.68
acceptance at token 2:0.39
acceptance at token 3:0.20
acceptance at token 4:0.11
acceptance at token 5:0.06
acceptance at token 6:0.03
acceptance at token 7:0.02

[WoosukKwon]

@zixi-qi My guess is that this can be related to vLLM's compilation cache...

[luyuzhe111]

@zixi-qi thanks for sharing the testing results. if you think there is something wrong with cuda graph implementation for EAGLE3, you can add VLLM_LOGGING_LEVEL=DEBUG for more checks (such as input address).

[luyuzhe111]

@zixi-qi oh actually, you forgot to capture the EAGLE3 model here. if you look at the profiling traces more closely, you probably should see there was no cuda graph launch for eagle3 previously.

[zixi-qi]

@zixi-qi oh actually, you forgot to capture the EAGLE3 model here. if you look at the profiling traces more closely, you probably should see there was no cuda graph launch for eagle3 previously.

Noob question which line were you referring to? This is trace from current PR

[luyuzhe111]

@zixi-qi i meant the dummy_run() method of eagle proposer, which is called here from the cuda graph capture method.

[WoosukKwon]

@luyuzhe111 @zixi-qi I've updated the PR. Thanks for catching the bug! Will merge once it passes the CI.

[zixi-qi]

@luyuzhe111 @WoosukKwon

Thanks so much for catching and fixing the issue! Also verified locally that the acceptance rate issue is fixed

--------------------------------------------------
mean acceptance length:         3.57
--------------------------------------------------
acceptance at token 0:1.00
acceptance at token 1:0.75
acceptance at token 2:0.55
acceptance at token 3:0.42
acceptance at token 4:0.31
acceptance at token 5:0.24
acceptance at token 6:0.18
acceptance at token 7:0.13