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 mut target = main_model.new_context(&backend, LlamaContextParams::default())?;
33//!
34//! // Draft: a SEPARATE EAGLE-3 draft model, also a default context.
35//! let mut 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(&mut target, &mut 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//! session.decode_target_and_process(&mut batch)?;
50//! let drafts = session.draft(0, n_past, last_token)?;
51//! // verify `drafts` against the target, count acceptances ...
52//! session.accept(0, n_accepted)?;
53//! ```
54//!
55//! # Hidden-state extraction
56//!
57//! EAGLE-3 needs the target model to expose internal hidden states. The
58//! session configures the required extraction on both contexts at construction
59//! time; [`need_embd`](Eagle3Session::need_embd) and
60//! [`need_embd_pre_norm`](Eagle3Session::need_embd_pre_norm) report which kind
61//! the active backend requested (rarely needed by callers).
62
63use std::marker::PhantomData;
64use std::ptr::NonNull;
65use std::rc::Rc;
66
67use crate::context::params::LlamaContextType;
68use crate::context::LlamaContext;
69use crate::llama_batch::LlamaBatch;
70use crate::speculative::MAX_SPECULATIVE_PROMPT_TOKENS;
71use crate::speculative::{
72    capture_state, restore_state, validate_config, validate_context_capacities,
73    SpeculativeContextCapacity, SpeculativeStateError,
74};
75use crate::token::LlamaToken;
76
77/// Errors raised by the EAGLE-3 draft session.
78#[derive(Debug, thiserror::Error)]
79pub enum Eagle3SessionError {
80    /// Returned when session init fails. The most common cause is that `draft`
81    /// was not built from a valid EAGLE-3 draft model (upstream expects a draft
82    /// model exposing exactly 3 target-extract layers), or that one of the
83    /// contexts is incompatible.
84    #[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")]
85    Init,
86
87    /// `process` returned false on the underlying speculative context.
88    #[error("EAGLE-3 process failed (see llama.cpp logs)")]
89    Process,
90
91    /// Native prompt initialization failed or raised a contained exception.
92    #[error("EAGLE-3 begin failed")]
93    Begin,
94
95    /// Native draft generation failed or raised a contained exception.
96    #[error("EAGLE-3 draft failed")]
97    Draft,
98
99    /// Native proposal acceptance failed or raised a contained exception.
100    #[error("EAGLE-3 accept failed")]
101    Accept,
102
103    /// Prompt storage exceeds the safe speculative-session bound.
104    #[error("prompt has {size} tokens, exceeding the {maximum}-token bound")]
105    PromptTooLong {
106        /// Caller-supplied prompt-token count.
107        size: usize,
108        /// Inclusive safe prompt-token bound.
109        maximum: usize,
110    },
111
112    /// The supplied contexts do not satisfy the native EAGLE-3 contract.
113    #[error("incompatible EAGLE-3 contexts: {0}")]
114    IncompatibleContexts(&'static str),
115
116    /// Caller passed a sequence id outside `[0, n_seq)`.
117    #[error("sequence id {seq_id} out of range (n_seq = {n_seq})")]
118    BadSeqId {
119        /// the offending seq id
120        seq_id: i32,
121        /// configured number of sequences
122        n_seq: u32,
123    },
124
125    /// Invalid session configuration (e.g. `n_draft_max <= 0`).
126    #[error("invalid EAGLE-3 session config: {0}")]
127    InvalidConfig(&'static str),
128
129    /// The target context failed to decode.
130    #[error("target decode failed: {0}")]
131    Decode(#[from] crate::DecodeError),
132
133    /// An operation requires all draft proposals to be completed first.
134    #[error("sequence {seq_id} still has an unaccepted draft proposal")]
135    ProposalPending {
136        /// Sequence with a pending proposal.
137        seq_id: i32,
138    },
139
140    /// `accept` was called without a preceding nonempty draft.
141    #[error("sequence {seq_id} has no draft proposal to accept")]
142    NoPendingProposal {
143        /// Sequence without a pending proposal.
144        seq_id: i32,
145    },
146
147    /// The accepted prefix exceeds the proposal length.
148    #[error("accepted {accepted} tokens from a {proposed}-token proposal")]
149    AcceptedTooMany {
150        /// Accepted prefix length.
151        accepted: u16,
152        /// Exact proposal length.
153        proposed: usize,
154    },
155
156    /// Exact speculative-state capture or restore failed.
157    #[error(transparent)]
158    State(#[from] SpeculativeStateError),
159}
160
161/// Parameters for [`Eagle3Session::new_with_config`].
162///
163/// Maps directly to upstream `common_params_speculative_draft`.
164#[derive(Debug, Clone, Copy, PartialEq)]
165pub struct Eagle3SessionConfig {
166    /// Number of concurrent sequences (usually `1`).
167    pub n_seq: u32,
168    /// Maximum tokens drafted per [`Eagle3Session::draft`] call (`n_max` upstream).
169    pub n_draft_max: i32,
170    /// Minimum draft tokens to propose (`n_min` upstream, default `0`).
171    pub n_min: i32,
172    /// Greedy probability floor; drafts below this are dropped (`p_min` upstream, default `0.0`).
173    pub p_min: f32,
174}
175
176impl Eagle3SessionConfig {
177    /// Build a config with upstream-aligned defaults for `n_min` (`0`) and
178    /// `p_min` (`0.0`).
179    #[must_use]
180    pub fn new(n_seq: u32, n_draft_max: i32) -> Self {
181        Self {
182            n_seq,
183            n_draft_max,
184            n_min: 0,
185            p_min: 0.0,
186        }
187    }
188
189    /// Set minimum draft tokens (`n_min` upstream).
190    #[must_use]
191    pub fn with_n_min(mut self, n_min: i32) -> Self {
192        self.n_min = n_min;
193        self
194    }
195
196    /// Set draft probability floor (`p_min` upstream).
197    ///
198    /// Draft tokens whose greedy probability falls below this value are dropped.
199    #[must_use]
200    pub fn with_p_min(mut self, p_min: f32) -> Self {
201        self.p_min = p_min;
202        self
203    }
204}
205
206/// Owned EAGLE-3 draft session.
207///
208/// Drops the underlying speculative context when freed.
209///
210/// Both contexts are exclusively borrowed for the session lifetime. The
211/// wrapper retains no manually enforced lifetime and is neither `Send` nor
212/// `Sync`.
213pub struct Eagle3Session<'ctx, 'target_model, 'draft_model> {
214    raw: NonNull<llama_cpp_sys_4::mtp_session>,
215    config: Eagle3SessionConfig,
216    target: &'ctx mut LlamaContext<'target_model>,
217    draft: &'ctx mut LlamaContext<'draft_model>,
218    pending_proposals: Vec<Option<usize>>,
219    not_send_sync: PhantomData<Rc<()>>,
220}
221
222impl<'ctx, 'target_model, 'draft_model> Eagle3Session<'ctx, 'target_model, 'draft_model> {
223    /// Construct an EAGLE-3 draft session with upstream defaults for `n_min`
224    /// and `p_min`.
225    ///
226    /// Equivalent to `new_with_config(target, draft, Eagle3SessionConfig::new(n_seq, n_draft_max))`.
227    ///
228    /// # Errors
229    ///
230    /// Returns [`Eagle3SessionError::Init`] or [`Eagle3SessionError::InvalidConfig`].
231    pub fn new(
232        target: &'ctx mut LlamaContext<'target_model>,
233        draft: &'ctx mut LlamaContext<'draft_model>,
234        n_seq: u32,
235        n_draft_max: i32,
236    ) -> Result<Self, Eagle3SessionError> {
237        Self::new_with_config(target, draft, Eagle3SessionConfig::new(n_seq, n_draft_max))
238    }
239
240    /// Construct an EAGLE-3 draft session with full speculative draft
241    /// parameters.
242    ///
243    /// `target` must be a
244    /// [`LlamaContextType::Default`](crate::context::params::LlamaContextType::Default)
245    /// context over the main model. `draft` must be a `Default` context over a
246    /// **separate EAGLE-3 draft model** trained against that target.
247    ///
248    /// # Errors
249    ///
250    /// Returns [`Eagle3SessionError::Init`] (e.g. the draft model is not a
251    /// valid EAGLE-3 model) or [`Eagle3SessionError::InvalidConfig`].
252    pub fn new_with_config(
253        target: &'ctx mut LlamaContext<'target_model>,
254        draft: &'ctx mut LlamaContext<'draft_model>,
255        config: Eagle3SessionConfig,
256    ) -> Result<Self, Eagle3SessionError> {
257        validate_config(config.n_seq, config.n_draft_max, config.n_min, config.p_min)
258            .map_err(Eagle3SessionError::InvalidConfig)?;
259        validate_contexts(target, draft, config)?;
260        let sequence_slots = usize::try_from(config.n_seq)
261            .map_err(|_| Eagle3SessionError::InvalidConfig("n_seq exceeds usize"))?;
262
263        // `MTP_SPEC_TYPE_*` is `c_uint` under clang/gcc and `c_int` under MSVC;
264        // `as i32` compiles on both. The allow covers the clang/gcc case.
265        #[allow(clippy::cast_possible_wrap)]
266        let c_config = llama_cpp_sys_4::mtp_session_config {
267            n_seq: config.n_seq,
268            n_draft_max: config.n_draft_max,
269            n_min: config.n_min,
270            p_min: config.p_min,
271            spec_type: llama_cpp_sys_4::MTP_SPEC_TYPE_EAGLE3 as i32,
272        };
273
274        let raw = unsafe {
275            llama_cpp_sys_4::mtp_session_new(
276                target.context.as_ptr(),
277                draft.context.as_ptr(),
278                &raw const c_config,
279            )
280        };
281        let raw = NonNull::new(raw).ok_or(Eagle3SessionError::Init)?;
282        Ok(Self {
283            raw,
284            config,
285            target,
286            draft,
287            pending_proposals: vec![None; sequence_slots],
288            not_send_sync: PhantomData,
289        })
290    }
291
292    /// Session configuration passed at construction.
293    #[must_use]
294    pub fn config(&self) -> Eagle3SessionConfig {
295        self.config
296    }
297
298    /// True when the speculative backend needs post-norm embeddings on the
299    /// target context (`llama_set_embeddings`).
300    #[must_use]
301    pub fn need_embd(&self) -> bool {
302        unsafe { llama_cpp_sys_4::mtp_session_need_embd(self.raw.as_ptr()) }
303    }
304
305    /// True when the speculative backend needs pre-norm hidden states on the
306    /// target context (`llama_set_embeddings_pre_norm`).
307    ///
308    /// Configured automatically during session init; callers normally do not
309    /// need to set it manually.
310    #[must_use]
311    pub fn need_embd_pre_norm(&self) -> bool {
312        unsafe { llama_cpp_sys_4::mtp_session_need_embd_pre_norm(self.raw.as_ptr()) }
313    }
314
315    /// Configured maximum number of tokens drafted per [`draft`](Self::draft) call.
316    #[must_use]
317    pub fn n_draft_max(&self) -> i32 {
318        self.config.n_draft_max
319    }
320
321    /// Configured minimum draft tokens (`n_min`).
322    #[must_use]
323    pub fn n_min(&self) -> i32 {
324        self.config.n_min
325    }
326
327    /// Configured draft probability floor (`p_min`).
328    #[must_use]
329    pub fn p_min(&self) -> f32 {
330        self.config.p_min
331    }
332
333    /// Configured number of sequences.
334    #[must_use]
335    pub fn n_seq(&self) -> u32 {
336        self.config.n_seq
337    }
338
339    /// Returns shared access to the target context for reading logits,
340    /// embeddings, and model metadata.
341    #[must_use]
342    pub fn target_context(&self) -> &LlamaContext<'target_model> {
343        self.target
344    }
345
346    /// Returns exclusive access to the target context while this wrapper
347    /// retains native pointer ownership.
348    #[must_use]
349    pub fn target_context_mut(&mut self) -> &mut LlamaContext<'target_model> {
350        self.target
351    }
352
353    /// Returns shared access to the draft context for metadata inspection.
354    #[must_use]
355    pub fn draft_context(&self) -> &LlamaContext<'draft_model> {
356        self.draft
357    }
358
359    /// Returns exclusive access to the draft context while this wrapper
360    /// retains native pointer ownership.
361    #[must_use]
362    pub fn draft_context_mut(&mut self) -> &mut LlamaContext<'draft_model> {
363        self.draft
364    }
365
366    /// Decodes on the target and immediately harvests the same batch into
367    /// EAGLE-3.
368    ///
369    /// # Errors
370    ///
371    /// Returns a target [`crate::DecodeError`] or native process failure.
372    pub fn decode_target_and_process(
373        &mut self,
374        batch: &mut LlamaBatch,
375    ) -> Result<(), Eagle3SessionError> {
376        self.decode_target(batch)?;
377        self.process(batch)
378    }
379
380    /// Decodes one batch on the exclusively held target context.
381    ///
382    /// Use [`Self::decode_target_and_process`] unless mechanics must run
383    /// between target decode and draft-state harvesting. This method remains
384    /// available while a draft proposal is pending because that is the target
385    /// verification phase; proposal creation, begin, and state access retain
386    /// their stricter lifecycle checks.
387    ///
388    /// # Errors
389    ///
390    /// Returns a target [`crate::DecodeError`].
391    pub fn decode_target(&mut self, batch: &mut LlamaBatch) -> Result<(), Eagle3SessionError> {
392        self.target.decode(batch)?;
393        Ok(())
394    }
395
396    /// Log speculative-decoding statistics (draft/accept counts and timings)
397    /// via llama.cpp `LOG_INF`. Install a log callback with [`crate::log_set`]
398    /// to capture output.
399    pub fn print_stats(&self) {
400        unsafe { llama_cpp_sys_4::mtp_session_print_stats(self.raw.as_ptr()) }
401    }
402
403    /// Optional: call once at the start of a fresh generation with the prompt
404    /// tokens that were just decoded into the target context.
405    ///
406    /// # Errors
407    ///
408    /// Returns [`Eagle3SessionError::BadSeqId`] if `seq_id` is out of range.
409    pub fn begin(&mut self, seq_id: i32, prompt: &[LlamaToken]) -> Result<(), Eagle3SessionError> {
410        self.check_seq(seq_id)?;
411        self.require_quiescent()?;
412        if prompt.len() > MAX_SPECULATIVE_PROMPT_TOKENS {
413            return Err(Eagle3SessionError::PromptTooLong {
414                size: prompt.len(),
415                maximum: MAX_SPECULATIVE_PROMPT_TOKENS,
416            });
417        }
418        let ok = unsafe {
419            llama_cpp_sys_4::mtp_session_begin(
420                self.raw.as_ptr(),
421                seq_id,
422                prompt.as_ptr().cast(),
423                prompt.len(),
424            )
425        };
426        if !ok {
427            return Err(Eagle3SessionError::Begin);
428        }
429        Ok(())
430    }
431
432    /// Hand the session a batch that was just decoded on the target context.
433    ///
434    /// Call this after every successful `target.decode(batch)` so upstream can
435    /// harvest the target hidden states EAGLE-3 drafts from.
436    ///
437    /// # Errors
438    ///
439    /// Returns [`Eagle3SessionError::Process`] if the underlying call fails.
440    pub fn process(&mut self, batch: &LlamaBatch) -> Result<(), Eagle3SessionError> {
441        let ok = unsafe {
442            llama_cpp_sys_4::mtp_session_process(self.raw.as_ptr(), &raw const batch.llama_batch)
443        };
444        if ok {
445            Ok(())
446        } else {
447            Err(Eagle3SessionError::Process)
448        }
449    }
450
451    /// Generate up to [`n_draft_max`](Self::n_draft_max) speculative tokens.
452    ///
453    /// `n_past` is the number of tokens already in the target KV cache for
454    /// `seq_id`. `id_last` is the last token accepted on the target (usually
455    /// the token you just sampled).
456    ///
457    /// # Errors
458    ///
459    /// Returns [`Eagle3SessionError::BadSeqId`] if `seq_id` is out of range.
460    pub fn draft(
461        &mut self,
462        seq_id: i32,
463        n_past: i32,
464        id_last: LlamaToken,
465    ) -> Result<Vec<LlamaToken>, Eagle3SessionError> {
466        self.check_seq(seq_id)?;
467        let sequence_index = self.sequence_index(seq_id)?;
468        if self.pending_proposals[sequence_index].is_some() {
469            return Err(Eagle3SessionError::ProposalPending { seq_id });
470        }
471
472        let cap = usize::try_from(self.config.n_draft_max.max(0)).unwrap_or(0);
473        let mut buf: Vec<i32> = vec![0; cap];
474        let mut out_n = i32::try_from(cap).unwrap_or(i32::MAX);
475
476        let ok = unsafe {
477            llama_cpp_sys_4::mtp_session_draft(
478                self.raw.as_ptr(),
479                seq_id,
480                n_past,
481                id_last.0,
482                buf.as_mut_ptr(),
483                &raw mut out_n,
484            )
485        };
486        if !ok {
487            return Err(Eagle3SessionError::Draft);
488        }
489
490        let n = usize::try_from(out_n.max(0)).unwrap_or(0);
491        buf.truncate(n);
492        if n > 0 {
493            self.pending_proposals[sequence_index] = Some(n);
494        }
495        Ok(buf.into_iter().map(LlamaToken).collect())
496    }
497
498    /// Inform the session how many draft tokens the target verifier accepted.
499    ///
500    /// Pass `0` when every draft was rejected.
501    ///
502    /// # Errors
503    ///
504    /// Returns [`Eagle3SessionError::BadSeqId`] if `seq_id` is out of range.
505    pub fn accept(&mut self, seq_id: i32, n_accepted: u16) -> Result<(), Eagle3SessionError> {
506        self.check_seq(seq_id)?;
507        let sequence_index = self.sequence_index(seq_id)?;
508        let proposed = self.pending_proposals[sequence_index]
509            .ok_or(Eagle3SessionError::NoPendingProposal { seq_id })?;
510        if usize::from(n_accepted) > proposed {
511            return Err(Eagle3SessionError::AcceptedTooMany {
512                accepted: n_accepted,
513                proposed,
514            });
515        }
516        let ok =
517            unsafe { llama_cpp_sys_4::mtp_session_accept(self.raw.as_ptr(), seq_id, n_accepted) };
518        if !ok {
519            return Err(Eagle3SessionError::Accept);
520        }
521        self.pending_proposals[sequence_index] = None;
522        Ok(())
523    }
524
525    /// Returns `true` when every draft proposal has been completed.
526    #[must_use]
527    pub fn is_quiescent(&self) -> bool {
528        self.pending_proposals.iter().all(Option::is_none)
529            && unsafe { llama_cpp_sys_4::mtp_session_is_quiescent(self.raw.as_ptr()) }
530    }
531
532    /// Captures versioned per-sequence speculative continuation state.
533    ///
534    /// Target and draft context bytes are separate and must be checkpointed at
535    /// the same quiescent boundary.
536    ///
537    /// # Errors
538    ///
539    /// Returns an error for an invalid sequence, pending proposal, incomplete
540    /// native support, or excessive state.
541    pub fn speculative_state(&self, seq_id: i32) -> Result<Vec<u8>, Eagle3SessionError> {
542        self.check_seq(seq_id)?;
543        self.require_quiescent()?;
544        Ok(capture_state(self.raw, seq_id)?)
545    }
546
547    /// Restores versioned per-sequence speculative continuation state.
548    ///
549    /// Restore the corresponding target and draft context bytes before calling
550    /// this method.
551    ///
552    /// # Errors
553    ///
554    /// Returns an error for an invalid sequence, pending proposal, excessive
555    /// input, or any version/configuration/state mismatch.
556    pub fn restore_speculative_state(
557        &mut self,
558        seq_id: i32,
559        state: &[u8],
560    ) -> Result<(), Eagle3SessionError> {
561        self.check_seq(seq_id)?;
562        self.require_quiescent()?;
563        restore_state(self.raw, seq_id, state)?;
564        Ok(())
565    }
566
567    /// Removes a target-context KV range.
568    ///
569    /// # Errors
570    ///
571    /// Returns a conversion error when an identifier or position exceeds
572    /// native `i32` bounds.
573    pub fn clear_target_kv_cache_seq(
574        &mut self,
575        seq_id: Option<u32>,
576        p0: Option<u32>,
577        p1: Option<u32>,
578    ) -> Result<bool, crate::context::kv_cache::KvCacheConversionError> {
579        self.target.clear_kv_cache_seq(seq_id, p0, p1)
580    }
581
582    /// Removes a draft-context KV range.
583    ///
584    /// # Errors
585    ///
586    /// Returns a conversion error when an identifier or position exceeds
587    /// native `i32` bounds.
588    pub fn clear_draft_kv_cache_seq(
589        &mut self,
590        seq_id: Option<u32>,
591        p0: Option<u32>,
592        p1: Option<u32>,
593    ) -> Result<bool, crate::context::kv_cache::KvCacheConversionError> {
594        self.draft.clear_kv_cache_seq(seq_id, p0, p1)
595    }
596
597    /// Returns the target context's exact sequence-state byte count.
598    pub fn target_state_seq_get_size_ext(&mut self, seq_id: i32, flags: u32) -> usize {
599        self.target.state_seq_get_size_ext(seq_id, flags)
600    }
601
602    /// Copies target context sequence state with exact native flags.
603    pub fn target_state_seq_get_data_ext(
604        &mut self,
605        dst: &mut [u8],
606        seq_id: i32,
607        flags: u32,
608    ) -> usize {
609        self.target.state_seq_get_data_ext(dst, seq_id, flags)
610    }
611
612    /// Restores target context sequence state with exact native flags.
613    pub fn target_state_seq_set_data_ext(&mut self, src: &[u8], seq_id: i32, flags: u32) -> usize {
614        self.target.state_seq_set_data_ext(src, seq_id, flags)
615    }
616
617    /// Returns the draft context's exact sequence-state byte count.
618    pub fn draft_state_seq_get_size_ext(&mut self, seq_id: i32, flags: u32) -> usize {
619        self.draft.state_seq_get_size_ext(seq_id, flags)
620    }
621
622    /// Copies draft context sequence state with exact native flags.
623    pub fn draft_state_seq_get_data_ext(
624        &mut self,
625        dst: &mut [u8],
626        seq_id: i32,
627        flags: u32,
628    ) -> usize {
629        self.draft.state_seq_get_data_ext(dst, seq_id, flags)
630    }
631
632    /// Restores draft context sequence state with exact native flags.
633    pub fn draft_state_seq_set_data_ext(&mut self, src: &[u8], seq_id: i32, flags: u32) -> usize {
634        self.draft.state_seq_set_data_ext(src, seq_id, flags)
635    }
636
637    fn require_quiescent(&self) -> Result<(), Eagle3SessionError> {
638        if let Some((index, _)) = self
639            .pending_proposals
640            .iter()
641            .enumerate()
642            .find(|(_, proposal)| proposal.is_some())
643        {
644            return Err(Eagle3SessionError::ProposalPending {
645                seq_id: i32::try_from(index).unwrap_or(i32::MAX),
646            });
647        }
648        if !unsafe { llama_cpp_sys_4::mtp_session_is_quiescent(self.raw.as_ptr()) } {
649            return Err(Eagle3SessionError::State(
650                SpeculativeStateError::NotQuiescent,
651            ));
652        }
653        Ok(())
654    }
655
656    fn check_seq(&self, seq_id: i32) -> Result<(), Eagle3SessionError> {
657        if seq_id < 0 || seq_id.cast_unsigned() >= self.config.n_seq {
658            return Err(Eagle3SessionError::BadSeqId {
659                seq_id,
660                n_seq: self.config.n_seq,
661            });
662        }
663        Ok(())
664    }
665
666    fn sequence_index(&self, seq_id: i32) -> Result<usize, Eagle3SessionError> {
667        self.check_seq(seq_id)?;
668        usize::try_from(seq_id)
669            .map_err(|_| Eagle3SessionError::InvalidConfig("sequence id exceeds usize"))
670    }
671}
672
673fn validate_contexts(
674    target: &LlamaContext<'_>,
675    draft: &LlamaContext<'_>,
676    config: Eagle3SessionConfig,
677) -> Result<(), Eagle3SessionError> {
678    if target.context_type() != LlamaContextType::Default
679        || draft.context_type() != LlamaContextType::Default
680    {
681        return Err(Eagle3SessionError::IncompatibleContexts(
682            "target and draft must both be Default contexts",
683        ));
684    }
685    if target.n_seq_max() < config.n_seq || draft.n_seq_max() != config.n_seq {
686        return Err(Eagle3SessionError::IncompatibleContexts(
687            "target sequence capacity is too small or draft capacity differs from n_seq",
688        ));
689    }
690    let required_draft = u32::try_from(config.n_draft_max)
691        .map_err(|_| Eagle3SessionError::InvalidConfig("n_draft_max exceeds u32"))?;
692    validate_context_capacities(
693        SpeculativeContextCapacity {
694            batch: target.n_batch(),
695            micro_batch: target.n_ubatch(),
696            recurrent_slots: target.n_rs_seq(),
697            recurrent_or_hybrid: target.model.is_recurrent() || target.model.is_hybrid(),
698        },
699        SpeculativeContextCapacity {
700            batch: draft.n_batch(),
701            micro_batch: draft.n_ubatch(),
702            recurrent_slots: draft.n_rs_seq(),
703            recurrent_or_hybrid: draft.model.is_recurrent() || draft.model.is_hybrid(),
704        },
705        required_draft,
706    )
707    .map_err(Eagle3SessionError::IncompatibleContexts)?;
708    let target_layers = target.model.n_layer();
709    let target_architecture = target
710        .model
711        .meta_val_str("general.architecture", 64)
712        .map_err(|_| {
713            Eagle3SessionError::IncompatibleContexts(
714                "target model architecture metadata is unavailable",
715            )
716        })?;
717    if !valid_target_layer_ids(
718        draft.model.target_layer_ids(),
719        target_layers,
720        target_architecture == "gpt-oss",
721    ) {
722        return Err(Eagle3SessionError::IncompatibleContexts(
723            "draft must name exactly three supported target extraction sites",
724        ));
725    }
726    Ok(())
727}
728
729fn valid_target_layer_ids(
730    layer_ids: &[i32],
731    target_layers: i32,
732    terminal_nextn_site: bool,
733) -> bool {
734    target_layers > 0
735        && layer_ids.len() == 3
736        && layer_ids.iter().all(|&layer| {
737            layer >= 0 && (layer < target_layers || (layer == target_layers && terminal_nextn_site))
738        })
739}
740
741impl Drop for Eagle3Session<'_, '_, '_> {
742    fn drop(&mut self) {
743        unsafe { llama_cpp_sys_4::mtp_session_free(self.raw.as_ptr()) }
744    }
745}
746
747impl std::fmt::Debug for Eagle3Session<'_, '_, '_> {
748    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
749        f.debug_struct("Eagle3Session")
750            .field("config", &self.config)
751            .finish_non_exhaustive()
752    }
753}
754
755#[cfg(test)]
756mod tests {
757    use super::valid_target_layer_ids;
758
759    #[test]
760    fn validates_transformer_and_terminal_nextn_sites() {
761        assert!(valid_target_layer_ids(&[1, 4, 7], 8, false));
762        assert!(!valid_target_layer_ids(&[1, 4], 8, false));
763        assert!(!valid_target_layer_ids(&[1, -1, 7], 8, false));
764        assert!(!valid_target_layer_ids(&[1, 4, 8], 8, false));
765        assert!(valid_target_layer_ids(&[1, 4, 8], 8, true));
766        assert!(!valid_target_layer_ids(&[1, 4, 9], 8, true));
767    }
768}
llama_cpp_4/eagle.rs

llama_cpp_4/
eagle.rs
