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

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:
2026-07-01 18:01:55 +08:00
parent a24621fa6a
commit 14925154a3
3 changed files with 328 additions and 30 deletions

View File

@@ -1113,6 +1113,123 @@ impl Qwen3 {
matmul_batched_gemv(&x, &self.lm_head_t)
}
/// Like `forward_verify_paged_decode_attention`, but also captures hidden
/// states at 3 layer indices (per position). Returns
/// (logits [new_tokens, vocab], hooks [3][new_tokens, hidden]). Used by
/// EAGLE3 speculative γ≥2 verify path so we can seed the next round's
/// EAGLE draft with target's real hidden states at the accepted position.
pub fn forward_verify_paged_decode_attention_with_hidden(
&self,
token_ids: &[u32],
slot: usize,
paged_cache: &mut PagedKVCache,
hook_layers: &[usize; 3],
) -> (Tensor, [Tensor; 3]) {
let new_tokens = token_ids.len();
let pos_offset = paged_cache.seq_len(slot);
let num_heads = self.local_num_heads;
let num_kv_heads = self.local_num_kv_heads;
let head_dim = self.config.head_dim();
let eps = self.config.rms_norm_eps.unwrap_or(1e-6) as f32;
paged_cache.ensure_capacity(slot, pos_offset + new_tokens);
paged_cache.advance_seq_len(slot, new_tokens);
let positions: Vec<u32> = (pos_offset..pos_offset + new_tokens)
.map(|p| p as u32)
.collect();
let kv_lens: Vec<i32> = (0..new_tokens)
.map(|i| (pos_offset + i + 1) as i32)
.collect();
let slots = vec![slot; new_tokens];
paged_cache.sync_active_batch_with_lens(&slots, &kv_lens);
let bt_ptr = paged_cache.block_table_gpu().as_ptr() as *const i32;
let cl_ptr = paged_cache.context_lens_gpu().as_ptr() as *const i32;
let max_blocks = paged_cache.max_blocks_per_seq();
let mut x = embedding(&self.embed_tokens, token_ids);
let mut hooks: [Option<Tensor>; 3] = [None, None, None];
for (layer_idx, layer) in self.layers.iter().enumerate() {
let residual = x.clone();
let normed = rmsnorm(&x, &layer.input_norm, eps);
let qkv = matmul_batched_gemv(&normed, &layer.qkv_proj_wt);
let q_dim = num_heads * head_dim;
let kv_dim = num_kv_heads * head_dim;
let q_all = qkv.narrow(1, 0, q_dim);
let k_all = qkv.narrow(1, q_dim, kv_dim);
let v_all = qkv.narrow(1, q_dim + kv_dim, kv_dim);
let q_flat = q_all
.contiguous()
.reshape(&[new_tokens * num_heads, head_dim]);
let k_flat = k_all
.contiguous()
.reshape(&[new_tokens * num_kv_heads, head_dim]);
let q_normed = rmsnorm(&q_flat, &layer.q_norm, eps);
let k_normed = rmsnorm(&k_flat, &layer.k_norm, eps);
let q_3d = q_normed.reshape(&[new_tokens, num_heads, head_dim]);
let k_3d = k_normed.reshape(&[new_tokens, num_kv_heads, head_dim]);
rope_inplace(&q_3d, &self.rope_cache, &positions);
rope_inplace(&k_3d, &self.rope_cache, &positions);
let v_3d = v_all
.contiguous()
.reshape(&[new_tokens, num_kv_heads, head_dim]);
paged_cache.append_tokens_batched(layer_idx, &k_3d, &v_3d, new_tokens);
let q_decode = q_3d.reshape(&[new_tokens, num_heads, 1, head_dim]);
let k_pool_ptr = paged_cache.k_pool(layer_idx).as_ptr() as *const std::ffi::c_void;
let v_pool_ptr = paged_cache.v_pool(layer_idx).as_ptr() as *const std::ffi::c_void;
let attn_out = xserv_kernels::paged_decode_attention(
&q_decode,
k_pool_ptr,
v_pool_ptr,
bt_ptr,
cl_ptr,
new_tokens,
num_heads,
num_kv_heads,
head_dim,
max_blocks,
);
let attn_merged = attn_out.reshape(&[new_tokens, num_heads * head_dim]);
let attn_proj = matmul_batched_gemv(&attn_merged, &layer.o_proj_wt);
self.all_reduce(&attn_proj);
let (normed, x_new) =
xserv_kernels::add_rmsnorm(&attn_proj, &residual, &layer.post_norm, eps);
let residual = x_new.clone();
let gate_up = matmul_batched_gemv(&normed, &layer.gate_up_proj_wt);
let ffn_dim = gate_up.shape()[1] / 2;
let gate = gate_up.narrow(1, 0, ffn_dim).contiguous();
let up = gate_up.narrow(1, ffn_dim, ffn_dim).contiguous();
let hidden_states = xserv_kernels::silu_mul(&gate, &up);
let down = matmul_batched_gemv(&hidden_states, &layer.down_proj_wt);
self.all_reduce(&down);
x = add_any(&residual, &down);
for (h_idx, &h_layer) in hook_layers.iter().enumerate() {
if layer_idx == h_layer {
hooks[h_idx] = Some(x.clone());
}
}
}
let x = rmsnorm(&x, &self.norm, eps);
let logits = matmul_batched_gemv(&x, &self.lm_head_t);
let hidden_arr = [
hooks[0].take().expect("hook layer 0 not reached"),
hooks[1].take().expect("hook layer 1 not reached"),
hooks[2].take().expect("hook layer 2 not reached"),
];
(logits, hidden_arr)
}
/// Forward with GPU-resident KV cache and GPU transpose/reshape kernels.
pub fn forward_gpu_cache(&self, token_ids: &[u32], cache: &mut GpuKVCache) -> Tensor {
let new_tokens = token_ids.len();