eagle3: γ≥2 correctness fixes + per-slot diagnostic

Two subtle bugs found and fixed in the γ≥2 speculative loop:

1. Wrong position handling: cache.truncate_sequence(round_pos - 1) was
   dropping the K/V of pending_prev, then verify OVERWROTE that slot with
   the wrong token. Removed the truncate: verify now starts at
   cache.seq_len (== position of pending_prev) and writes γ+1 tokens
   forward. Also fixed EAGLE draft positions: pending_prev is at position
   p, so step 0 uses position=p (not p+1).

2. EAGLE KV cache accumulated rejected drafts' K/V: each round writes γ
   entries to EAGLE's cache regardless of how many drafts were accepted.
   Added eagle.truncate_to(new_len) API. After each round, truncate to
   eagle_len_before + k + 1 (pending_prev + k accepted drafts).

Also expose Eagle3Head::current_len() getter and Eagle3Head::truncate_to().

Additionally: return the PRE-norm hidden state as aux (matching vllm's
llama_eagle3.py default `norm_output=False`). Was returning the normed
version.

Result: matched=true across the full γ sweep. speedup_e2e remains <1:

  γ=1 (single-decode verify): accept=22.7%, speedup=0.95x
  γ=1 (batched verify):       accept=20.6%, speedup=0.75x
  γ=2:                         accept=12.6%, speedup=0.59x
  γ=4:                         accept=7.6%,  speedup=0.41x
  γ=8:                         accept=4.1%,  speedup=0.27x

Per-slot diagnostic shows d[0]≈15%, d[1]≈8%, d[2..γ-1] varies. d[0] is
lower than γ=1's 20% because batched verify introduces small numerical
differences vs single-token decode.

Larger γ hurts because:
- verify_cost scales roughly linearly with γ+1 (batched matmul at
  batch=γ+1 costs ~γ+1× a single decode).
- accepted tokens per round grows sub-linearly (recursive EAGLE degrades).
- speedup ≈ (1 + accepted_avg) / verify_cost → below 1 across the sweep.

Path forward for speedup > 1 requires EITHER: (a) faster batched verify
(closer to single-decode cost per query row via better GPU utilization),
OR (b) better draft accuracy (tree-based drafting to explore multiple
candidates per position, larger EAGLE head, or a differently-trained
EAGLE variant).
This commit is contained in:
2026-07-01 19:16:31 +08:00
parent 14925154a3
commit d2c55c47b2
2 changed files with 168 additions and 59 deletions

View File

@@ -150,6 +150,18 @@ impl Eagle3Head {
self.current_len = 0;
}
/// Truncate the internal KV cache to `new_len` entries. Used to discard
/// K/V of rejected drafts after a speculative round.
pub fn truncate_to(&mut self, new_len: usize) {
assert!(new_len <= self.current_len);
self.current_len = new_len;
}
/// Current number of committed K/V entries in the internal EAGLE cache.
pub fn current_len(&self) -> usize {
self.current_len
}
/// One draft step: produce a token in target vocabulary space.
///
/// - `target_hidden`: 3 tensors [1, hidden_size] from target hook layers
@@ -248,12 +260,14 @@ impl Eagle3Head {
let hidden = silu_mul(&gate, &up);
let down = matmul_2d(&hidden, &self.down_proj_wt);
let (x, _) = add_rmsnorm(&down, &residual, &self.norm, eps);
let (x, prenorm) = add_rmsnorm(&down, &residual, &self.norm, eps);
let logits = matmul_2d(&x, &self.lm_head_wt);
let draft_id = argmax_bf16_single(&logits);
let target_id = (draft_id as i64 + self.d2t[draft_id as usize]) as u32;
(target_id, logits, x)
// aux for recursive drafting = PRE-norm hidden (default norm_output=False
// in vllm/llama_eagle3.py). Feeding the pre-norm state matches training.
(target_id, logits, prenorm)
}
/// Write new K/V rows (shape [1, num_kv_heads, head_dim]) at position