gahow/xserv
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

eagle3: γ≥2 recursive drafting + batched verify with hooks

Browse Source 
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.
This commit is contained in:
Gahow Wang
2026-07-01 18:01:55 +08:00
parent a24621fa6a
commit 14925154a3
3 changed files with 328 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

59
crates/xserv-model/src/eagle3.rs
View File

@@ -165,6 +165,44 @@ impl Eagle3Head {
prev_token: u32,
position: usize,
) -> (u32, Tensor) {
let (id, logits, _) = self.step_with_aux(target_hidden, embed_table, prev_token, position);
(id, logits)
}
/// Like `step`, but also returns the final hidden state (aux) usable as
/// the fused_h for a subsequent recursive draft step via `step_recursive`.
pub fn step_with_aux(
&mut self,
target_hidden: &[Tensor; 3],
embed_table: &Tensor,
prev_token: u32,
position: usize,
) -> (u32, Tensor, Tensor) {
// Fuse 3 target hidden states into fused_h via fc.
let h_cat = concat_hidden(target_hidden);
let fused_h = matmul_2d(&h_cat, &self.fc_wt);
self.forward_from_fused(fused_h, embed_table, prev_token, position)
}
/// Recursive draft step: reuses the previous EAGLE step's aux as fused_h,
/// bypassing the fc+3-hidden fusion. Used for γ≥2 chained drafts.
pub fn step_recursive(
&mut self,
fused_h: Tensor,
embed_table: &Tensor,
prev_token: u32,
position: usize,
) -> (u32, Tensor, Tensor) {
self.forward_from_fused(fused_h, embed_table, prev_token, position)
}
fn forward_from_fused(
&mut self,
fused_h: Tensor,
embed_table: &Tensor,
prev_token: u32,
position: usize,
) -> (u32, Tensor, Tensor) {
let eps = 1e-6f32;
assert!(
self.current_len < self.max_seq_len,
@@ -173,20 +211,12 @@ impl Eagle3Head {
self.max_seq_len
);
// 1. Fuse target hidden states: concat [h_low, h_mid, h_high] → fc
let h_cat = concat_hidden(target_hidden);
let fused_h = matmul_2d(&h_cat, &self.fc_wt); // [1, hidden]
// 2. Embed previous token (shared with target)
let emb = embedding(embed_table, &[prev_token]); // [1, hidden]
// 3. Norm both, concat, remember residual = fused_h (pre-norm).
let emb = embedding(embed_table, &[prev_token]);
let residual = fused_h.clone();
let emb_normed = rmsnorm(&emb, &self.input_layernorm, eps);
let h_normed = rmsnorm(&fused_h, &self.hidden_norm, eps);
let attn_in = concat_last_dim(&emb_normed, &h_normed); // [1, 2*hidden]
let attn_in = concat_last_dim(&emb_normed, &h_normed);
// 4. Q/K/V projection then RoPE (position from caller).
let q = matmul_2d(&attn_in, &self.q_proj_wt);
let k = matmul_2d(&attn_in, &self.k_proj_wt);
let v = matmul_2d(&attn_in, &self.v_proj_wt);
@@ -197,8 +227,6 @@ impl Eagle3Head {
rope_inplace(&q_3d, &self.rope_cache, &positions);
rope_inplace(&k_3d, &self.rope_cache, &positions);
// 5. Append new K/V to the internal cache at slot `current_len`, then
// build a contiguous view [1, num_kv_heads, current_len+1, head_dim].
let v_3d = v.reshape(&[1, self.num_kv_heads, self.head_dim]);
self.append_to_kv_cache(&k_3d, &v_3d);
self.current_len += 1;
@@ -206,31 +234,26 @@ impl Eagle3Head {
let k_view = self.k_cache.narrow(2, 0, kv_len).contiguous();
let v_view = self.v_cache.narrow(2, 0, kv_len).contiguous();
// 6. Attention: q [1, num_q_heads, 1, head_dim] × k/v [1, num_kv_heads, kv_len, head_dim]
let q_4d = q_3d.reshape(&[1, self.num_heads, 1, self.head_dim]);
let attn_out = decode_attention(&q_4d, &k_view, &v_view);
// 7. Merge heads and o_proj.
let attn_merged = attn_out.reshape(&[1, self.num_heads * self.head_dim]);
let attn_proj = matmul_2d(&attn_merged, &self.o_proj_wt);
// 8. Post-attn fused add_rmsnorm.
let (mlp_in, residual) =
add_rmsnorm(&attn_proj, &residual, &self.post_attention_layernorm, eps);
// 9. MLP.
let gate = matmul_2d(&mlp_in, &self.gate_proj_wt);
let up = matmul_2d(&mlp_in, &self.up_proj_wt);
let hidden = silu_mul(&gate, &up);
let down = matmul_2d(&hidden, &self.down_proj_wt);
// 10. Final fused add_rmsnorm → lm_head.
let (x, _) = 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)
(target_id, logits, x)
}
/// Write new K/V rows (shape [1, num_kv_heads, head_dim]) at position