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:
2026-07-01 17:23:22 +08:00
parent 6485c87c5b
commit e04a8ffb18
5 changed files with 577 additions and 1 deletions

View File

@@ -0,0 +1,152 @@
//! EAGLE3 sanity check: load weights, run one draft step, print top-5 predictions.
//!
//! This verifies that:
//! - Eagle3Head weights load without shape mismatches
//! - Target hidden states can be captured via decode_core_with_hidden
//! - Eagle3Head::step produces a valid token id (in target vocab)
//!
//! Does NOT measure speedup — that requires a full γ≥2 speculative loop, which
//! is more complex integration work.
use std::path::PathBuf;
use xserv_model::eagle3::{EAGLE_HOOK_LAYERS, Eagle3Head};
use xserv_model::{BLOCK_SIZE, ModelConfig, PagedKVCache, Qwen3, loader};
use xserv_tensor::{DType, Device, Tensor};
use xserv_tokenizer::Tokenizer;
fn main() {
let args: Vec<String> = std::env::args().collect();
if args.len() < 3 {
eprintln!("Usage: check-eagle3 <target-model-dir> <eagle3-model-dir> [prompt]");
std::process::exit(1);
}
let target_dir = PathBuf::from(&args[1]);
let eagle_dir = PathBuf::from(&args[2]);
let prompt = args
.get(3)
.cloned()
.unwrap_or_else(|| "The capital of France is".to_string());
let device: u32 = 0;
xserv_cuda::device::set_device(device).unwrap();
let target_config = ModelConfig::from_file(&target_dir.join("config.json"));
eprintln!("Loading target Qwen3-8B...");
let target_weights = loader::load_model_dir(&target_dir, Device::Cuda(device));
let target = Qwen3::from_weights(target_config.clone(), target_weights);
xserv_cuda::allocator::cached_trim();
eprintln!("Loading EAGLE3 head from {}", eagle_dir.display());
let eagle = Eagle3Head::load(&eagle_dir, device);
xserv_cuda::allocator::cached_trim();
let tokenizer = Tokenizer::from_file(&target_dir.join("tokenizer.json"));
let embed_tokens = target.embed_tokens_tensor();
let ids = tokenizer.encode(&prompt);
let max_seq_len = 512;
let num_blocks = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE + 2;
let mut cache = PagedKVCache::new(
&target_config,
num_blocks,
0,
1,
num_blocks,
DType::BF16,
device,
);
cache.register_sequence(0).unwrap();
// Prefill target.
let logits = target.forward_prefill_paged(&ids, 0, &mut cache);
let target_first = *xserv_kernels::argmax_bf16_to_host(&logits).last().unwrap();
let target_first_text = tokenizer.decode(&[target_first]);
println!("Prompt: {:?}", prompt);
println!(
"Target argmax after prefill: {} ({:?})",
target_first, target_first_text
);
// Now run one target decode step with target_first to get hidden states at the
// hook layers.
let pos = cache.seq_len(0);
target.decode_prepare(&[pos], &[0], &mut cache);
let ids_gpu = upload_u32(&[target_first]);
let pos_gpu = upload_u32(&[pos as u32]);
let (target_next_logits, hooks) = target.decode_core_with_hidden(
ids_gpu.as_ptr() as *const std::ffi::c_void,
pos_gpu.as_ptr() as *const std::ffi::c_void,
1,
&[0],
&mut cache,
&EAGLE_HOOK_LAYERS,
);
let target_next = xserv_kernels::argmax_bf16_single(&target_next_logits);
let target_next_text = tokenizer.decode(&[target_next]);
println!(
"Target argmax after 1 decode step: {} ({:?})",
target_next, target_next_text
);
for (i, h) in hooks.iter().enumerate() {
println!(
"hook[{}] (layer {}): shape={:?} dtype={:?}",
i,
EAGLE_HOOK_LAYERS[i],
h.shape(),
h.dtype()
);
}
// Ask EAGLE what it thinks the NEXT token is (given target_first as prev_token
// and the hidden states from the position where target_first lives).
// EAGLE should predict target_next (or close to it) to be useful.
let (eagle_pred, eagle_logits) = eagle.step(&hooks, embed_tokens, target_first, pos);
let eagle_pred_text = tokenizer.decode(&[eagle_pred]);
println!(
"EAGLE draft prediction: {} ({:?})",
eagle_pred, eagle_pred_text
);
if eagle_pred == target_next {
println!("MATCH: EAGLE agrees with target on next token.");
} else {
println!(
"MISMATCH: EAGLE draft={} vs target={} (this is fine per-step; check top-5 below)",
eagle_pred, target_next
);
}
// Show top-5 from eagle logits (in draft vocab space, mapped to target).
print_top5(&eagle_logits, "EAGLE draft top-5", &eagle, &tokenizer);
}
fn upload_u32(vals: &[u32]) -> xserv_cuda::GpuBuffer {
let bytes = unsafe { std::slice::from_raw_parts(vals.as_ptr() as *const u8, vals.len() * 4) };
let mut buf = xserv_cuda::allocator::cached_alloc(bytes.len()).unwrap();
buf.copy_from_host(bytes).unwrap();
buf
}
fn print_top5(logits: &Tensor, label: &str, eagle: &Eagle3Head, tokenizer: &Tokenizer) {
use half::bf16;
let cpu = logits.to_device(Device::Cpu);
let data = cpu.as_slice::<bf16>();
let mut vals: Vec<(usize, f32)> = data
.iter()
.enumerate()
.map(|(i, v)| (i, v.to_f32()))
.collect();
vals.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
println!("{label}:");
for (i, val) in vals.iter().take(5) {
let target_id = eagle.map_draft_to_target(*i as u32);
let text = tokenizer.decode(&[target_id]);
println!(
" draft_id={} target_id={} val={:.3} text={:?}",
i, target_id, val, text
);
}
}