[Spec Decode] EAGLE 3: Micro-optimization of Draft to Target logit conversion

[ekagra-ranjan]

E3 needs to map the draft logit ids to target logit ids. Currently, the following operations happen:

1. base = torch.arange(self.config.draft_vocab_size, device=logits.device)
2. allocate new activation of dim [bs, target_vocab_size] in every fwd pass
3. fill the target_vocab_size with draft_vocab_size and remaning with -inf
4. argmax the tensor with target_vocab_size elements

This PR instead does:

1. argmax the tensor with draft_vocab_size elements (faster than step 4) since draft_vocab_size < target_vocab_size
2. find the offset in O(1) and adds the offset

It avoid extra allocation and computation for the Draft to Target logit conversion in Eagle 3.

The PR also does some refactoring so that there is more code sharing bw model definitions. E.g., any model implementing E3 does not need to reimplement the Draft to Target logit conversion logic in the model definition.

Test

Compute the AL and ensured it matches the AL w/o this PR.

E1:
python3 examples/offline_inference/spec_decode.py --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill

output

--------------------------------------------------
total_num_output_tokens: 16956
num_drafts: 7405
num_draft_tokens: 22215
num_accepted_tokens: 9557
mean acceptance length: 2.29
--------------------------------------------------

E3
python3 examples/offline_inference/spec_decode.py --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill

output

--------------------------------------------------                                                                                                             
total_num_output_tokens: 16930                                                                                                                                 
num_drafts: 6079                                                                                                                                               
num_draft_tokens: 18237                                                                                                                                        
num_accepted_tokens: 10840                                                                                                                                     
mean acceptance length: 2.78                                                                                                                                   
-------------------------------

Benchmark

BS 64:

• K=5: 17.82 -> 17.6
• K=7: 20.89 -> 20.62

1% improvement

[gemini-code-assist]

Code Review

This pull request introduces a micro-optimization for EAGLE 3's draft-to-target logit conversion, avoiding extra allocation and computation. It also refactors the code to improve sharing by introducing EagleMixin and Eagle3Mixin. While the optimization idea is sound, I've found two critical issues in the implementation. One is a TypeError due to an incorrect number of arguments in a method call. The other is a logical bug in propose_chain that will lead to an index out of bounds error during speculative token generation. Both issues need to be addressed.

[benchislett]

propose_chain feels like a poor name for this function, as it conflicts with EagleProposer.propose(). How about something closer to "get_draft_logits" or similar?

[benchislett]

Or "greedy_sample_draft_tokens" might be more appropriate since I now see that this is a shortcut for compute_logits + greedy sampling

[ekagra-ranjan]

There is a method called propose_tree so propose_chain is analogious to it. The propose_chain and propose_tree are called from within EagleProposer.propose(). The purpose of this is to get the chain draft so I am not sure get_draft_logits is a better name

[benchislett]

In what way is propose_chain here analogous to propose_tree? The latter seems like a much more comprehensive propose()-style function, whereas propose_chain here seems like a much simpler combination of compute_logits and greedy sampling?

[ekagra-ranjan]

In what way is propose_chain here analogous to propose_tree?

propose_tree() is called inside propose() and returns draft token ids at the current step
propose_chain() is called inside propose() and returns draft token ids at the current step

The functionality/role of propose_chain() is analogous to propose_tree(), i.e. its called by propose() and returns the draft tokens ids sampled at the current step.

For E1, its compute logits and argmax whereas in E3 it has an extra step of offset computation. This different behavior of E1 and E3 is abstracted away from propose() function which will call propose_chain().

[benchislett]

I see what you mean. I think I would still prefer something like sample_chain or propose_sample_greedy or similar, but I am okay to go with propose_chain if you think it is the best choice.

[ekagra-ranjan]

Alright, I have renamed it to sample_chain following your suggestion.

[benchislett]

does this optimization lead to any measurable speedup? Could you include benchmarking and/or profiling results as needed to illustrate the advantage?

[ekagra-ranjan]

does this optimization lead to any measurable speedup? Could you include benchmarking and/or profiling results as needed to illustrate the advantage?

It's a micro-optimization on the draft to logit conversion scope so it leads to ~1% for BS64.

[ekagra-ranjan]

@benchislett - updated PR description with some benchmark. The gains are not significant but it does improve the draft to target logit conversion and the effect is seen on higher BS and larger K value. The ITL goes from 17.82 -> 17.6 at BS 64 for K=5. The PR also consolidates the common functions in E3 which can be inherited thereby avoiding the need to be implemented in every model class.

[mergify]

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

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

[benchislett]

I think this is supposed to call compute_logits? Otherwise I don't see where Eagle3Mixin.compute_logits would ever get called here.

Suggested change

	logits = self.compute_draft_logits(hidden_states)
	logits = self.compute_logits(hidden_states)

[ekagra-ranjan]

compute_draft_logits() is the optimized version of compute_logits() and is used in chain sampling. The original compute_logits() is still called by the Tree sampling here in eagle.py

[rahul-tuli]

Hi @ekagra-ranjan are you still working on this?

[ekagra-ranjan]

@rahul-tuli - sorry for the late response. Yes, I will update this PR this weekend.

[mergify]

Hi @ekagra-ranjan, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[mergify]

Hi @ekagra-ranjan, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[mergify]

Hi @ekagra-ranjan, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[ekagra-ranjan]

@benchislett @rahul-tuli - I have resolved the comments and updated the PR. Pls have a look!

[mergify]

Hi @ekagra-ranjan, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[mergify]

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

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

[ekagra-ranjan]

synced and resolved conflict after #34049

[ekagra-ranjan]

this if block comes from recently merged: #34049 whereas the else part is what this PR originally added

[github-actions]

This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this pull request should remain open. Thank you!

[github-actions]

This pull request has been automatically closed due to inactivity. Please feel free to reopen if you intend to continue working on it. Thank you!