[V1][Spec Decode] Apply torch.compile & cudagraph to EAGLE #17211

luyuzhe111 · 2025-04-26T00:02:14Z

Task 8 of #15901

A few notes regarding the implementation.

Torch.compile

@support_torch_compile is convenient but requires a specific signature for the model. Changes in vllm/model_executor/models/llama_eagle.py address this requirement.
Further, to make torch.compile work, we need a separate cache dir for EAGLE model. Without edits in vllm/compilation/backends.py, we wouldn't be able to cache EAGLE's compilation properly since the EAGLE module was registered under vllm config from the target model (see here)
One notable bug related to torch.compile is this line. Essentially, the default data type for input ids is int32 and the EAGLE model was also compiled with this data type. However, tensor.argmax() returns int64 by default. Feeding int64 input ids to the compiled model will completely mess things up and lead to gibberish draft tokens. Currently the compiled model does not even give warnings when the input data type mismatches. Wonder if we can prevent similar bugs in the future by some more checks.

CudaGraph

Changes in vllm/v1/spec_decode/eagle.py and vllm/v1/worker/gpu_model_runner.py are mostly for CudaGraph. Nothing fancy other than registering additional persistent buffers and making sure to use them for EAGLE's forward pass. I do want to mention that with torch.compile & CudaGraph, the EAGLE model's forward pass has been drastically improved (2.5x faster), which makes the small but abundant torch operations look inefficient. Any advice to further optimize these overheads is greatly appreciated.

Finally, it would be great to have #17010 reviewed and merged so that we don't have to pull in other PRs to test the acceptance length.

Note that the current PR does not directly make torch.compile & cuda graph available for EAGLE3. I think it's worth a separate PR since the work is non-trivial due to the fact that EAGLE-3's input hidden states have dynamic shapes. Maybe @benchislett could help.

@WoosukKwon @LiuXiaoxuanPKU

Signed-off-by: Bryan Lu <yuzhelu@amazon.com>

github-actions · 2025-04-26T00:02:24Z

👋 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.

🚀

mergify · 2025-04-26T00:02:53Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @luyuzhe111.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

Signed-off-by: Bryan Lu <yuzhelu@amazon.com>

WoosukKwon · 2025-04-26T17:39:53Z

Thanks for the PR. This is so cool!
I’ll take a look, but it would be great if we could also get @youkaichao’s review.

WoosukKwon

@luyuzhe111 Thanks for the PR!

Left some comments. Please take a look!

vllm/model_executor/models/llama_eagle.py

vllm/v1/spec_decode/eagle.py

vllm/model_executor/models/llama_eagle.py

vllm/v1/spec_decode/eagle.py

vllm/v1/worker/gpu_model_runner.py

Signed-off-by: Bryan Lu <yuzhelu@amazon.com>

vllm/v1/worker/gpu_model_runner.py

Signed-off-by: Bryan Lu <yuzhelu@amazon.com>

WoosukKwon

@luyuzhe111 Thanks for addressing the comments! One last thing: Can you please add tests for this?

luyuzhe111 · 2025-04-29T05:05:40Z

Hi @ekagra-ranjan, thanks for this PR! but wondering if you could create a new PR later that makes sure the draft model takes in vllm_config instead of model_config, if you are interested? This is a requirement for compatibility with torch.compile decorator. for this PR I will just create a condition to handle eagle-1 and eagle-3 models separately. thanks! cc @WoosukKwon

luyuzhe111 · 2025-04-29T08:09:10Z

One last thing: Can you please add tests for this?

@WoosukKwon I feel like acceptance length tests are probably the most meaningful tests for this PR. I tested the acceptance length by cherry picking commits from this PR. You can see that the difference is less than 0.01.

With meta-llama/Llama-3.1-8B-Instruct, yuhuili/EAGLE-LLaMA3.1-Instruct-8B

On MT Bench

When max number generated tokens = 256

Number of Speculated Tokens	1	2	3	4	5
Eager	1.71	2.09	2.30	2.38	2.43
Compilation & CudaGraph	1.70	2.10	2.29	2.38	2.43

I can do a follow-up PR adding acceptance length tests after #17010 is merged.

vllm/vllm_flash_attn/__init__.py

ekagra-ranjan · 2025-04-29T16:03:40Z

Thank you @luyuzhe111 for the PR!
Wondering if you have any benchmark for speedup in output token/s for with and w/o cuda graph and torch compile for EAGLE llama 3.1 on MTBench TP1 BS1?

I do want to mention that with torch.compile & CudaGraph, the EAGLE model's forward pass has been drastically improved (2.5x faster), which makes the small but abundant torch operations look inefficient.

This is fantastic!
Could you share the script you used to measure just the EAGLE's fwd pass?
By "abundant torch operation", do you mean the torch ops outside the forward pass but within the propose() ?

Signed-off-by: Bryan Lu <yuzhelu@amazon.com>

luyuzhe111 · 2025-04-29T18:25:39Z

Hey @ekagra-ranjan Re:

I remember you had some numbers for benchmarking on mt bench. can you share the setup so that I can run the benchmarking again for torch.compile + cuda graph? thanks!
for the forward pass I had to look at the profiler, so no script here.
right I meant those operations that prepare inputs for the EAGLE model.

WoosukKwon

LGTM. @youkaichao Any final comments?

ekagra-ranjan · 2025-04-29T19:46:28Z

@luyuzhe111 - here is an example setup for benchmark. Looking fwd to the results!

luyuzhe111 · 2025-04-29T22:25:54Z

@ekagra-ranjan I got the following results:

Target model: meta-llama/Llama-3.1-8B-Instruct
EAGLE model: yuhuili/EAGLE-LLaMA3.1-Instruct-8B
Hardware: A100 (40GB)
Script: VLLM_USE_V1=1 python examples/offline_inference/eagle.py —dataset="./data/mt_bench/question.jsonl" —num_spec_tokens x —max_num_seqs 1 —num_prompts 80

Regular Decoding OTPS: 72

it looks like a further 10% speedup.

zou3519 · 2025-05-01T19:38:20Z

One notable bug related to torch.compile is this line. Essentially, the default data type for input ids is int32 and the EAGLE model was also compiled with this data type. However, tensor.argmax() returns int64 by default. Feeding int64 input ids to the compiled model will completely mess things up and lead to gibberish draft tokens. Currently the compiled model does not even give warnings when the input data type mismatches. Wonder if we can prevent similar bugs in the future by some more checks.

These checks lead to performance degradation so vLLM decided to drop all of them. Some of the lower compilation_levels (e.g. 1) should do the checking but those also have slightly different behavior for other things

zou3519 · 2025-05-01T19:42:32Z

vllm/compilation/backends.py

+        if compilation_counter.num_graphs_seen > 0:
+            cache_dir = self.compilation_config.cache_dir + \
+                f'-{compilation_counter.num_graphs_seen}'
+        else:
+            cache_dir = self.compilation_config.cache_dir
        os.makedirs(cache_dir, exist_ok=True)
+        self.compilation_config.cache_dir = cache_dir


@luyuzhe111 This is suspicious. Why do you need a different cache directory for each graph? Also, this looks like it modifies everything, even the models that don't use eagle.

If there isn't a good reason I would prefer going back to the "single cache directory" that we had previously.

@zou3519 thanks for reviewing! if there isn't a separate cache directory, the compiled code for the draft model (EAGLE) will not be saved at all. for models without EAGLE, my understanding is that the backend is invoked only once so this should not impact other models.

@luyuzhe111 thanks for the response and clarifying that. Woosuk also filled me in on some more details offline. I understand why we need a separate cache directory.

Which of the "original model" and the "eagle head" get compiled first? (I'm trying to figure out if the first cache dir is for the original model or for the eagle head)

@zou3519 original model should be compiled first! also if you wanna double check, the transformed code of EAGLE in the cache directory has a slightly different signature with hidden_states as an additional arg. if there is a more elegant solution, that would be great! I think my approach is a bit hacky indeed : )))

Thanks for the discussion! I added some comments and an assertion into #17662 , please take a look.

I think in the future we'll want a better way to handle multiple compiled regions in a vLLM model, but that will take some re-designing

@luyuzhe111 the asserts in #17662 triggered, which means that this PR does affect non-eagle models

@zou3519 Thanks for the catch! I guess a simple fix would be just to create a separate cache directory only for EAGLE, via looking at the vllm speculative config, for example?

Yeah that would work

…ect#17211) Signed-off-by: Bryan Lu <yuzhelu@amazon.com>

@luyuzhe111

I'm recording down my understanding of how eagle and the compilation cache works after discussing vllm-project#17211 with @luyuzhe111 and @WoosukKwon. In the future we likely will have a situation where we want to torch.compile multiple pieces of code (e.g. decoder and encoder separately) and then we'll need to refactor the system to support it (each compiled region needs its own cache directory with its own hash) But until then the current design seems fine. Signed-off-by: rzou <zou3519@gmail.com>

…ect#17211) Signed-off-by: Bryan Lu <yuzhelu@amazon.com> Signed-off-by: Mu Huai <tianbowen.tbw@antgroup.com>

…ect#17211) Signed-off-by: Bryan Lu <yuzhelu@amazon.com> Signed-off-by: Yuqi Zhang <yuqizhang@google.com>

youkaichao · 2025-06-02T07:57:44Z

vllm/compilation/backends.py

-        cache_dir = self.compilation_config.cache_dir
+        if compilation_counter.num_graphs_seen > 0:
+            cache_dir = self.compilation_config.cache_dir + \
+                f'-{compilation_counter.num_graphs_seen}'


I think we should be able to get the component's prefix to use as the cache directory, it could be more meaningful.

I'll make that change, @luyuzhe111 and I were discussing something similar above

…ect#17211) Signed-off-by: Bryan Lu <yuzhelu@amazon.com> Signed-off-by: minpeter <kali2005611@gmail.com>

apply torch.compile & cudagraph to EAGLE

644e254

Signed-off-by: Bryan Lu <yuzhelu@amazon.com>

luyuzhe111 requested review from WoosukKwon, robertgshaw2-redhat, njhill, ywang96, comaniac and alexm-redhat as code owners April 26, 2025 00:02

mergify bot added the v1 label Apr 26, 2025

mergify bot added the needs-rebase label Apr 26, 2025

luyuzhe111 changed the title ~~Apply torch.compile & cudagraph to EAGLE~~ [V1][Spec Decode] Apply torch.compile & cudagraph to EAGLE Apr 26, 2025

luyuzhe111 added 2 commits April 26, 2025 00:19

remove redundant code

262356d

Signed-off-by: Bryan Lu <yuzhelu@amazon.com>

resolve merge conflicts with upstream/main

92796e1

Signed-off-by: Bryan Lu <yuzhelu@amazon.com>

mergify bot added the documentation Improvements or additions to documentation label Apr 26, 2025

remove outdated comments

b6a5c3d

Signed-off-by: Bryan Lu <yuzhelu@amazon.com>

mergify bot removed the needs-rebase label Apr 26, 2025

WoosukKwon mentioned this pull request Apr 26, 2025

[Benchmark] Add single turn MTBench to Serving Bench #17202

Merged

WoosukKwon reviewed Apr 28, 2025

View reviewed changes

luyuzhe111 added 2 commits April 28, 2025 22:43

avoid pad token indices

7bea2d1

Signed-off-by: Bryan Lu <yuzhelu@amazon.com>

minor

a49c3e4

Signed-off-by: Bryan Lu <yuzhelu@amazon.com>

WoosukKwon reviewed Apr 28, 2025

View reviewed changes

vllm/v1/worker/gpu_model_runner.py Outdated Show resolved Hide resolved

dummy_run() method for eagle proposer

784d1b6

Signed-off-by: Bryan Lu <yuzhelu@amazon.com>

WoosukKwon reviewed Apr 29, 2025

View reviewed changes

luyuzhe111 requested a review from LucasWilkinson as a code owner April 29, 2025 05:39

markmc added the speculative-decoding label Apr 29, 2025

LucasWilkinson requested changes Apr 29, 2025

View reviewed changes

vllm/vllm_flash_attn/__init__.py Outdated Show resolved Hide resolved

resolve conflicts from upstream/main

85b3ade

Signed-off-by: Bryan Lu <yuzhelu@amazon.com>

luyuzhe111 force-pushed the eagle-non-eager branch from 0f3d045 to 85b3ade Compare April 29, 2025 18:15

WoosukKwon added the ready ONLY add when PR is ready to merge/full CI is needed label Apr 29, 2025

WoosukKwon approved these changes Apr 29, 2025

View reviewed changes

LucasWilkinson approved these changes Apr 29, 2025

View reviewed changes

WoosukKwon enabled auto-merge (squash) April 29, 2025 21:08

WoosukKwon merged commit 70788bd into vllm-project:main Apr 29, 2025
69 checks passed

zixi-qi mentioned this pull request Apr 30, 2025

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

Merged

luyuzhe111 mentioned this pull request May 1, 2025

[Bug]: [V1][Spec Dec] EAGLE TP > 1 leads to errors when using --enforce_eager #17513

Open

1 task

zou3519 reviewed May 1, 2025

View reviewed changes

radeksm pushed a commit to radeksm/vllm that referenced this pull request May 2, 2025

[V1][Spec Decode] Apply torch.compile & cudagraph to EAGLE (vllm-proj…

1846d64

…ect#17211) Signed-off-by: Bryan Lu <yuzhelu@amazon.com>

zou3519 mentioned this pull request May 5, 2025

Tighten compilation cache invariants around eagle #17662

Closed

ekagra-ranjan mentioned this pull request May 7, 2025

[Benchmark][V1][Spec Decode][EAGLE] Tracking benchmark for V1 EAGLE #17812

Open

vadiklyutiy mentioned this pull request May 12, 2025

[Bugfix] Enable torch.comple for 2 parts of model #14913

Closed

RichardoMrMu pushed a commit to RichardoMrMu/vllm that referenced this pull request May 12, 2025

[V1][Spec Decode] Apply torch.compile & cudagraph to EAGLE (vllm-proj…

4951fd6

…ect#17211) Signed-off-by: Bryan Lu <yuzhelu@amazon.com> Signed-off-by: Mu Huai <tianbowen.tbw@antgroup.com>

zzzyq pushed a commit to zzzyq/vllm that referenced this pull request May 24, 2025

[V1][Spec Decode] Apply torch.compile & cudagraph to EAGLE (vllm-proj…

1adf95f

…ect#17211) Signed-off-by: Bryan Lu <yuzhelu@amazon.com> Signed-off-by: Yuqi Zhang <yuqizhang@google.com>

youkaichao reviewed Jun 2, 2025

View reviewed changes

youkaichao mentioned this pull request Jun 3, 2025

[torch.compile] reorganize the cache directory to support compiling multiple models #19064

Merged

minpeter pushed a commit to minpeter/vllm that referenced this pull request Jun 24, 2025

[V1][Spec Decode] Apply torch.compile & cudagraph to EAGLE (vllm-proj…

84924e0

…ect#17211) Signed-off-by: Bryan Lu <yuzhelu@amazon.com> Signed-off-by: minpeter <kali2005611@gmail.com>

Uh oh!

[V1][Spec Decode] Apply torch.compile & cudagraph to EAGLE #17211

[V1][Spec Decode] Apply torch.compile & cudagraph to EAGLE #17211

Uh oh!

Conversation

luyuzhe111 commented Apr 26, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Torch.compile

CudaGraph

Uh oh!

github-actions bot commented Apr 26, 2025

Uh oh!

mergify bot commented Apr 26, 2025

Uh oh!

WoosukKwon commented Apr 26, 2025

Uh oh!

WoosukKwon left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

WoosukKwon left a comment

Choose a reason for hiding this comment

Uh oh!

luyuzhe111 commented Apr 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

luyuzhe111 commented Apr 29, 2025

Uh oh!

Uh oh!

ekagra-ranjan commented Apr 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

luyuzhe111 commented Apr 29, 2025

Uh oh!

WoosukKwon left a comment

Choose a reason for hiding this comment

Uh oh!

ekagra-ranjan commented Apr 29, 2025

Uh oh!

Uh oh!

luyuzhe111 commented Apr 29, 2025

Uh oh!

zou3519 commented May 1, 2025

Uh oh!

zou3519 May 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

luyuzhe111 commented Apr 26, 2025 •

edited by github-actions bot

Loading

luyuzhe111 commented Apr 29, 2025 •

edited

Loading

ekagra-ranjan commented Apr 29, 2025 •

edited

Loading

zou3519 May 1, 2025 •

edited

Loading