speculative: EAGLE3 draft head implementation (Phase 25 step 1)
- 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).
This commit is contained in:
@@ -825,6 +825,111 @@ impl Qwen3 {
|
||||
matmul_2d(&x, &self.lm_head_t)
|
||||
}
|
||||
|
||||
/// Like `decode_core` but also captures hidden states at 3 specified layer
|
||||
/// indices (after residual+MLP output). Used by EAGLE3 speculative drafting
|
||||
/// to feed the draft head with low/mid/high target representations.
|
||||
pub fn decode_core_with_hidden(
|
||||
&self,
|
||||
ids_gpu: *const std::ffi::c_void,
|
||||
pos_gpu: *const std::ffi::c_void,
|
||||
batch: usize,
|
||||
seq_slots: &[usize],
|
||||
paged_cache: &mut PagedKVCache,
|
||||
hook_layers: &[usize; 3],
|
||||
) -> (Tensor, [Tensor; 3]) {
|
||||
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;
|
||||
|
||||
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_device_ids(&self.embed_tokens, ids_gpu, batch);
|
||||
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_2d(&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(&[batch * num_heads, head_dim]);
|
||||
let k_flat = k_all
|
||||
.contiguous()
|
||||
.reshape(&[batch * 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(&[batch, num_heads, head_dim]);
|
||||
let k_3d = k_normed.reshape(&[batch, num_kv_heads, head_dim]);
|
||||
rope_inplace_device_pos(&q_3d, &self.rope_cache, pos_gpu);
|
||||
rope_inplace_device_pos(&k_3d, &self.rope_cache, pos_gpu);
|
||||
|
||||
let v_3d = v_all.contiguous().reshape(&[batch, num_kv_heads, head_dim]);
|
||||
|
||||
paged_cache.append_tokens_batched(layer_idx, &k_3d, &v_3d, batch);
|
||||
|
||||
let q_4d = q_3d.reshape(&[batch, 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_4d,
|
||||
k_pool_ptr,
|
||||
v_pool_ptr,
|
||||
bt_ptr,
|
||||
cl_ptr,
|
||||
batch,
|
||||
num_heads,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
max_blocks,
|
||||
);
|
||||
|
||||
let attn_merged = attn_out.reshape(&[batch, num_heads * head_dim]);
|
||||
let attn_proj = matmul_2d(&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_2d(&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_2d(&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());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for &slot in seq_slots {
|
||||
paged_cache.advance_seq_len(slot, 1);
|
||||
}
|
||||
|
||||
let x = rmsnorm(&x, &self.norm, eps);
|
||||
let logits = matmul_2d(&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)
|
||||
}
|
||||
|
||||
/// Paged prefill: write a sequence of `new_tokens` K/V into the paged
|
||||
/// cache for `slot`, run flash attention via gathered contiguous K/V.
|
||||
/// Returns logits [new_tokens, vocab_size].
|
||||
@@ -1074,6 +1179,12 @@ impl Qwen3 {
|
||||
matmul_2d(&x, &self.lm_head_t)
|
||||
}
|
||||
|
||||
/// Reference to the target's token embedding table. Shared (not copied)
|
||||
/// with speculative draft heads like EAGLE3.
|
||||
pub fn embed_tokens_tensor(&self) -> &Tensor {
|
||||
&self.embed_tokens
|
||||
}
|
||||
|
||||
/// Extract weight pointers for CUDA Graph capture.
|
||||
pub fn layer_weight_ptrs(&self) -> Vec<crate::decode_graph::LayerWeightPtrs> {
|
||||
self.layers
|
||||
|
||||
Reference in New Issue
Block a user