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.
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).
Adds infrastructure for γ≥2 EAGLE speculative decoding:
qwen3.rs:
- New forward_verify_paged_decode_attention_with_hidden: same as the
existing verify but also captures target hidden states at 3 hook
layers, one per verify position. Needed to seed next round's EAGLE.
eagle3.rs:
- step split into step (unchanged public API) + step_with_aux (also
returns final hidden state) + step_recursive (takes fused_h directly,
no fc+3-hidden combine). This mirrors the EAGLE3 paper: γ=1 uses
target hooks + fc; γ≥2 uses previous EAGLE aux as fused_h for
subsequent drafts, approximating target hidden.
bench-eagle3.rs:
- New run_eagle_gamma_multi function with --gamma CLI (default 2).
- Per round: recursive EAGLE γ drafts, verify [prev_token, d0..d_{γ-1}]
in one target forward, accept longest prefix, correction via 1 more
target decode.
- max_seqs bumped to 16 in the paged cache so verify can batch up to
16 rows.
γ=2 test result (5 prompts × 32 tokens, dash5):
matched=false — sequences diverge
acceptance_rate = 29.8% at γ=2 (~1.1 tokens accepted per draft)
speedup_e2e = 0.52x (SLOWER than baseline)
The divergence bug is in the verify's re-writing of prev_token's K/V
at position round_pos-1. In principle matmul_batched_gemv at row-0
should be bit-exact with the seed decode's launch_gemv_bf16, but the
sequence output diverges so something is off. Investigation pending
(likely the correction decode step or seed_hooks position offset).
γ=1 path still works correctly (matched=true, acceptance 20%,
speedup 0.95x) from the previous commit. The γ≥2 path is scaffolded
but not yet correct — next step is to debug the verify-write path,
then measure real speedup.
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.
The initial [11, 23, 35] (equally-spaced) guess was wrong — EAGLE3 heads
are trained against specific target layer indices, and using different
ones at inference gives wrong outputs. Correct values come from vLLM
speculators' training config for Qwen3-8B:
https://github.com/vllm-project/speculators/blob/main/examples/train/
dflash_qwen3_8b_sharegpt_online_5k.sh
which pins target_layer_ids to "2 18 33". Re-running check-eagle3 with
the fix produces coherent top-5 for "The capital of France is":
Old ([11,23,35]): "," / " Paris" / " Madrid" / "." / " Berlin"
New ([2,18,33]): " Paris" / " Tokyo" / " Madrid" / "," / "."
Top-1 still differs from target's next token, but that's because EAGLE
compares (state_that_produced_prev, prev_token) → next, and the exact
pairing convention may need one more offset check when integrated into
the full speculative loop.
- eagle3.rs: Eagle3Head struct loads AngelSlim/Qwen3-8B_eagle3 safetensors,
runs a single draft step via fc(concat(h_low, h_mid, h_high)) +
concat(input_norm(emb), hidden_norm(fused_h)) → 1 midlayer → norm →
lm_head → argmax in draft_vocab(32000) → d2t → target_vocab.
- qwen3.rs: new decode_core_with_hidden method that mirrors decode_core
but captures hidden states at 3 configurable layer indices (default
[11, 23, 35] for the 36-layer Qwen3-8B). Also expose embed_tokens_tensor
and (in eagle3) map_draft_to_target as public accessors.
- loader.rs: make_tensor now pub(crate) so eagle3 can reuse it.
- bin/check-eagle3.rs: sanity binary that loads target + EAGLE, runs one
prefill + one decode + one EAGLE step, prints the top-5 EAGLE predictions.
Verified on dash5 with prompt "The capital of France is":
target says: " Paris" then "."
EAGLE top-5: "," / " Paris" / " Madrid" / "." / " Berlin"
Weights load correctly, d2t mapping works, hidden state hooks are the
right shape ([1, 4096]), and EAGLE produces thematically-relevant tokens.
The top-1 pick "," doesn't match target's "." at this position, but
that's expected: this test uses hidden states from a single decode step
with no recursive chaining. A full speculative loop still needs the
γ≥2 verify + accept path wired up (next step).
