eagle3: proper residual chain + stateful KV cache
Two fixes to bring EAGLE3 forward in line with vllm's llama_eagle3.py
reference:
1. Residual chain: previously the residual added into post_attention_layernorm
was the token embedding (wrong). Reference uses _norm_after_residual:
residual = fused_h (pre-norm)
hidden_states = hidden_norm(fused_h)
Then post_attention_layernorm is a fused add_rmsnorm(attn_out, residual),
and the final norm is another add_rmsnorm(mlp_out, residual_after_attn).
Neither residual carries the embedding — both carry fused_h forward.
2. KV cache: previously the attention was approximated as "output = V"
because seq_len=1 (no cache), effectively giving EAGLE no history.
Add a real per-Eagle3Head KV cache (1 layer × [1, num_kv_heads,
max_seq_len, head_dim] BF16) that grows as we call step(). Use the
existing decode_attention kernel with a fresh contiguous slice of the
cache each step. reset() clears current_len for a new sequence.
Result on 10 prompts × 32 tokens (γ=1, no batched verify yet):
matched=true across all prompts
acceptance_rate = 20.0% (was 4.7% before residual fix, 1.3% originally)
- Prompt 00 "The capital of France is": 60% (18/30) — best case
- Other prompts: 10-25% — matches EAGLE paper's observation that
structured/factual prompts get higher acceptance
Sanity check (check-eagle3) on Paris prompt now shows:
EAGLE top-5 pairing A: "." / " is" / "," / " Paris" / ".\n"
MATCH: EAGLE agrees with target on next token.
speedup_e2e still 0.95x because γ=1 does 1 target decode per token
regardless of acceptance. Real speedup requires γ≥2 with a single
batched target-verify covering all γ draft tokens; that's the next step.
This commit is contained in:
@@ -38,7 +38,7 @@ fn main() {
|
||||
xserv_cuda::allocator::cached_trim();
|
||||
|
||||
eprintln!("Loading EAGLE3 head from {}", eagle_dir.display());
|
||||
let eagle = Eagle3Head::load(&eagle_dir, device);
|
||||
let mut eagle = Eagle3Head::load(&eagle_dir, device);
|
||||
xserv_cuda::allocator::cached_trim();
|
||||
|
||||
let tokenizer = Tokenizer::from_file(&target_dir.join("tokenizer.json"));
|
||||
@@ -103,6 +103,7 @@ fn main() {
|
||||
// Ask EAGLE what it thinks the NEXT token is (given target_first as prev_token
|
||||
// and the hidden states from the position where target_first lives).
|
||||
// EAGLE should predict target_next (or close to it) to be useful.
|
||||
eagle.reset();
|
||||
let (eagle_pred, eagle_logits) = eagle.step(&hooks, embed_tokens, target_first, pos);
|
||||
let eagle_pred_text = tokenizer.decode(&[eagle_pred]);
|
||||
println!(
|
||||
@@ -129,6 +130,7 @@ fn main() {
|
||||
|
||||
// Alternative pairing B: pair hooks with target_next (the token those hooks produced
|
||||
// via lm_head), predict token after target_next. Position advances by 1.
|
||||
eagle.reset();
|
||||
let (eagle_pred_b, eagle_logits_b) = eagle.step(&hooks, embed_tokens, target_next, pos + 1);
|
||||
let eagle_pred_b_text = tokenizer.decode(&[eagle_pred_b]);
|
||||
println!(
|
||||
|
||||
Reference in New Issue
Block a user