eagle3: coverage + top-3 diagnostic; acceptance ceiling analysis

Add t2d bool tensor loading and per-slot top-3 rate tracking to
bench-eagle3 so we can distinguish three failure modes:
- Not covered: target's argmax not in EAGLE's 32k-vocab (upper bound).
- Not top-3: target's argmax not in EAGLE's top-3 (drafting quality).
- Not top-1: target's argmax not EAGLE's argmax (final acceptance rule).

Measured on 50 prompts × 64 tokens γ=2:
  d[0]: correct=27%, top-3=42%, covered=98% → EAGLE covers vocab well
                                              but often ranks target
                                              answer below top-1.
  d[1]: correct=9%,  top-3=17%, covered=100% → recursive draft even
                                               weaker.

Coverage is essentially not a bottleneck (98%+). The bottleneck is
that EAGLE ranks the true target answer only ~27% of the time at slot
0. Top-3 rate (~42%) shows the correct answer is often in EAGLE's
distribution but not the highest-scored candidate.

To exploit the top-3 headroom would require tree-based verify (multiple
candidates per position, tree-aware attention masking). Each candidate
attends only to its own branch, not siblings. Current paged_decode_
attention writes K/V at unique per-batch positions and does not
support tree causal masks.

Speedup formula analysis (from bench-verify-cost):
  γ=2: verify_cost=1.11×, round_yield=1.34 → theoretical speedup=1.21×,
       observed 1.10× (0.11× lost to EAGLE draft cost + bookkeeping).
  γ=4: verify_cost=1.12×, round_yield=1.36 → theoretical=1.21×,
       observed 1.02×.

Current numbers are near-optimal given measured acceptance. Further
gains require either tree drafting (unlocks top-K acceptance) or a
better-trained EAGLE head. Neither is a small change.
This commit is contained in:
2026-07-01 20:19:28 +08:00
parent cc3bc2188c
commit 10a98539d0
2 changed files with 74 additions and 13 deletions

View File

@@ -43,6 +43,10 @@ pub struct Eagle3Head {
norm: Tensor, // [hidden] final
lm_head_wt: Tensor, // [draft_vocab, hidden]
d2t: Vec<i64>, // [draft_vocab] offset mapping
/// t2d[target_id] = true iff target_id has a corresponding draft-vocab id
/// (i.e. can potentially be produced by EAGLE). Used to measure the
/// coverage cap on acceptance.
t2d: Vec<bool>,
hidden_size: usize,
num_heads: usize,
num_kv_heads: usize,
@@ -59,7 +63,7 @@ pub struct Eagle3Head {
impl Eagle3Head {
pub fn load(dir: &Path, device: u32) -> Self {
let (weights, d2t) = load_eagle3_weights(dir, device);
let (weights, d2t, t2d) = load_eagle3_weights(dir, device);
let hidden_size = 4096;
let num_heads = 32;
let num_kv_heads = 8;
@@ -133,6 +137,7 @@ impl Eagle3Head {
norm,
lm_head_wt,
d2t,
t2d,
hidden_size,
num_heads,
num_kv_heads,
@@ -290,6 +295,12 @@ impl Eagle3Head {
pub fn map_draft_to_target(&self, draft_id: u32) -> u32 {
(draft_id as i64 + self.d2t[draft_id as usize]) as u32
}
/// Returns true iff `target_id` is representable in the draft vocabulary
/// (i.e., EAGLE could in principle produce it).
pub fn target_id_in_draft_vocab(&self, target_id: u32) -> bool {
self.t2d.get(target_id as usize).copied().unwrap_or(false)
}
}
fn d2d(dst: *mut u8, src: *const u8, bytes: usize) {
@@ -349,8 +360,8 @@ fn repeat_kv_for_single_token(kv: &Tensor, repeats: usize) -> Tensor {
out
}
/// Load EAGLE3 weights from safetensors, handling int64 d2t specially.
fn load_eagle3_weights(dir: &Path, device: u32) -> (HashMap<String, Tensor>, Vec<i64>) {
/// Load EAGLE3 weights from safetensors, handling int64 d2t + bool t2d specially.
fn load_eagle3_weights(dir: &Path, device: u32) -> (HashMap<String, Tensor>, Vec<i64>, Vec<bool>) {
let st_path = dir.join("model.safetensors");
assert!(
st_path.exists(),
@@ -368,9 +379,13 @@ fn load_eagle3_weights(dir: &Path, device: u32) -> (HashMap<String, Tensor>, Vec
let mut tensors = HashMap::new();
let mut d2t_vec: Vec<i64> = Vec::new();
let mut t2d_vec: Vec<bool> = Vec::new();
for (name, view) in st.tensors() {
if name == "t2d" {
let raw = view.data();
assert_eq!(view.dtype(), safetensors::Dtype::BOOL);
t2d_vec = raw.iter().map(|&b| b != 0).collect();
continue;
}
if name == "d2t" {
@@ -402,5 +417,9 @@ fn load_eagle3_weights(dir: &Path, device: u32) -> (HashMap<String, Tensor>, Vec
!d2t_vec.is_empty(),
"d2t tensor not found in eagle3 weights"
);
(tensors, d2t_vec)
assert!(
!t2d_vec.is_empty(),
"t2d tensor not found in eagle3 weights"
);
(tensors, d2t_vec, t2d_vec)
}