[Bug]: EAGLE / MTP Doesn't Overwrite Approximated Hidden States / KV Cache, 8%- 15% Acceptance Length Degradation #14649
Labels
bug
Something isn't working
Comments
This was referenced Mar 12, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Your current environment
The output of `python collect_env.py`
🐛 Describe the bug
This bug, incurring significant degradations to EAGLE's acceptance rate, is actually quite subtle, so I would like to first provide some contexts.
As described in its paper, EAGLE tries to predict the hidden states of future tokens, and uses them autoregressively during drafting. The important point is that these hidden states are approximated, rather than ground truths from the target model. As a result, even though some draft tokens are accepted, their corresponding KV cache is to some extent, not accurate. Let$h$ be the hidden state from the target model, $\hat{h}$ be the hidden state predicted by the EAGLE model during drafting, and $e$ be the embedding of the accepted draft token. Now, the KV cache should be from some transformation of the concatenation of ($h$ , $e$ ), not ($\hat{h}$ , $e$ ). This is not immediately clear solely from the EAGLE paper, but is quite evident if we look at the codebase. In EAGLE’s codebase, the authors keep track of stable_kv which are KV cache from hidden states of the target model, which are retrieved here. In contrast, the KV cache during drafting, from predicted hidden states (past_key_values), were never saved to the next forward pass. However, in vLLM, we never overwrite EAGLE’s approximated KV cache if the drafted token is accepted. Since we always use the target model’s ground truth hidden states during training, having these approximated KV cache creates problems when we speculated more tokens for longer horizons. Concretely, let’s say at a certain step we speculate 5 tokens, and 3 of them accepted. Because we have these approximated KV cache that’s not quite right, we will have a harder time drafting accurately later since we have 3 approximated KV cache in memory. In the following table, I show the acceptance rate drop between EAGLE repo and vLLM’s EAGLE as we increase the number of speculated tokens. I computed the average accepted length using features added in this PR. We can see that generally, with more speculated tokens, we observe larger acceptance length discrepancies. In vLLM, with 4 and 5 speculated tokens, we even see lower acceptance length. The reason is likely due to the accumulation of approximated KV cache in EAGLE, which ends up hurting the acceptance rate for all positions.
This issue also bothers DeepSeek MTP when number of MTP modules > 1 PR in a similar manner.
cc @LiuXiaoxuanPKU @benchislett
Before submitting a new issue...
The text was updated successfully, but these errors were encountered: