pub struct Eagle3Draft {Show 24 fields
pub fc: GpuTensor,
pub input_layernorm: GpuTensor,
pub hidden_norm: GpuTensor,
pub q_proj: GpuTensor,
pub k_proj: GpuTensor,
pub v_proj: GpuTensor,
pub o_proj: GpuTensor,
pub post_attention_layernorm: GpuTensor,
pub gate_proj: GpuTensor,
pub up_proj: GpuTensor,
pub down_proj: GpuTensor,
pub norm: GpuTensor,
pub lm_head: GpuTensor,
pub d2t: Vec<i64>,
pub n_embd: usize,
pub n_head: usize,
pub n_head_kv: usize,
pub head_dim: usize,
pub n_ff: usize,
pub draft_vocab: usize,
pub rope_dim_count: usize,
pub rope_theta: f32,
pub eps: f32,
pub aux_layers: Vec<usize>,
}Expand description
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.
Fields§
§fc: GpuTensor§input_layernorm: GpuTensor§q_proj: GpuTensor§k_proj: GpuTensor§v_proj: GpuTensor§o_proj: GpuTensor§post_attention_layernorm: GpuTensor§gate_proj: GpuTensor§up_proj: GpuTensor§down_proj: GpuTensor§norm: GpuTensor§lm_head: GpuTensor§d2t: Vec<i64>§n_embd: usize§n_head: usize§n_head_kv: usize§head_dim: usize§n_ff: usize§draft_vocab: usize§rope_dim_count: usize§rope_theta: f32§eps: f32§aux_layers: Vec<usize>Implementations§
Source§impl Eagle3Draft
impl Eagle3Draft
Sourcepub fn load(e: &Engine, path: &Path) -> Result<Self, Box<dyn Error>>
pub fn load(e: &Engine, path: &Path) -> Result<Self, Box<dyn Error>>
Load the EAGLE3 draft from a checkpoint directory (config.json + model.safetensors) or a
direct path to the .safetensors. Reads the geometry/rope params from the sibling config.json.
aux_layers is the trunk layer-id list from eagle_config.eagle_aux_hidden_state_layer_ids.
Sourcepub fn d2t_map(&self, draft_id: u32) -> u32
pub fn d2t_map(&self, draft_id: u32) -> u32
Map a DRAFT-vocab id to a TARGET-vocab id (d2t is a DELTA: target = draft + d2t[draft]).
Sourcepub fn encode(
&self,
e: &Engine,
aux: &[CudaSlice<f32>],
) -> Result<CudaSlice<f32>, Box<dyn Error>>
pub fn encode( &self, e: &Engine, aux: &[CudaSlice<f32>], ) -> Result<CudaSlice<f32>, Box<dyn Error>>
ENCODE (once per round, EAGLE-PLAN N3): g = fc @ concat(aux0, aux1, aux2). aux are the 3
trunk residual hiddens of the just-committed token (decode_step_aux / decode_step_t_aux),
in ascending-layer order. Returns the recurrent draft hidden g [n_embd].
Sourcepub fn draft_token(
&self,
e: &Engine,
target: &HybridModel,
prev_tok: u32,
g: &CudaSlice<f32>,
scratch: &mut Eagle3Scratch,
pos: usize,
) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn Error>>
pub fn draft_token( &self, e: &Engine, target: &HybridModel, prev_tok: u32, g: &CudaSlice<f32>, scratch: &mut Eagle3Scratch, pos: usize, ) -> Result<(Vec<f32>, CudaSlice<f32>), Box<dyn Error>>
One DRAFT-token forward (EAGLE-PLAN N4, T=1). prev_tok = the TARGET token id to predict
from (last committed or previous draft). g = the recurrent draft hidden (encode() output
on round entry, then the previous step’s g_next). Returns (draft_logits[draft_vocab] host,
g_next dev). Mirrors the vLLM op-sequence documented at the top of this file.