llama_cpp_4/
mtp.rs

1//! Safe wrapper around the C++ MTP draft session.
2//!
3//! [`MtpSession`] pairs a target [`LlamaContext`] with an MTP draft
4//! [`LlamaContext`] (built with
5//! [`crate::context::params::LlamaContextType::Mtp`]) and drives the
6//! multi-token-prediction speculative-decoding loop introduced in upstream
7//! llama.cpp [PR #22673](https://github.com/ggml-org/llama.cpp/pull/22673).
8//!
9//! The draft algorithm lives in upstream's
10//! `common/speculative.cpp` (`common_speculative_impl_draft_mtp`). This module
11//! wraps it through a stable C shim in `llama-cpp-sys-4/mtp_shim/`.
12//!
13//! # Upstream behaviour (llama.cpp #23269+)
14//!
15//! After [MTP clean-up #23269](https://github.com/ggml-org/llama.cpp/pull/23269):
16//!
17//! - Draft sampling uses `top_k = 10` inside upstream (not configurable from Rust).
18//! - [`MtpSessionConfig::p_min`] filters low-confidence draft tokens (default `0.0`).
19//! - Upstream CLI default for `n_max` is `3`; set [`MtpSessionConfig::n_draft_max`]
20//!   explicitly — optimal values are model/quant dependent ([`MTP.md`] on GitHub).
21//!
22//! [`MTP.md`]: https://github.com/eugenehp/llama-cpp-rs/blob/main/MTP.md
23//!
24//! # Quick start
25//!
26//! ```ignore
27//! use llama_cpp_4::context::params::{LlamaContextParams, LlamaContextType};
28//! use llama_cpp_4::mtp::{MtpSession, MtpSessionConfig};
29//!
30//! let n_draft_max = 3;
31//!
32//! let target = model.new_context(&backend, LlamaContextParams::default())?;
33//! let draft = model.new_context(
34//!     &backend,
35//!     LlamaContextParams::default()
36//!         .with_ctx_type(LlamaContextType::Mtp)
37//!         .with_n_rs_seq(n_draft_max.max(4)),
38//! )?;
39//!
40//! let config = MtpSessionConfig::new(1, n_draft_max).with_p_min(0.0);
41//! let mut session = MtpSession::new_with_config(&target, &draft, config)?;
42//! ```
43//!
44//! # Speculative loop
45//!
46//! For each generation step, after decoding on the **target** context:
47//!
48//! ```ignore
49//! // 1. Target prefill or verify decode (you build the batch)
50//! target.decode(&mut batch)?;
51//!
52//! // 2. Tell MTP about the batch just decoded on the target
53//! session.process(&batch)?;
54//!
55//! // 3. Ask for draft tokens starting from the last accepted token
56//! let drafts = session.draft(0, n_past, last_token)?;
57//!
58//! // 4. Verify drafts on the target (compare logits / sample — your code)
59//! let n_accepted: u16 = /* ... */;
60//!
61//! // 5. Sync draft recurrent state with what the target accepted
62//! session.accept(0, n_accepted)?;
63//! ```
64//!
65//! Call [`MtpSession::begin`] once per fresh generation if you want upstream
66//! prompt tracking (optional for MTP). Call [`MtpSession::print_stats`] when
67//! finished to log draft/accept counters via llama.cpp's log callback.
68//!
69//! A full runnable implementation is in `examples/mtp/`.
70//!
71//! # Embedding requirements
72//!
73//! | Method | MTP typical value | Meaning |
74//! |---|---|---|
75//! | [`MtpSession::need_embd_pre_norm`] | `true` | Next-n hidden states (upstream name) |
76//! | [`MtpSession::need_embd`] | `false` | Post-norm / seq embeddings not used |
77//!
78//! # Multi-head `NextN` (Step3.5+)
79//!
80//! When [`crate::model::LlamaModel::n_layer_nextn`] returns a value greater than `1`, set the
81//! draft context head before each [`MtpSession::draft`] call:
82//!
83//! ```ignore
84//! for head in 0..model.n_layer_nextn() {
85//!     draft.set_nextn_layer_offset(head);
86//!     let drafts = session.draft(0, n_past, last_token)?;
87//!     // verify on target ...
88//! }
89//! draft.set_nextn_layer_offset(0); // restore default
90//! ```
91//!
92
93use std::ptr::NonNull;
94
95use crate::context::LlamaContext;
96use crate::llama_batch::LlamaBatch;
97use crate::token::LlamaToken;
98
99/// Errors raised by the MTP draft session.
100#[derive(Debug, thiserror::Error)]
101pub enum MtpSessionError {
102    /// Returned when `mtp_session_new` fails (typically: model lacks MTP heads,
103    /// or one of the contexts is incompatible).
104    #[error("failed to create MTP draft session — check that ctx_dft was built with LlamaContextType::Mtp and the model has MTP heads")]
105    Init,
106
107    /// `mtp_session_process` returned false.
108    #[error("mtp_session_process failed (see llama.cpp logs)")]
109    Process,
110
111    /// Caller passed a sequence id outside `[0, n_seq)`.
112    #[error("sequence id {seq_id} out of range (n_seq = {n_seq})")]
113    BadSeqId {
114        /// the offending seq id
115        seq_id: i32,
116        /// configured number of sequences
117        n_seq: u32,
118    },
119
120    /// Invalid session configuration (e.g. `n_draft_max <= 0`).
121    #[error("invalid MTP session config: {0}")]
122    InvalidConfig(&'static str),
123}
124
125/// Parameters for [`MtpSession::new_with_config`].
126///
127/// Maps directly to upstream `common_params_speculative_draft`.
128///
129/// # Examples
130///
131/// ```ignore
132/// // Defaults: n_min = 0, p_min = 0.0 (aligned with upstream #23269+)
133/// let cfg = MtpSessionConfig::new(1, 3);
134///
135/// // Stricter drafts: skip tokens below 10% draft-model probability
136/// let cfg = MtpSessionConfig::new(1, 1).with_p_min(0.10);
137/// ```
138#[derive(Debug, Clone, Copy, PartialEq)]
139pub struct MtpSessionConfig {
140    /// Number of concurrent sequences (usually `1`).
141    pub n_seq: u32,
142    /// Maximum tokens drafted per [`MtpSession::draft`] call (`n_max` upstream).
143    pub n_draft_max: i32,
144    /// Minimum draft tokens to propose (`n_min` upstream, default `0`).
145    pub n_min: i32,
146    /// Greedy probability floor; drafts below this are dropped (`p_min` upstream, default `0.0`).
147    pub p_min: f32,
148}
149
150impl MtpSessionConfig {
151    /// Build config with upstream-aligned defaults for `n_min` (`0`) and `p_min` (`0.0`).
152    ///
153    /// # Examples
154    ///
155    /// ```ignore
156    /// let cfg = MtpSessionConfig::new(1, 3); // one sequence, up to 3 draft tokens
157    /// ```
158    #[must_use]
159    pub fn new(n_seq: u32, n_draft_max: i32) -> Self {
160        Self {
161            n_seq,
162            n_draft_max,
163            n_min: 0,
164            p_min: 0.0,
165        }
166    }
167
168    /// Set minimum draft tokens (`n_min` upstream).
169    #[must_use]
170    pub fn with_n_min(mut self, n_min: i32) -> Self {
171        self.n_min = n_min;
172        self
173    }
174
175    /// Set draft probability floor (`p_min` upstream).
176    ///
177    /// Draft tokens whose greedy probability falls below this value are dropped.
178    /// Upstream default is `0.0` after #23269 (was `0.75` in older builds).
179    ///
180    /// # Examples
181    ///
182    /// ```ignore
183    /// let cfg = MtpSessionConfig::new(1, 1).with_p_min(0.10);
184    /// ```
185    #[must_use]
186    pub fn with_p_min(mut self, p_min: f32) -> Self {
187        self.p_min = p_min;
188        self
189    }
190}
191
192/// Owned MTP draft session.
193///
194/// Drops the underlying `mtp_session *` (and the C++ `common_speculative *`
195/// it holds) when freed.
196///
197/// # Lifetime contract (manual)
198///
199/// The session holds raw pointers to both the target and draft
200/// [`LlamaContext`]s. **The caller must keep both contexts alive (i.e. not
201/// drop them) for as long as the session exists.**
202pub struct MtpSession {
203    raw: NonNull<llama_cpp_sys_4::mtp_session>,
204    config: MtpSessionConfig,
205}
206
207// SAFETY: the underlying C++ session owns its own state and is not tied to
208// any TLS. Concurrent calls from multiple threads are NOT safe.
209unsafe impl Send for MtpSession {}
210
211impl MtpSession {
212    /// Construct an MTP draft session with upstream defaults for `n_min` and
213    /// `p_min`.
214    ///
215    /// Equivalent to `new_with_config(MtpSessionConfig::new(n_seq, n_draft_max))`.
216    ///
217    /// # Examples
218    ///
219    /// ```ignore
220    /// let mut session = MtpSession::new(&target, &draft, 1, 3)?;
221    /// ```
222    ///
223    /// # Errors
224    ///
225    /// Returns [`MtpSessionError::Init`] or [`MtpSessionError::InvalidConfig`].
226    pub fn new(
227        target: &LlamaContext<'_>,
228        draft: &LlamaContext<'_>,
229        n_seq: u32,
230        n_draft_max: i32,
231    ) -> Result<Self, MtpSessionError> {
232        Self::new_with_config(target, draft, MtpSessionConfig::new(n_seq, n_draft_max))
233    }
234
235    /// Construct an MTP draft session with full speculative draft parameters.
236    ///
237    /// `target` must be a [`LlamaContextType::Default`](crate::context::params::LlamaContextType::Default) context.
238    /// `draft` must be a [`LlamaContextType::Mtp`](crate::context::params::LlamaContextType::Mtp) context from the same model,
239    /// with [`LlamaContextParams::with_n_rs_seq`](crate::context::params::LlamaContextParams::with_n_rs_seq)
240    /// `>= config.n_draft_max`.
241    ///
242    /// # Examples
243    ///
244    /// ```ignore
245    /// let config = MtpSessionConfig::new(1, 1)
246    ///     .with_p_min(0.0); // match upstream default after #23269
247    /// let session = MtpSession::new_with_config(&target, &draft, config)?;
248    /// ```
249    ///
250    /// # Errors
251    ///
252    /// Returns [`MtpSessionError::Init`] or [`MtpSessionError::InvalidConfig`].
253    pub fn new_with_config(
254        target: &LlamaContext<'_>,
255        draft: &LlamaContext<'_>,
256        config: MtpSessionConfig,
257    ) -> Result<Self, MtpSessionError> {
258        if config.n_seq == 0 {
259            return Err(MtpSessionError::InvalidConfig("n_seq must be > 0"));
260        }
261        if config.n_draft_max <= 0 {
262            return Err(MtpSessionError::InvalidConfig("n_draft_max must be > 0"));
263        }
264
265        // `MTP_SPEC_TYPE_*` is `c_uint` under clang/gcc and `c_int` under MSVC;
266        // `as i32` compiles on both. The allow covers the clang/gcc case.
267        #[allow(clippy::cast_possible_wrap)]
268        let c_config = llama_cpp_sys_4::mtp_session_config {
269            n_seq: config.n_seq,
270            n_draft_max: config.n_draft_max,
271            n_min: config.n_min,
272            p_min: config.p_min,
273            spec_type: llama_cpp_sys_4::MTP_SPEC_TYPE_MTP as i32,
274        };
275
276        let raw = unsafe {
277            llama_cpp_sys_4::mtp_session_new(
278                target.context.as_ptr(),
279                draft.context.as_ptr(),
280                &raw const c_config,
281            )
282        };
283        let raw = NonNull::new(raw).ok_or(MtpSessionError::Init)?;
284        Ok(Self { raw, config })
285    }
286
287    /// Session configuration passed at construction.
288    #[must_use]
289    pub fn config(&self) -> MtpSessionConfig {
290        self.config
291    }
292
293    /// True when the speculative backend needs post-norm embeddings on the
294    /// target context (`llama_set_embeddings`).
295    ///
296    /// MTP returns **false**; use [`Self::need_embd_pre_norm`] for MTP.
297    #[must_use]
298    pub fn need_embd(&self) -> bool {
299        unsafe { llama_cpp_sys_4::mtp_session_need_embd(self.raw.as_ptr()) }
300    }
301
302    /// True when the speculative backend needs pre-norm hidden states on the
303    /// target context (`llama_set_embeddings_pre_norm`).
304    ///
305    /// MTP returns **true**. Upstream configures this on both contexts during
306    /// session init; callers normally do not need to set it manually.
307    #[must_use]
308    pub fn need_embd_pre_norm(&self) -> bool {
309        unsafe { llama_cpp_sys_4::mtp_session_need_embd_pre_norm(self.raw.as_ptr()) }
310    }
311
312    /// Configured maximum number of tokens drafted per [`draft`](Self::draft)
313    /// call.
314    #[must_use]
315    pub fn n_draft_max(&self) -> i32 {
316        self.config.n_draft_max
317    }
318
319    /// Configured minimum draft tokens (`n_min`).
320    #[must_use]
321    pub fn n_min(&self) -> i32 {
322        self.config.n_min
323    }
324
325    /// Configured draft probability floor (`p_min`).
326    #[must_use]
327    pub fn p_min(&self) -> f32 {
328        self.config.p_min
329    }
330
331    /// Configured number of sequences.
332    #[must_use]
333    pub fn n_seq(&self) -> u32 {
334        self.config.n_seq
335    }
336
337    /// Log speculative-decoding statistics (draft/accept counts and timings) via
338    /// llama.cpp `LOG_INF`. Install a log callback with [`crate::log_set`] to
339    /// capture output.
340    ///
341    /// # Examples
342    ///
343    /// ```ignore
344    /// // After your generation loop:
345    /// session.print_stats();
346    /// ```
347    pub fn print_stats(&self) {
348        unsafe { llama_cpp_sys_4::mtp_session_print_stats(self.raw.as_ptr()) }
349    }
350
351    /// Optional: call once at the start of a fresh generation with the
352    /// prompt tokens that were just decoded into the target context.
353    ///
354    /// Upstream uses this for prompt tracking; MTP speculative loops often
355    /// work without it if you call [`Self::process`] after every target decode.
356    ///
357    /// # Examples
358    ///
359    /// ```ignore
360    /// session.begin(0, &prompt_tokens)?;
361    /// ```
362    ///
363    /// # Errors
364    ///
365    /// Returns [`MtpSessionError::BadSeqId`] if `seq_id` is out of range.
366    pub fn begin(&mut self, seq_id: i32, prompt: &[LlamaToken]) -> Result<(), MtpSessionError> {
367        self.check_seq(seq_id)?;
368        unsafe {
369            llama_cpp_sys_4::mtp_session_begin(
370                self.raw.as_ptr(),
371                seq_id,
372                prompt.as_ptr().cast(),
373                prompt.len(),
374            );
375        }
376        Ok(())
377    }
378
379    /// Hand the session a batch that was just decoded on the target context.
380    ///
381    /// Call this after every successful `target.decode(batch)` so upstream can
382    /// sync draft recurrent state with the target KV cache.
383    ///
384    /// # Examples
385    ///
386    /// ```ignore
387    /// target.decode(&mut batch)?;
388    /// session.process(&batch)?;
389    /// ```
390    ///
391    /// # Errors
392    ///
393    /// Returns [`MtpSessionError::Process`] when upstream rejects the batch.
394    pub fn process(&mut self, batch: &LlamaBatch) -> Result<(), MtpSessionError> {
395        let ok = unsafe {
396            llama_cpp_sys_4::mtp_session_process(self.raw.as_ptr(), &raw const batch.llama_batch)
397        };
398        if ok {
399            Ok(())
400        } else {
401            Err(MtpSessionError::Process)
402        }
403    }
404
405    /// Generate up to [`n_draft_max`](Self::n_draft_max) speculative tokens.
406    ///
407    /// `n_past` is the number of tokens already in the target KV cache for
408    /// `seq_id`. `id_last` is the last token accepted on the target (usually
409    /// the token you just sampled).
410    ///
411    /// # Examples
412    ///
413    /// ```ignore
414    /// let drafts = session.draft(0, n_past, last_token)?;
415    /// for draft in &drafts {
416    ///     // verify each draft against target logits ...
417    /// }
418    /// ```
419    ///
420    /// # Errors
421    ///
422    /// Returns [`MtpSessionError::BadSeqId`] if `seq_id` is out of range.
423    pub fn draft(
424        &mut self,
425        seq_id: i32,
426        n_past: i32,
427        id_last: LlamaToken,
428    ) -> Result<Vec<LlamaToken>, MtpSessionError> {
429        self.check_seq(seq_id)?;
430
431        let cap = usize::try_from(self.config.n_draft_max.max(0)).unwrap_or(0);
432        let mut buf: Vec<i32> = vec![0; cap];
433        let mut out_n = i32::try_from(cap).unwrap_or(i32::MAX);
434
435        unsafe {
436            llama_cpp_sys_4::mtp_session_draft(
437                self.raw.as_ptr(),
438                seq_id,
439                n_past,
440                id_last.0,
441                buf.as_mut_ptr(),
442                &raw mut out_n,
443            );
444        }
445
446        let n = usize::try_from(out_n.max(0)).unwrap_or(0);
447        buf.truncate(n);
448        Ok(buf.into_iter().map(LlamaToken).collect())
449    }
450
451    /// Inform the session how many draft tokens the target verifier accepted.
452    ///
453    /// Pass `0` when every draft was rejected. Upstream rolls back draft
454    /// recurrent state accordingly.
455    ///
456    /// # Examples
457    ///
458    /// ```ignore
459    /// session.accept(0, n_accepted)?;
460    /// ```
461    ///
462    /// # Errors
463    ///
464    /// Returns [`MtpSessionError::BadSeqId`] if `seq_id` is out of range.
465    pub fn accept(&mut self, seq_id: i32, n_accepted: u16) -> Result<(), MtpSessionError> {
466        self.check_seq(seq_id)?;
467        unsafe {
468            llama_cpp_sys_4::mtp_session_accept(self.raw.as_ptr(), seq_id, n_accepted);
469        }
470        Ok(())
471    }
472
473    fn check_seq(&self, seq_id: i32) -> Result<(), MtpSessionError> {
474        if seq_id < 0 || seq_id.cast_unsigned() >= self.config.n_seq {
475            return Err(MtpSessionError::BadSeqId {
476                seq_id,
477                n_seq: self.config.n_seq,
478            });
479        }
480        Ok(())
481    }
482}
483
484impl Drop for MtpSession {
485    fn drop(&mut self) {
486        unsafe { llama_cpp_sys_4::mtp_session_free(self.raw.as_ptr()) }
487    }
488}
489
490impl std::fmt::Debug for MtpSession {
491    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
492        f.debug_struct("MtpSession")
493            .field("config", &self.config)
494            .field("need_embd_pre_norm", &self.need_embd_pre_norm())
495            .finish_non_exhaustive()
496    }
497}
llama_cpp_4/mtp.rs

llama_cpp_4/
mtp.rs
