llama_cpp_4/
eagle.rs

1//! Safe wrapper around the C++ EAGLE-3 draft session.
2//!
3//! [`Eagle3Session`] drives **EAGLE-3** speculative decoding
4//! (`COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3` in upstream llama.cpp). EAGLE-3
5//! pairs a target model with a small, separately-trained **EAGLE-3 draft
6//! model** that predicts the next tokens from hidden states extracted out of
7//! the target model.
8//!
9//! The draft algorithm lives in upstream's `common/speculative.cpp`
10//! (`common_speculative_impl_draft_eagle3`). This module wraps it through the
11//! same stable C shim used for MTP (`llama-cpp-sys-4/mtp_shim/`); the two
12//! techniques share an identical session lifecycle and differ only in how the
13//! draft context is built.
14//!
15//! # EAGLE-3 vs MTP
16//!
17//! | | EAGLE-3 ([`Eagle3Session`]) | MTP ([`crate::mtp::MtpSession`]) |
18//! |---|---|---|
19//! | Draft weights | a **separate** EAGLE-3 draft model | the **same** model as the target |
20//! | Draft context type | [`LlamaContextType::Default`](crate::context::params::LlamaContextType::Default) | [`LlamaContextType::Mtp`](crate::context::params::LlamaContextType::Mtp) |
21//! | Requirement | draft model must expose 3 target-extract layers | target model must have MTP heads |
22//!
23//! # Setup
24//!
25//! ```ignore
26//! use llama_cpp_4::context::params::LlamaContextParams;
27//! use llama_cpp_4::eagle::{Eagle3Session, Eagle3SessionConfig};
28//!
29//! let n_draft_max = 3;
30//!
31//! // Target: the main model, a normal (default) context.
32//! let target = main_model.new_context(&backend, LlamaContextParams::default())?;
33//!
34//! // Draft: a SEPARATE EAGLE-3 draft model, also a default context.
35//! let draft = eagle3_model.new_context(&backend, LlamaContextParams::default())?;
36//!
37//! let config = Eagle3SessionConfig::new(1, n_draft_max);
38//! let mut session = Eagle3Session::new_with_config(&target, &draft, config)?;
39//! ```
40//!
41//! # Speculative loop
42//!
43//! Identical in shape to MTP: after each decode on the **target** context call
44//! [`process`](Eagle3Session::process), then [`draft`](Eagle3Session::draft)
45//! to get candidate tokens, verify them on the target, and report how many
46//! were accepted with [`accept`](Eagle3Session::accept).
47//!
48//! ```ignore
49//! target.decode(&mut batch)?;
50//! session.process(&batch)?;
51//! let drafts = session.draft(0, n_past, last_token)?;
52//! // verify `drafts` against the target, count acceptances ...
53//! session.accept(0, n_accepted)?;
54//! ```
55//!
56//! # Hidden-state extraction
57//!
58//! EAGLE-3 needs the target model to expose internal hidden states. The
59//! session configures the required extraction on both contexts at construction
60//! time; [`need_embd`](Eagle3Session::need_embd) and
61//! [`need_embd_pre_norm`](Eagle3Session::need_embd_pre_norm) report which kind
62//! the active backend requested (rarely needed by callers).
63
64use std::ptr::NonNull;
65
66use crate::context::LlamaContext;
67use crate::llama_batch::LlamaBatch;
68use crate::token::LlamaToken;
69
70/// Errors raised by the EAGLE-3 draft session.
71#[derive(Debug, thiserror::Error)]
72pub enum Eagle3SessionError {
73    /// Returned when session init fails. The most common cause is that `draft`
74    /// was not built from a valid EAGLE-3 draft model (upstream expects a draft
75    /// model exposing exactly 3 target-extract layers), or that one of the
76    /// contexts is incompatible.
77    #[error("failed to create EAGLE-3 draft session — check that `draft` is a context over a valid EAGLE-3 draft model (3 extract layers) built from the same target")]
78    Init,
79
80    /// `process` returned false on the underlying speculative context.
81    #[error("EAGLE-3 process failed (see llama.cpp logs)")]
82    Process,
83
84    /// Caller passed a sequence id outside `[0, n_seq)`.
85    #[error("sequence id {seq_id} out of range (n_seq = {n_seq})")]
86    BadSeqId {
87        /// the offending seq id
88        seq_id: i32,
89        /// configured number of sequences
90        n_seq: u32,
91    },
92
93    /// Invalid session configuration (e.g. `n_draft_max <= 0`).
94    #[error("invalid EAGLE-3 session config: {0}")]
95    InvalidConfig(&'static str),
96}
97
98/// Parameters for [`Eagle3Session::new_with_config`].
99///
100/// Maps directly to upstream `common_params_speculative_draft`.
101#[derive(Debug, Clone, Copy, PartialEq)]
102pub struct Eagle3SessionConfig {
103    /// Number of concurrent sequences (usually `1`).
104    pub n_seq: u32,
105    /// Maximum tokens drafted per [`Eagle3Session::draft`] call (`n_max` upstream).
106    pub n_draft_max: i32,
107    /// Minimum draft tokens to propose (`n_min` upstream, default `0`).
108    pub n_min: i32,
109    /// Greedy probability floor; drafts below this are dropped (`p_min` upstream, default `0.0`).
110    pub p_min: f32,
111}
112
113impl Eagle3SessionConfig {
114    /// Build a config with upstream-aligned defaults for `n_min` (`0`) and
115    /// `p_min` (`0.0`).
116    #[must_use]
117    pub fn new(n_seq: u32, n_draft_max: i32) -> Self {
118        Self {
119            n_seq,
120            n_draft_max,
121            n_min: 0,
122            p_min: 0.0,
123        }
124    }
125
126    /// Set minimum draft tokens (`n_min` upstream).
127    #[must_use]
128    pub fn with_n_min(mut self, n_min: i32) -> Self {
129        self.n_min = n_min;
130        self
131    }
132
133    /// Set draft probability floor (`p_min` upstream).
134    ///
135    /// Draft tokens whose greedy probability falls below this value are dropped.
136    #[must_use]
137    pub fn with_p_min(mut self, p_min: f32) -> Self {
138        self.p_min = p_min;
139        self
140    }
141}
142
143/// Owned EAGLE-3 draft session.
144///
145/// Drops the underlying speculative context when freed.
146///
147/// # Lifetime contract (manual)
148///
149/// The session holds raw pointers to both the target and draft
150/// [`LlamaContext`]s. **The caller must keep both contexts alive (i.e. not
151/// drop them) for as long as the session exists.**
152pub struct Eagle3Session {
153    raw: NonNull<llama_cpp_sys_4::mtp_session>,
154    config: Eagle3SessionConfig,
155}
156
157// SAFETY: the underlying C++ session owns its own state and is not tied to any
158// TLS. Concurrent calls from multiple threads are NOT safe.
159unsafe impl Send for Eagle3Session {}
160
161impl Eagle3Session {
162    /// Construct an EAGLE-3 draft session with upstream defaults for `n_min`
163    /// and `p_min`.
164    ///
165    /// Equivalent to `new_with_config(target, draft, Eagle3SessionConfig::new(n_seq, n_draft_max))`.
166    ///
167    /// # Errors
168    ///
169    /// Returns [`Eagle3SessionError::Init`] or [`Eagle3SessionError::InvalidConfig`].
170    pub fn new(
171        target: &LlamaContext<'_>,
172        draft: &LlamaContext<'_>,
173        n_seq: u32,
174        n_draft_max: i32,
175    ) -> Result<Self, Eagle3SessionError> {
176        Self::new_with_config(target, draft, Eagle3SessionConfig::new(n_seq, n_draft_max))
177    }
178
179    /// Construct an EAGLE-3 draft session with full speculative draft
180    /// parameters.
181    ///
182    /// `target` must be a
183    /// [`LlamaContextType::Default`](crate::context::params::LlamaContextType::Default)
184    /// context over the main model. `draft` must be a `Default` context over a
185    /// **separate EAGLE-3 draft model** trained against that target.
186    ///
187    /// # Errors
188    ///
189    /// Returns [`Eagle3SessionError::Init`] (e.g. the draft model is not a
190    /// valid EAGLE-3 model) or [`Eagle3SessionError::InvalidConfig`].
191    pub fn new_with_config(
192        target: &LlamaContext<'_>,
193        draft: &LlamaContext<'_>,
194        config: Eagle3SessionConfig,
195    ) -> Result<Self, Eagle3SessionError> {
196        if config.n_seq == 0 {
197            return Err(Eagle3SessionError::InvalidConfig("n_seq must be > 0"));
198        }
199        if config.n_draft_max <= 0 {
200            return Err(Eagle3SessionError::InvalidConfig("n_draft_max must be > 0"));
201        }
202
203        // `MTP_SPEC_TYPE_*` is `c_uint` under clang/gcc and `c_int` under MSVC;
204        // `as i32` compiles on both. The allow covers the clang/gcc case.
205        #[allow(clippy::cast_possible_wrap)]
206        let c_config = llama_cpp_sys_4::mtp_session_config {
207            n_seq: config.n_seq,
208            n_draft_max: config.n_draft_max,
209            n_min: config.n_min,
210            p_min: config.p_min,
211            spec_type: llama_cpp_sys_4::MTP_SPEC_TYPE_EAGLE3 as i32,
212        };
213
214        let raw = unsafe {
215            llama_cpp_sys_4::mtp_session_new(
216                target.context.as_ptr(),
217                draft.context.as_ptr(),
218                &raw const c_config,
219            )
220        };
221        let raw = NonNull::new(raw).ok_or(Eagle3SessionError::Init)?;
222        Ok(Self { raw, config })
223    }
224
225    /// Session configuration passed at construction.
226    #[must_use]
227    pub fn config(&self) -> Eagle3SessionConfig {
228        self.config
229    }
230
231    /// True when the speculative backend needs post-norm embeddings on the
232    /// target context (`llama_set_embeddings`).
233    #[must_use]
234    pub fn need_embd(&self) -> bool {
235        unsafe { llama_cpp_sys_4::mtp_session_need_embd(self.raw.as_ptr()) }
236    }
237
238    /// True when the speculative backend needs pre-norm hidden states on the
239    /// target context (`llama_set_embeddings_pre_norm`).
240    ///
241    /// Configured automatically during session init; callers normally do not
242    /// need to set it manually.
243    #[must_use]
244    pub fn need_embd_pre_norm(&self) -> bool {
245        unsafe { llama_cpp_sys_4::mtp_session_need_embd_pre_norm(self.raw.as_ptr()) }
246    }
247
248    /// Configured maximum number of tokens drafted per [`draft`](Self::draft) call.
249    #[must_use]
250    pub fn n_draft_max(&self) -> i32 {
251        self.config.n_draft_max
252    }
253
254    /// Configured minimum draft tokens (`n_min`).
255    #[must_use]
256    pub fn n_min(&self) -> i32 {
257        self.config.n_min
258    }
259
260    /// Configured draft probability floor (`p_min`).
261    #[must_use]
262    pub fn p_min(&self) -> f32 {
263        self.config.p_min
264    }
265
266    /// Configured number of sequences.
267    #[must_use]
268    pub fn n_seq(&self) -> u32 {
269        self.config.n_seq
270    }
271
272    /// Log speculative-decoding statistics (draft/accept counts and timings)
273    /// via llama.cpp `LOG_INF`. Install a log callback with [`crate::log_set`]
274    /// to capture output.
275    pub fn print_stats(&self) {
276        unsafe { llama_cpp_sys_4::mtp_session_print_stats(self.raw.as_ptr()) }
277    }
278
279    /// Optional: call once at the start of a fresh generation with the prompt
280    /// tokens that were just decoded into the target context.
281    ///
282    /// # Errors
283    ///
284    /// Returns [`Eagle3SessionError::BadSeqId`] if `seq_id` is out of range.
285    pub fn begin(&mut self, seq_id: i32, prompt: &[LlamaToken]) -> Result<(), Eagle3SessionError> {
286        self.check_seq(seq_id)?;
287        unsafe {
288            llama_cpp_sys_4::mtp_session_begin(
289                self.raw.as_ptr(),
290                seq_id,
291                prompt.as_ptr().cast(),
292                prompt.len(),
293            );
294        }
295        Ok(())
296    }
297
298    /// Hand the session a batch that was just decoded on the target context.
299    ///
300    /// Call this after every successful `target.decode(batch)` so upstream can
301    /// harvest the target hidden states EAGLE-3 drafts from.
302    ///
303    /// # Errors
304    ///
305    /// Returns [`Eagle3SessionError::Process`] if the underlying call fails.
306    pub fn process(&mut self, batch: &LlamaBatch) -> Result<(), Eagle3SessionError> {
307        let ok = unsafe {
308            llama_cpp_sys_4::mtp_session_process(self.raw.as_ptr(), &raw const batch.llama_batch)
309        };
310        if ok {
311            Ok(())
312        } else {
313            Err(Eagle3SessionError::Process)
314        }
315    }
316
317    /// Generate up to [`n_draft_max`](Self::n_draft_max) speculative tokens.
318    ///
319    /// `n_past` is the number of tokens already in the target KV cache for
320    /// `seq_id`. `id_last` is the last token accepted on the target (usually
321    /// the token you just sampled).
322    ///
323    /// # Errors
324    ///
325    /// Returns [`Eagle3SessionError::BadSeqId`] if `seq_id` is out of range.
326    pub fn draft(
327        &mut self,
328        seq_id: i32,
329        n_past: i32,
330        id_last: LlamaToken,
331    ) -> Result<Vec<LlamaToken>, Eagle3SessionError> {
332        self.check_seq(seq_id)?;
333
334        let cap = usize::try_from(self.config.n_draft_max.max(0)).unwrap_or(0);
335        let mut buf: Vec<i32> = vec![0; cap];
336        let mut out_n = i32::try_from(cap).unwrap_or(i32::MAX);
337
338        unsafe {
339            llama_cpp_sys_4::mtp_session_draft(
340                self.raw.as_ptr(),
341                seq_id,
342                n_past,
343                id_last.0,
344                buf.as_mut_ptr(),
345                &raw mut out_n,
346            );
347        }
348
349        let n = usize::try_from(out_n.max(0)).unwrap_or(0);
350        buf.truncate(n);
351        Ok(buf.into_iter().map(LlamaToken).collect())
352    }
353
354    /// Inform the session how many draft tokens the target verifier accepted.
355    ///
356    /// Pass `0` when every draft was rejected.
357    ///
358    /// # Errors
359    ///
360    /// Returns [`Eagle3SessionError::BadSeqId`] if `seq_id` is out of range.
361    pub fn accept(&mut self, seq_id: i32, n_accepted: u16) -> Result<(), Eagle3SessionError> {
362        self.check_seq(seq_id)?;
363        unsafe {
364            llama_cpp_sys_4::mtp_session_accept(self.raw.as_ptr(), seq_id, n_accepted);
365        }
366        Ok(())
367    }
368
369    fn check_seq(&self, seq_id: i32) -> Result<(), Eagle3SessionError> {
370        if seq_id < 0 || seq_id.cast_unsigned() >= self.config.n_seq {
371            return Err(Eagle3SessionError::BadSeqId {
372                seq_id,
373                n_seq: self.config.n_seq,
374            });
375        }
376        Ok(())
377    }
378}
379
380impl Drop for Eagle3Session {
381    fn drop(&mut self) {
382        unsafe { llama_cpp_sys_4::mtp_session_free(self.raw.as_ptr()) }
383    }
384}
385
386impl std::fmt::Debug for Eagle3Session {
387    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
388        f.debug_struct("Eagle3Session")
389            .field("config", &self.config)
390            .finish_non_exhaustive()
391    }
392}
llama_cpp_4/eagle.rs

llama_cpp_4/
eagle.rs
