Skip to main content

Module eagle

Module eagle 

Source
Expand description

EAGLE3.1 greedy-chain speculative decode (research/basics/EAGLE-PLAN.md, N1-N7).

Greedy spec decode is MATHEMATICALLY EXACT: the accepted+bonus token stream is token-for-token identical to plain greedy generate (decode.rs). EAGLE differs from MTP (spec.rs) ONLY in the DRAFT step: instead of the trunk-coupled NextN head, EAGLE drafts with a SEPARATE 1-layer model (own vocab, own RoPE, untied lm_head) fed the trunk’s hidden states from 3 aux layers [1,15,28] fused through an encoder fc. The verify / accept-prefix / snapshot / rollback are REUSED VERBATIM from spec.rs (decode_step_t, the greedy accept walk, cache.snapshot/rollback).

On-disk draft (eagle3-qwen35-9b/model.safetensors, bf16, ground-truthed at impl time): fc.weight [4096, 12288] (3n_embd -> n_embd encoder) midlayer.input_layernorm.weight [4096] (RMSNorm of the prev-token EMBED) midlayer.hidden_norm.weight [4096] (RMSNorm of the recurrent hidden g) midlayer.self_attn.{q,k,v}_proj q[4096,8192] k/v[1024,8192] (in = 2n_embd!) midlayer.self_attn.o_proj [4096, 4096] midlayer.post_attention_layernorm [4096] midlayer.mlp.{gate,up}_proj [12288,4096] down [4096,12288] norm.weight [4096] (final RMSNorm before lm_head) lm_head.weight [32000, 4096] (DRAFT vocab) d2t [32000] i64 target_id = draft_id + d2t[draft_id] t2d [248320] bool (unused on the chain-greedy decode path)

Op-sequence (authoritative: vLLM llama_eagle3.py LlamaDecoderLayer layer_idx==0, this ckpt’s flags norm_before_residual=false, norm_before_fc=false, fc_norm=false, norm_output=false): ENCODE (once/round): g = fc @ concat(aux[1], aux[15], aux[28]) -> [n_embd] DRAFT step (T=1): e = embed(prev_tok) (TARGET embedding; EAGLE3 shares it) eN = RMSNorm(e, input_layernorm) res = g (_norm_after_residual: residual is PRE-norm g) gN = RMSNorm(g, hidden_norm) cat = [eN ; gN] -> [2*n_embd] attn= o_proj @ SDPA( q,k,v = {q,k,v}_proj @ cat ; partial RoPE 64/256 @ theta 1e7 ; GQA16:4 ) x1 = attn + res z = RMSNorm(x1, post_attention_layernorm) mlp = down @ silu(gate @ z) * (up @ z) gsum= mlp + x1 (the model’s final fused-add residual) dl = lm_head @ RMSNorm(gsum, norm) -> draft_logits[32000] g_next = gsum (EAGLE recurrence: pre-norm residual)

Structs§

Eagle3Draft
The EAGLE3 draft model: encoder fc + ONE Llama-style decoder layer + untied lm_head + d2t. All weights are bf16 -> dequant to f32 GpuTensor::Float (the draft is ~0.8 GB; the matmuls go through cuBLASLt linear). The draft attention is PLAIN Llama (no QK-norm, no output gate), distinct from the trunk’s gated/QK-normed full-attn.
Eagle3Scratch
Tiny scratch KV for the EAGLE3 draft layer (one full-attn layer). Reset each draft round. Uses the SAME q8_0-K / q5_1-V quantized layout as the trunk KV (head_dim%32==0 holds: 256).