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:
2026-07-01 17:50:49 +08:00
parent 68b55fa1e6
commit a24621fa6a
3 changed files with 90 additions and 34 deletions

View File

@@ -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!(