osobhandClaude Sonnet 5
4aaa36a57a
Whole-workspace rustfmt pass picked up while iterating on Mamba GPU backward work. Verified formatting-only via diff sampling; no logic changed. Co-Authored-By: Claude Sonnet 5 <[email protected]>
929 lines
34 KiB
Rust
929 lines
34 KiB
Rust
//! EAGLE-1 autoregressive draft head (Li et al. 2024, arXiv:2401.15077).
|
||
//!
|
||
//! EAGLE drafts multiple future tokens by running a single small draft head
|
||
//! *autoregressively* over a fused feature stream: at each step the head
|
||
//! consumes a fusion of the current hidden state and the embedding of the
|
||
//! previously drafted token, and produces both the next hidden state and a
|
||
//! logits distribution over the vocabulary for the next token. This is in
|
||
//! contrast to Medusa, which uses `K` independent heads each predicting a
|
||
//! fixed future offset.
|
||
//!
|
||
//! # Architecture
|
||
//!
|
||
//! Each [`EagleDraftHead`] layer computes:
|
||
//!
|
||
//! ```text
|
||
//! fused = fuse(hidden [H], token_embedding [H]) // per FusionMethod
|
||
//! act = SiLU(W1 @ fused + b1) // [H]
|
||
//! hidden' = W_h @ act + b_h // [H] next hidden state
|
||
//! logits = W_v @ act + b_v // [V] vocab logits
|
||
//! ```
|
||
//!
|
||
//! [`EagleHeads`] wraps a stack of `num_draft_layers` such heads plus a small
|
||
//! token-embedding table, and drives the autoregressive draft loop described
|
||
//! in the EAGLE paper:
|
||
//!
|
||
//! 1. Fuse the current hidden state with the embedding of the last token.
|
||
//! 2. Run the draft head to get a new hidden state and vocab logits.
|
||
//! 3. Pick the top token (argmax) — or the best of a top-k list for tree
|
||
//! construction — feed it back in as the "last token" for the next step.
|
||
//! 4. Repeat for `steps` iterations.
|
||
//!
|
||
//! # Example
|
||
//!
|
||
//! ```rust
|
||
//! use rtx_inference::eagle::{EagleHeads, EagleHeadsConfig};
|
||
//!
|
||
//! let config = EagleHeadsConfig::default();
|
||
//! let heads = EagleHeads::new(config);
|
||
//! let hidden = vec![0.0_f32; 64];
|
||
//! let draft = heads.draft(&hidden, &[1, 2, 3], 4);
|
||
//! assert_eq!(draft.len(), 4);
|
||
//! ```
|
||
|
||
use crate::InferenceResult;
|
||
use crate::speculative::{EagleConfig, EagleDraftModel, FusionMethod, Token};
|
||
|
||
// ── LCG pseudo-random number generator ────────────────────────────────────────
|
||
|
||
/// Minimal LCG PRNG used for reproducible weight initialisation (no external deps).
|
||
///
|
||
/// Mirrors the generator in [`crate::medusa`] so seeded init is deterministic
|
||
/// and reproducible across runs without pulling in `rand`.
|
||
struct Lcg(u64);
|
||
|
||
impl Lcg {
|
||
fn new(seed: u64) -> Self {
|
||
Self(seed ^ 0x1234_5678_9abc_def0)
|
||
}
|
||
|
||
/// Advance the state and return a value in `[0, 1)`.
|
||
fn next_f32(&mut self) -> f32 {
|
||
self.0 = self
|
||
.0
|
||
.wrapping_mul(6_364_136_223_846_793_005)
|
||
.wrapping_add(1_442_695_040_888_963_407);
|
||
let bits = 0x3f80_0000_u32 | ((self.0 >> 41) as u32 & 0x007f_ffff);
|
||
f32::from_bits(bits) - 1.0
|
||
}
|
||
|
||
/// Return a value uniformly in `[-scale, scale]`.
|
||
fn next_scaled(&mut self, scale: f32) -> f32 {
|
||
(self.next_f32() * 2.0 - 1.0) * scale
|
||
}
|
||
}
|
||
|
||
// ── Fusion helper ─────────────────────────────────────────────────────────────
|
||
|
||
/// Fuse a hidden state with a token embedding per [`FusionMethod`].
|
||
///
|
||
/// * [`FusionMethod::Add`] and [`FusionMethod::Attention`] return a vector of
|
||
/// length `hidden_dim` (elementwise combination).
|
||
/// * [`FusionMethod::Concat`] returns a vector of length `2 * hidden_dim`.
|
||
///
|
||
/// `Attention` uses a lightweight scalar-gate approximation: the gate is the
|
||
/// sigmoid of the dot product between `hidden` and `embedding`, scaled by
|
||
/// `1/sqrt(hidden_dim)` (a single-head, single-query attention score), and is
|
||
/// used to interpolate between `hidden` and `embedding`. This keeps the
|
||
/// draft head architecture-agnostic to the fusion strategy without requiring
|
||
/// a full multi-head attention implementation for a draft-only path.
|
||
fn fuse_features(hidden: &[f32], embedding: &[f32], method: FusionMethod) -> Vec<f32> {
|
||
debug_assert_eq!(
|
||
hidden.len(),
|
||
embedding.len(),
|
||
"hidden/embedding dims must match"
|
||
);
|
||
match method {
|
||
FusionMethod::Concat => {
|
||
let mut out = Vec::with_capacity(hidden.len() + embedding.len());
|
||
out.extend_from_slice(hidden);
|
||
out.extend_from_slice(embedding);
|
||
out
|
||
}
|
||
FusionMethod::Add => hidden
|
||
.iter()
|
||
.zip(embedding.iter())
|
||
.map(|(h, e)| h + e)
|
||
.collect(),
|
||
FusionMethod::Attention => {
|
||
let dim = hidden.len().max(1);
|
||
let score: f32 = hidden
|
||
.iter()
|
||
.zip(embedding.iter())
|
||
.map(|(h, e)| h * e)
|
||
.sum::<f32>()
|
||
/ (dim as f32).sqrt();
|
||
let gate = 1.0 / (1.0 + (-score).exp());
|
||
hidden
|
||
.iter()
|
||
.zip(embedding.iter())
|
||
.map(|(h, e)| gate * h + (1.0 - gate) * e)
|
||
.collect()
|
||
}
|
||
}
|
||
}
|
||
|
||
/// The input dimension a draft head layer expects for a given `hidden_dim`
|
||
/// and [`FusionMethod`].
|
||
fn fused_dim(hidden_dim: usize, method: FusionMethod) -> usize {
|
||
match method {
|
||
FusionMethod::Concat => hidden_dim * 2,
|
||
FusionMethod::Add | FusionMethod::Attention => hidden_dim,
|
||
}
|
||
}
|
||
|
||
// ── EagleDraftHead ────────────────────────────────────────────────────────────
|
||
|
||
/// One autoregressive EAGLE draft head layer.
|
||
///
|
||
/// Consumes fused features of length `fused_dim` and produces:
|
||
/// * a new hidden state of length `hidden_dim`
|
||
/// * vocab logits of length `vocab_size`
|
||
#[derive(Debug, Clone)]
|
||
pub struct EagleDraftHead {
|
||
/// W1 weight matrix, row-major `[hidden_dim × fused_dim]`.
|
||
pub w1: Vec<f32>,
|
||
/// Bias for layer 1, `[hidden_dim]`.
|
||
pub b1: Vec<f32>,
|
||
/// Next-hidden-state projection, row-major `[hidden_dim × hidden_dim]`.
|
||
pub w_hidden: Vec<f32>,
|
||
/// Bias for the hidden-state projection, `[hidden_dim]`.
|
||
pub b_hidden: Vec<f32>,
|
||
/// Vocab logits projection, row-major `[vocab_size × hidden_dim]`.
|
||
pub w_vocab: Vec<f32>,
|
||
/// Bias for the vocab projection, `[vocab_size]`.
|
||
pub b_vocab: Vec<f32>,
|
||
/// Fused input dimension (depends on `fusion_method`).
|
||
pub fused_dim: usize,
|
||
/// Hidden dimension `H`.
|
||
pub hidden_dim: usize,
|
||
/// Vocabulary size `V`.
|
||
pub vocab_size: usize,
|
||
}
|
||
|
||
/// Output of a single [`EagleDraftHead::forward`] call.
|
||
#[derive(Debug, Clone)]
|
||
pub struct EagleHeadOutput {
|
||
/// The next hidden state, length `hidden_dim`.
|
||
pub hidden: Vec<f32>,
|
||
/// Vocab logits, length `vocab_size`.
|
||
pub logits: Vec<f32>,
|
||
}
|
||
|
||
impl EagleDraftHead {
|
||
/// Initialise with LCG-seeded random weights.
|
||
///
|
||
/// Weights are drawn from `U[-scale, scale]` where `scale = sqrt(2 /
|
||
/// fan_in)` (Xavier-like) for each projection independently. Biases are
|
||
/// zeroed.
|
||
pub fn new_random(fused_dim: usize, hidden_dim: usize, vocab_size: usize, seed: u64) -> Self {
|
||
let mut rng = Lcg::new(seed);
|
||
|
||
let scale_w1 = (2.0_f32 / fused_dim.max(1) as f32).sqrt();
|
||
let w1: Vec<f32> = (0..hidden_dim * fused_dim)
|
||
.map(|_| rng.next_scaled(scale_w1))
|
||
.collect();
|
||
let b1 = vec![0.0_f32; hidden_dim];
|
||
|
||
let scale_wh = (2.0_f32 / hidden_dim.max(1) as f32).sqrt();
|
||
let w_hidden: Vec<f32> = (0..hidden_dim * hidden_dim)
|
||
.map(|_| rng.next_scaled(scale_wh))
|
||
.collect();
|
||
let b_hidden = vec![0.0_f32; hidden_dim];
|
||
|
||
let scale_wv = (2.0_f32 / hidden_dim.max(1) as f32).sqrt();
|
||
let w_vocab: Vec<f32> = (0..vocab_size * hidden_dim)
|
||
.map(|_| rng.next_scaled(scale_wv))
|
||
.collect();
|
||
let b_vocab = vec![0.0_f32; vocab_size];
|
||
|
||
Self {
|
||
w1,
|
||
b1,
|
||
w_hidden,
|
||
b_hidden,
|
||
w_vocab,
|
||
b_vocab,
|
||
fused_dim,
|
||
hidden_dim,
|
||
vocab_size,
|
||
}
|
||
}
|
||
|
||
/// Run the head forward pass over pre-fused input features.
|
||
///
|
||
/// `fused` must have length `fused_dim`. Returns the next hidden state
|
||
/// (`hidden_dim`) and vocab logits (`vocab_size`).
|
||
///
|
||
/// # Panics
|
||
///
|
||
/// Panics in debug builds if `fused.len() != self.fused_dim`.
|
||
pub fn forward(&self, fused: &[f32]) -> EagleHeadOutput {
|
||
debug_assert_eq!(
|
||
fused.len(),
|
||
self.fused_dim,
|
||
"fused feature length must equal fused_dim"
|
||
);
|
||
let h = self.hidden_dim;
|
||
let f = self.fused_dim;
|
||
|
||
// Layer 1: act = SiLU(W1 @ fused + b1)
|
||
let mut act = vec![0.0_f32; h];
|
||
for i in 0..h {
|
||
let dot: f32 = (0..f).map(|j| self.w1[i * f + j] * fused[j]).sum::<f32>() + self.b1[i];
|
||
act[i] = dot * (1.0 / (1.0 + (-dot).exp()));
|
||
}
|
||
|
||
// Next hidden state: hidden' = W_hidden @ act + b_hidden
|
||
let mut hidden = vec![0.0_f32; h];
|
||
for i in 0..h {
|
||
hidden[i] = (0..h)
|
||
.map(|j| self.w_hidden[i * h + j] * act[j])
|
||
.sum::<f32>()
|
||
+ self.b_hidden[i];
|
||
}
|
||
|
||
// Vocab logits: logits = W_vocab @ act + b_vocab
|
||
let v = self.vocab_size;
|
||
let mut logits = vec![0.0_f32; v];
|
||
for i in 0..v {
|
||
logits[i] = (0..h)
|
||
.map(|j| self.w_vocab[i * h + j] * act[j])
|
||
.sum::<f32>()
|
||
+ self.b_vocab[i];
|
||
}
|
||
|
||
EagleHeadOutput { hidden, logits }
|
||
}
|
||
|
||
/// Return the indices of the top-`k` tokens sorted descending by logit.
|
||
///
|
||
/// `k` is clamped to `vocab_size` if it exceeds it.
|
||
pub fn top_k(&self, logits: &[f32], k: usize) -> Vec<(u32, f32)> {
|
||
let k = k.min(logits.len());
|
||
let mut idx: Vec<u32> = (0..logits.len() as u32).collect();
|
||
idx.sort_unstable_by(|&a, &b| {
|
||
logits[b as usize]
|
||
.partial_cmp(&logits[a as usize])
|
||
.unwrap_or(std::cmp::Ordering::Equal)
|
||
});
|
||
idx[..k].iter().map(|&i| (i, logits[i as usize])).collect()
|
||
}
|
||
}
|
||
|
||
// ── EagleHeadsConfig ───────────────────────────────────────────────────────────
|
||
|
||
/// Configuration for the [`EagleHeads`] draft model.
|
||
#[derive(Debug, Clone)]
|
||
pub struct EagleHeadsConfig {
|
||
/// Hidden dimension shared by the draft head(s) and the base model.
|
||
pub hidden_dim: usize,
|
||
/// Vocabulary size.
|
||
pub vocab_size: usize,
|
||
/// Number of stacked draft-head layers (EAGLE-1 typically uses 1).
|
||
pub num_draft_layers: usize,
|
||
/// How to fuse the hidden state with the previous token's embedding.
|
||
pub fusion_method: FusionMethod,
|
||
/// Number of candidate tokens kept per autoregressive step.
|
||
pub top_k: usize,
|
||
}
|
||
|
||
impl Default for EagleHeadsConfig {
|
||
fn default() -> Self {
|
||
Self {
|
||
hidden_dim: 64,
|
||
vocab_size: 256,
|
||
num_draft_layers: 1,
|
||
fusion_method: FusionMethod::Concat,
|
||
top_k: 3,
|
||
}
|
||
}
|
||
}
|
||
|
||
impl From<&EagleConfig> for EagleHeadsConfig {
|
||
fn from(cfg: &EagleConfig) -> Self {
|
||
Self {
|
||
hidden_dim: cfg.hidden_dim,
|
||
vocab_size: 32_000, // orchestration EagleConfig has no vocab_size; use a safe default
|
||
num_draft_layers: cfg.num_draft_layers.max(1),
|
||
fusion_method: cfg.fusion_method,
|
||
top_k: cfg.top_k,
|
||
}
|
||
}
|
||
}
|
||
|
||
// ── EagleDraftStep / EagleLossResult ──────────────────────────────────────────
|
||
|
||
/// One step of an autoregressive draft, as returned by [`EagleHeads::draft`].
|
||
#[derive(Debug, Clone)]
|
||
pub struct EagleDraftStep {
|
||
/// Top token chosen for this step (argmax of `candidates`).
|
||
pub token: u32,
|
||
/// Probability-proxy (raw logit) of the chosen token.
|
||
pub score: f32,
|
||
/// Top-`top_k` candidates at this step, `(token_id, logit)`, sorted
|
||
/// descending by logit — usable to build a candidate tree.
|
||
pub candidates: Vec<(u32, f32)>,
|
||
/// Hidden state produced at this step (fed forward to the next step).
|
||
pub hidden: Vec<f32>,
|
||
}
|
||
|
||
/// Cross-entropy loss result returned by [`EagleHeads::training_loss`].
|
||
#[derive(Debug, Clone)]
|
||
pub struct EagleLossResult {
|
||
/// Per-step cross-entropy loss, indexed by autoregressive step.
|
||
pub per_step_loss: Vec<f32>,
|
||
/// Mean cross-entropy across all steps.
|
||
pub total_loss: f32,
|
||
}
|
||
|
||
// ── EagleHeads ─────────────────────────────────────────────────────────────────
|
||
|
||
/// The concrete EAGLE-1 draft model: a stack of [`EagleDraftHead`] layers plus
|
||
/// a token-embedding table, driving the autoregressive draft loop described
|
||
/// in arXiv:2401.15077.
|
||
///
|
||
/// Unlike Medusa's `K` independent heads (each predicting a fixed future
|
||
/// offset from the *same* base hidden state), EAGLE re-uses (a stack of) one
|
||
/// small head autoregressively: the hidden state produced at step `t` is fed
|
||
/// back in, fused with the embedding of the token drafted at step `t`, to
|
||
/// produce step `t + 1`. This lets the draft model track feature-level
|
||
/// uncertainty from step to step, which the EAGLE paper shows improves
|
||
/// acceptance length over Medusa's independent-head approach.
|
||
pub struct EagleHeads {
|
||
/// Shared configuration.
|
||
pub config: EagleHeadsConfig,
|
||
/// Stacked draft-head layers (only the first is used autoregressively in
|
||
/// the paper's EAGLE-1 design; deeper stacks are supported for
|
||
/// experimentation with `num_draft_layers > 1`).
|
||
pub layers: Vec<EagleDraftHead>,
|
||
/// Token embedding table, row-major `[vocab_size × hidden_dim]`.
|
||
pub embedding: Vec<f32>,
|
||
}
|
||
|
||
impl EagleHeads {
|
||
/// Create a new draft model with random weights.
|
||
pub fn new(config: EagleHeadsConfig) -> Self {
|
||
let fdim = fused_dim(config.hidden_dim, config.fusion_method);
|
||
let layers: Vec<EagleDraftHead> = (0..config.num_draft_layers.max(1))
|
||
.map(|i| {
|
||
EagleDraftHead::new_random(
|
||
fdim,
|
||
config.hidden_dim,
|
||
config.vocab_size,
|
||
(i as u64)
|
||
.wrapping_mul(0x9e37_79b9_7f4a_7c15)
|
||
.wrapping_add(7),
|
||
)
|
||
})
|
||
.collect();
|
||
|
||
let mut emb_rng = Lcg::new(0xe5a1_e000);
|
||
let emb_scale = (1.0_f32 / config.hidden_dim.max(1) as f32).sqrt();
|
||
let embedding: Vec<f32> = (0..config.vocab_size * config.hidden_dim)
|
||
.map(|_| emb_rng.next_scaled(emb_scale))
|
||
.collect();
|
||
|
||
Self {
|
||
config,
|
||
layers,
|
||
embedding,
|
||
}
|
||
}
|
||
|
||
/// Look up the embedding row for `token`, clamped to a valid vocab index
|
||
/// (out-of-range tokens map to index 0, so drafting never panics on an
|
||
/// unexpected context token id).
|
||
fn embed(&self, token: u32) -> &[f32] {
|
||
let h = self.config.hidden_dim;
|
||
let idx = (token as usize).min(self.config.vocab_size.saturating_sub(1));
|
||
&self.embedding[idx * h..idx * h + h]
|
||
}
|
||
|
||
/// The draft head used for autoregressive stepping (the first layer).
|
||
fn head(&self) -> &EagleDraftHead {
|
||
&self.layers[0]
|
||
}
|
||
|
||
/// Run the EAGLE autoregressive draft loop.
|
||
///
|
||
/// `hidden` is the base model's last hidden state (`[hidden_dim]`).
|
||
/// `context_tokens` supplies the most recently generated token (its last
|
||
/// element seeds the first fusion step); if empty, token `0` is used.
|
||
/// `steps` is the number of autoregressive draft tokens to generate.
|
||
///
|
||
/// Returns one `(token, score)` pair per step — the argmax token and its
|
||
/// raw logit — where `score` approximates a probability (higher is more
|
||
/// confident); use [`EagleHeads::draft_steps`] for the full per-step
|
||
/// detail (candidates, hidden states) needed to build a candidate tree.
|
||
pub fn draft(&self, hidden: &[f32], context_tokens: &[u32], steps: usize) -> Vec<(u32, f32)> {
|
||
self.draft_steps(hidden, context_tokens, steps)
|
||
.into_iter()
|
||
.map(|s| (s.token, s.score))
|
||
.collect()
|
||
}
|
||
|
||
/// Like [`EagleHeads::draft`], but returns full per-step detail
|
||
/// (top-`k` candidates and the resulting hidden state at each step) so a
|
||
/// caller can build a candidate tree (cf. Medusa's `generate_tree` /
|
||
/// EAGLE-3's dynamic tree).
|
||
pub fn draft_steps(
|
||
&self,
|
||
hidden: &[f32],
|
||
context_tokens: &[u32],
|
||
steps: usize,
|
||
) -> Vec<EagleDraftStep> {
|
||
let head = self.head();
|
||
let mut cur_hidden = hidden.to_vec();
|
||
let mut last_token = context_tokens.last().copied().unwrap_or(0);
|
||
|
||
let mut out = Vec::with_capacity(steps);
|
||
for _ in 0..steps {
|
||
let embedding = self.embed(last_token);
|
||
let fused = fuse_features(&cur_hidden, embedding, self.config.fusion_method);
|
||
let result = head.forward(&fused);
|
||
|
||
let candidates = head.top_k(&result.logits, self.config.top_k.max(1));
|
||
let (top_token, top_score) = candidates
|
||
.first()
|
||
.copied()
|
||
.unwrap_or((0, f32::NEG_INFINITY));
|
||
|
||
cur_hidden = result.hidden.clone();
|
||
last_token = top_token;
|
||
|
||
out.push(EagleDraftStep {
|
||
token: top_token,
|
||
score: top_score,
|
||
candidates,
|
||
hidden: result.hidden,
|
||
});
|
||
}
|
||
out
|
||
}
|
||
|
||
/// Compute per-step cross-entropy training loss.
|
||
///
|
||
/// `hidden` is the base model's last hidden state (`[hidden_dim]`).
|
||
/// `context_tokens` seeds the first fusion step, matching [`Self::draft`].
|
||
/// `targets[t]` is the ground-truth token the model should draft at
|
||
/// autoregressive step `t` (teacher-forced: the *target* token, not the
|
||
/// model's own prediction, is fed forward into the next step's fusion —
|
||
/// this is the standard teacher-forcing training regime for EAGLE).
|
||
///
|
||
/// Loss is computed with the numerically stable log-sum-exp trick.
|
||
pub fn training_loss(
|
||
&self,
|
||
hidden: &[f32],
|
||
context_tokens: &[u32],
|
||
targets: &[u32],
|
||
) -> EagleLossResult {
|
||
let head = self.head();
|
||
let mut cur_hidden = hidden.to_vec();
|
||
let mut last_token = context_tokens.last().copied().unwrap_or(0);
|
||
|
||
let mut per_step_loss = Vec::with_capacity(targets.len());
|
||
for &target in targets {
|
||
let embedding = self.embed(last_token);
|
||
let fused = fuse_features(&cur_hidden, embedding, self.config.fusion_method);
|
||
let result = head.forward(&fused);
|
||
|
||
let logits = &result.logits;
|
||
let max_l = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
|
||
let log_sum = max_l + logits.iter().map(|l| (l - max_l).exp()).sum::<f32>().ln();
|
||
let ce = -(logits[target as usize] - log_sum);
|
||
per_step_loss.push(ce);
|
||
|
||
// Teacher forcing: advance state using the *target* token.
|
||
cur_hidden = result.hidden;
|
||
last_token = target;
|
||
}
|
||
|
||
let total_loss = if per_step_loss.is_empty() {
|
||
0.0
|
||
} else {
|
||
per_step_loss.iter().sum::<f32>() / per_step_loss.len() as f32
|
||
};
|
||
|
||
EagleLossResult {
|
||
per_step_loss,
|
||
total_loss,
|
||
}
|
||
}
|
||
}
|
||
|
||
// ── EagleDraftModel trait impl ─────────────────────────────────────────────────
|
||
|
||
#[async_trait::async_trait]
|
||
impl EagleDraftModel for EagleHeads {
|
||
async fn draft_from_hidden(
|
||
&self,
|
||
hidden_states: &[f32],
|
||
context_tokens: &[u32],
|
||
config: &EagleConfig,
|
||
) -> InferenceResult<Vec<Token>> {
|
||
let steps = self.draft_steps(hidden_states, context_tokens, config.draft_steps.max(1));
|
||
Ok(steps
|
||
.into_iter()
|
||
.map(|step| {
|
||
// Softmax over the top-k candidates for a normalised probability.
|
||
let max_l = step
|
||
.candidates
|
||
.iter()
|
||
.map(|&(_, l)| l)
|
||
.fold(f32::NEG_INFINITY, f32::max);
|
||
let sum_exp: f32 = step
|
||
.candidates
|
||
.iter()
|
||
.map(|&(_, l)| (l - max_l).exp())
|
||
.sum();
|
||
let prob = if sum_exp > 0.0 {
|
||
((step.score - max_l).exp()) / sum_exp
|
||
} else {
|
||
0.0
|
||
};
|
||
Token {
|
||
id: step.token,
|
||
text: format!("tok_{}", step.token),
|
||
logits: step.candidates.iter().map(|&(_, l)| l).collect(),
|
||
probability: prob,
|
||
}
|
||
})
|
||
.collect())
|
||
}
|
||
|
||
fn hidden_dim(&self) -> usize {
|
||
self.config.hidden_dim
|
||
}
|
||
}
|
||
|
||
// ── Tests ─────────────────────────────────────────────────────────────────────
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
// ── helpers ──
|
||
|
||
fn default_model() -> EagleHeads {
|
||
EagleHeads::new(EagleHeadsConfig::default())
|
||
}
|
||
|
||
fn zero_hidden(dim: usize) -> Vec<f32> {
|
||
vec![0.0_f32; dim]
|
||
}
|
||
|
||
fn ones_hidden(dim: usize) -> Vec<f32> {
|
||
vec![1.0_f32; dim]
|
||
}
|
||
|
||
// ── config ───────────────────────────────────────────────────────────────
|
||
|
||
#[test]
|
||
fn test_default_config() {
|
||
let cfg = EagleHeadsConfig::default();
|
||
assert_eq!(cfg.hidden_dim, 64);
|
||
assert_eq!(cfg.vocab_size, 256);
|
||
assert_eq!(cfg.num_draft_layers, 1);
|
||
assert_eq!(cfg.fusion_method, FusionMethod::Concat);
|
||
assert_eq!(cfg.top_k, 3);
|
||
}
|
||
|
||
// ── fuse_features ────────────────────────────────────────────────────────
|
||
|
||
#[test]
|
||
fn test_fuse_concat_length() {
|
||
let h = vec![1.0_f32, 2.0, 3.0];
|
||
let e = vec![4.0_f32, 5.0, 6.0];
|
||
let fused = fuse_features(&h, &e, FusionMethod::Concat);
|
||
assert_eq!(fused.len(), 6);
|
||
assert_eq!(fused, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
|
||
}
|
||
|
||
#[test]
|
||
fn test_fuse_add_length_and_values() {
|
||
let h = vec![1.0_f32, 2.0, 3.0];
|
||
let e = vec![4.0_f32, 5.0, 6.0];
|
||
let fused = fuse_features(&h, &e, FusionMethod::Add);
|
||
assert_eq!(fused.len(), 3);
|
||
assert_eq!(fused, vec![5.0, 7.0, 9.0]);
|
||
}
|
||
|
||
#[test]
|
||
fn test_fuse_attention_length_and_finite() {
|
||
let h = vec![1.0_f32, 2.0, 3.0];
|
||
let e = vec![0.5_f32, -1.0, 2.0];
|
||
let fused = fuse_features(&h, &e, FusionMethod::Attention);
|
||
assert_eq!(fused.len(), 3);
|
||
assert!(fused.iter().all(|v| v.is_finite()));
|
||
}
|
||
|
||
#[test]
|
||
fn test_fused_dim_matches_method() {
|
||
assert_eq!(fused_dim(64, FusionMethod::Concat), 128);
|
||
assert_eq!(fused_dim(64, FusionMethod::Add), 64);
|
||
assert_eq!(fused_dim(64, FusionMethod::Attention), 64);
|
||
}
|
||
|
||
// ── EagleDraftHead::forward ──────────────────────────────────────────────
|
||
|
||
#[test]
|
||
fn test_head_forward_shape() {
|
||
let cfg = EagleHeadsConfig::default();
|
||
let fdim = fused_dim(cfg.hidden_dim, cfg.fusion_method);
|
||
let head = EagleDraftHead::new_random(fdim, cfg.hidden_dim, cfg.vocab_size, 1);
|
||
let fused = zero_hidden(fdim);
|
||
let out = head.forward(&fused);
|
||
assert_eq!(out.hidden.len(), cfg.hidden_dim);
|
||
assert_eq!(out.logits.len(), cfg.vocab_size);
|
||
}
|
||
|
||
#[test]
|
||
fn test_head_forward_finite() {
|
||
let cfg = EagleHeadsConfig::default();
|
||
let fdim = fused_dim(cfg.hidden_dim, cfg.fusion_method);
|
||
let head = EagleDraftHead::new_random(fdim, cfg.hidden_dim, cfg.vocab_size, 2);
|
||
let fused = ones_hidden(fdim);
|
||
let out = head.forward(&fused);
|
||
assert!(out.hidden.iter().all(|v| v.is_finite()));
|
||
assert!(out.logits.iter().all(|v| v.is_finite()));
|
||
}
|
||
|
||
// ── EagleDraftHead::top_k ────────────────────────────────────────────────
|
||
|
||
#[test]
|
||
fn test_head_top_k_length_and_sorted() {
|
||
let cfg = EagleHeadsConfig::default();
|
||
let fdim = fused_dim(cfg.hidden_dim, cfg.fusion_method);
|
||
let head = EagleDraftHead::new_random(fdim, cfg.hidden_dim, cfg.vocab_size, 3);
|
||
let logits = head.forward(&ones_hidden(fdim)).logits;
|
||
let result = head.top_k(&logits, 5);
|
||
assert_eq!(result.len(), 5);
|
||
for w in result.windows(2) {
|
||
assert!(w[0].1 >= w[1].1, "top-k must be sorted descending by logit");
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_head_top_k_clamped_to_vocab() {
|
||
let cfg = EagleHeadsConfig::default();
|
||
let fdim = fused_dim(cfg.hidden_dim, cfg.fusion_method);
|
||
let head = EagleDraftHead::new_random(fdim, cfg.hidden_dim, cfg.vocab_size, 4);
|
||
let logits = head.forward(&zero_hidden(fdim)).logits;
|
||
let result = head.top_k(&logits, cfg.vocab_size + 50);
|
||
assert_eq!(result.len(), cfg.vocab_size);
|
||
}
|
||
|
||
// ── EagleHeads::new ──────────────────────────────────────────────────────
|
||
|
||
#[test]
|
||
fn test_new_builds_correct_layer_count() {
|
||
let cfg = EagleHeadsConfig {
|
||
num_draft_layers: 3,
|
||
..EagleHeadsConfig::default()
|
||
};
|
||
let model = EagleHeads::new(cfg.clone());
|
||
assert_eq!(model.layers.len(), 3);
|
||
}
|
||
|
||
#[test]
|
||
fn test_new_embedding_table_shape() {
|
||
let cfg = EagleHeadsConfig::default();
|
||
let model = EagleHeads::new(cfg.clone());
|
||
assert_eq!(model.embedding.len(), cfg.vocab_size * cfg.hidden_dim);
|
||
}
|
||
|
||
#[test]
|
||
fn test_seeded_init_is_deterministic() {
|
||
let cfg = EagleHeadsConfig::default();
|
||
let a = EagleHeads::new(cfg.clone());
|
||
let b = EagleHeads::new(cfg);
|
||
assert_eq!(a.layers[0].w1, b.layers[0].w1);
|
||
assert_eq!(a.embedding, b.embedding);
|
||
}
|
||
|
||
// ── EagleHeads::draft / draft_steps ──────────────────────────────────────
|
||
|
||
#[test]
|
||
fn test_draft_length() {
|
||
let model = default_model();
|
||
let hidden = zero_hidden(model.config.hidden_dim);
|
||
let draft = model.draft(&hidden, &[1, 2, 3], 4);
|
||
assert_eq!(draft.len(), 4);
|
||
}
|
||
|
||
#[test]
|
||
fn test_draft_zero_steps() {
|
||
let model = default_model();
|
||
let hidden = zero_hidden(model.config.hidden_dim);
|
||
let draft = model.draft(&hidden, &[1, 2, 3], 0);
|
||
assert!(draft.is_empty());
|
||
}
|
||
|
||
#[test]
|
||
fn test_draft_empty_context_does_not_panic() {
|
||
let model = default_model();
|
||
let hidden = zero_hidden(model.config.hidden_dim);
|
||
let draft = model.draft(&hidden, &[], 2);
|
||
assert_eq!(draft.len(), 2);
|
||
}
|
||
|
||
#[test]
|
||
fn test_draft_out_of_range_context_token_does_not_panic() {
|
||
let model = default_model();
|
||
let hidden = zero_hidden(model.config.hidden_dim);
|
||
// Token id far beyond vocab_size must be clamped, not panic.
|
||
let draft = model.draft(&hidden, &[u32::MAX], 2);
|
||
assert_eq!(draft.len(), 2);
|
||
}
|
||
|
||
#[test]
|
||
fn test_draft_steps_candidates_and_hidden_shapes() {
|
||
let model = default_model();
|
||
let hidden = zero_hidden(model.config.hidden_dim);
|
||
let steps = model.draft_steps(&hidden, &[5], 3);
|
||
assert_eq!(steps.len(), 3);
|
||
for step in &steps {
|
||
assert_eq!(step.hidden.len(), model.config.hidden_dim);
|
||
assert_eq!(step.candidates.len(), model.config.top_k);
|
||
assert!(step.score.is_finite());
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_draft_steps_top_candidate_matches_token() {
|
||
let model = default_model();
|
||
let hidden = ones_hidden(model.config.hidden_dim);
|
||
let steps = model.draft_steps(&hidden, &[0], 2);
|
||
for step in &steps {
|
||
assert_eq!(step.candidates[0].0, step.token);
|
||
assert!((step.candidates[0].1 - step.score).abs() < 1e-6);
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_draft_all_fusion_methods() {
|
||
for method in [
|
||
FusionMethod::Concat,
|
||
FusionMethod::Add,
|
||
FusionMethod::Attention,
|
||
] {
|
||
let cfg = EagleHeadsConfig {
|
||
fusion_method: method,
|
||
..EagleHeadsConfig::default()
|
||
};
|
||
let model = EagleHeads::new(cfg.clone());
|
||
let hidden = zero_hidden(cfg.hidden_dim);
|
||
let draft = model.draft(&hidden, &[1, 2], 3);
|
||
assert_eq!(
|
||
draft.len(),
|
||
3,
|
||
"fusion method {method:?} produced wrong draft length"
|
||
);
|
||
assert!(draft.iter().all(|&(_, s)| s.is_finite()));
|
||
}
|
||
}
|
||
|
||
// ── EagleHeads::training_loss ────────────────────────────────────────────
|
||
|
||
#[test]
|
||
fn test_training_loss_shape() {
|
||
let model = default_model();
|
||
let hidden = zero_hidden(model.config.hidden_dim);
|
||
let targets = vec![0_u32, 1, 2];
|
||
let result = model.training_loss(&hidden, &[1], &targets);
|
||
assert_eq!(result.per_step_loss.len(), targets.len());
|
||
}
|
||
|
||
#[test]
|
||
fn test_training_loss_non_negative_and_finite() {
|
||
let model = default_model();
|
||
let hidden = ones_hidden(model.config.hidden_dim);
|
||
let targets = vec![0_u32; 4];
|
||
let result = model.training_loss(&hidden, &[2], &targets);
|
||
for (i, &loss) in result.per_step_loss.iter().enumerate() {
|
||
assert!(
|
||
loss >= 0.0,
|
||
"step {i}: CE loss must be non-negative, got {loss}"
|
||
);
|
||
assert!(loss.is_finite());
|
||
}
|
||
assert!(result.total_loss.is_finite());
|
||
}
|
||
|
||
#[test]
|
||
fn test_training_loss_total_is_mean() {
|
||
let model = default_model();
|
||
let hidden = zero_hidden(model.config.hidden_dim);
|
||
let targets = vec![1_u32; 5];
|
||
let result = model.training_loss(&hidden, &[0], &targets);
|
||
let expected_mean =
|
||
result.per_step_loss.iter().sum::<f32>() / result.per_step_loss.len() as f32;
|
||
assert!((result.total_loss - expected_mean).abs() < 1e-5);
|
||
}
|
||
|
||
#[test]
|
||
fn test_training_loss_empty_targets() {
|
||
let model = default_model();
|
||
let hidden = zero_hidden(model.config.hidden_dim);
|
||
let result = model.training_loss(&hidden, &[0], &[]);
|
||
assert!(result.per_step_loss.is_empty());
|
||
assert_eq!(result.total_loss, 0.0);
|
||
}
|
||
|
||
#[test]
|
||
fn test_training_loss_near_zero_for_dominant_logit() {
|
||
// Build a single-layer head whose vocab row for token 0 dominates,
|
||
// so CE loss for target=0 should be close to 0.
|
||
let hidden_dim = 4;
|
||
let vocab_size = 8;
|
||
let fdim = fused_dim(hidden_dim, FusionMethod::Add);
|
||
let target: u32 = 0;
|
||
|
||
let mut w_vocab = vec![0.0_f32; vocab_size * hidden_dim];
|
||
for j in 0..hidden_dim {
|
||
w_vocab[target as usize * hidden_dim + j] = 100.0;
|
||
}
|
||
let head = EagleDraftHead {
|
||
w1: {
|
||
// Identity-like mapping from fused (Add => hidden_dim) to hidden_dim.
|
||
let mut m = vec![0.0_f32; hidden_dim * fdim];
|
||
for i in 0..hidden_dim {
|
||
m[i * fdim + i] = 1.0;
|
||
}
|
||
m
|
||
},
|
||
b1: vec![0.0; hidden_dim],
|
||
w_hidden: vec![0.0; hidden_dim * hidden_dim],
|
||
b_hidden: vec![0.0; hidden_dim],
|
||
w_vocab,
|
||
b_vocab: vec![0.0; vocab_size],
|
||
fused_dim: fdim,
|
||
hidden_dim,
|
||
vocab_size,
|
||
};
|
||
|
||
let fused = vec![1.0_f32; fdim];
|
||
let out = head.forward(&fused);
|
||
let logits = out.logits;
|
||
let max_l = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
|
||
let log_sum = max_l + logits.iter().map(|l| (l - max_l).exp()).sum::<f32>().ln();
|
||
let ce = -(logits[target as usize] - log_sum);
|
||
assert!(
|
||
ce < 0.01,
|
||
"CE loss should be near 0 when target logit dominates, got {ce}"
|
||
);
|
||
}
|
||
|
||
// ── EagleDraftModel trait impl ───────────────────────────────────────────
|
||
|
||
#[tokio::test]
|
||
async fn test_trait_draft_from_hidden_smoke() {
|
||
let model = default_model();
|
||
let hidden = zero_hidden(model.config.hidden_dim);
|
||
let eagle_cfg = EagleConfig {
|
||
draft_steps: 3,
|
||
hidden_dim: model.config.hidden_dim,
|
||
num_draft_layers: 1,
|
||
fusion_method: model.config.fusion_method,
|
||
top_k: model.config.top_k,
|
||
};
|
||
|
||
let tokens = (&model as &dyn EagleDraftModel)
|
||
.draft_from_hidden(&hidden, &[1, 2, 3], &eagle_cfg)
|
||
.await
|
||
.expect("draft_from_hidden must succeed");
|
||
|
||
assert_eq!(tokens.len(), 3);
|
||
for t in &tokens {
|
||
assert!(t.probability.is_finite());
|
||
assert!(!t.text.is_empty());
|
||
}
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_trait_hidden_dim_accessor() {
|
||
let model = default_model();
|
||
let dim = (&model as &dyn EagleDraftModel).hidden_dim();
|
||
assert_eq!(dim, model.config.hidden_dim);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_trait_zero_draft_steps_yields_empty() {
|
||
let model = default_model();
|
||
let hidden = zero_hidden(model.config.hidden_dim);
|
||
let eagle_cfg = EagleConfig {
|
||
draft_steps: 0,
|
||
..EagleConfig::default()
|
||
};
|
||
// `.max(1)` inside draft_from_hidden ensures at least 1 step even if
|
||
// config requests 0, matching the crate-wide convention of never
|
||
// silently no-op'ing a draft call.
|
||
let tokens = (&model as &dyn EagleDraftModel)
|
||
.draft_from_hidden(&hidden, &[0], &eagle_cfg)
|
||
.await
|
||
.expect("draft_from_hidden must succeed");
|
||
assert_eq!(tokens.len(), 1);
|
||
}
|
||
}
|