Struct MtpSession 

pub struct MtpSession<'ctx, 'model> { /* private fields */ }

Owned MTP draft session.

Drops the underlying mtp_session * (and the C++ common_speculative * it holds) when freed.

The session exclusively borrows both contexts for its Rust lifetime, so neither can be moved, accessed mutably, or dropped while native code retains their pointers. It is deliberately neither Send nor Sync.

Implementations

impl<'ctx, 'model> MtpSession<'ctx, 'model>

pub fn new( target: &'ctx mut LlamaContext<'model>, draft: &'ctx mut LlamaContext<'model>, n_seq: u32, n_draft_max: i32, ) -> Result<Self, MtpSessionError>

Construct an MTP draft session with upstream defaults for n_min and p_min.

Equivalent to new_with_config(MtpSessionConfig::new(n_seq, n_draft_max)).

Examples

let mut session = MtpSession::new(&mut target, &mut draft, 1, 3)?;

Errors

Returns MtpSessionError::Init or MtpSessionError::InvalidConfig.

pub fn new_with_config( target: &'ctx mut LlamaContext<'model>, draft: &'ctx mut LlamaContext<'model>, config: MtpSessionConfig, ) -> Result<Self, MtpSessionError>

Construct an MTP draft session with full speculative draft parameters.

target must be a LlamaContextType::Default context. draft must be a LlamaContextType::Mtp context from the same model, with LlamaContextParams::with_n_rs_seq >= config.n_draft_max.

Examples

let config = MtpSessionConfig::new(1, 1)
    .with_p_min(0.0); // match upstream default after #23269
let session = MtpSession::new_with_config(&mut target, &mut draft, config)?;

Errors

Returns MtpSessionError::Init or MtpSessionError::InvalidConfig.

pub fn config(&self) -> MtpSessionConfig

Session configuration passed at construction.

pub fn need_embd(&self) -> bool

True when the speculative backend needs post-norm embeddings on the target context (llama_set_embeddings).

MTP returns false; use Self::need_embd_pre_norm for MTP.

pub fn need_embd_pre_norm(&self) -> bool

True when the speculative backend needs pre-norm hidden states on the target context (llama_set_embeddings_pre_norm).

MTP returns true. Upstream configures this on both contexts during session init; callers normally do not need to set it manually.

pub fn n_draft_max(&self) -> i32

Configured maximum number of tokens drafted per draft call.

pub fn n_min(&self) -> i32

Configured minimum draft tokens (n_min).

pub fn p_min(&self) -> f32

Configured draft probability floor (p_min).

pub fn n_seq(&self) -> u32

Configured number of sequences.

pub fn target_context(&self) -> &LlamaContext<'model>

Returns shared access to the target context for reading logits, embeddings, and model metadata.

pub fn target_context_mut(&mut self) -> &mut LlamaContext<'model>

Returns exclusive access to the target context while this wrapper retains native pointer ownership.

pub fn draft_context(&self) -> &LlamaContext<'model>

Returns shared access to the draft context for metadata inspection.

pub fn draft_context_mut(&mut self) -> &mut LlamaContext<'model>

Returns exclusive access to the draft context while this wrapper retains native pointer ownership.

pub fn decode_target_and_process( &mut self, batch: &mut LlamaBatch, ) -> Result<(), MtpSessionError>

Decodes on the target and immediately harvests the same batch into MTP.

This is the causal target boundary required by the native draft implementation. A failed decode is never passed to process.

Errors

Returns a target crate::DecodeError or native process failure.

pub fn decode_target( &mut self, batch: &mut LlamaBatch, ) -> Result<(), MtpSessionError>

Decodes one batch on the exclusively held target context.

Use Self::decode_target_and_process unless mechanics must run between target decode and draft-state harvesting. This method remains available while a draft proposal is pending because that is the target verification phase; proposal creation, begin, and state access retain their stricter lifecycle checks.

Errors

Returns a target crate::DecodeError.

pub fn print_stats(&self)

Log speculative-decoding statistics (draft/accept counts and timings) via llama.cpp LOG_INF. Install a log callback with crate::log_set to capture output.

Examples

// After your generation loop:
session.print_stats();

pub fn begin( &mut self, seq_id: i32, prompt: &[LlamaToken], ) -> Result<(), MtpSessionError>

Optional: call once at the start of a fresh generation with the prompt tokens that were just decoded into the target context.

Upstream uses this for prompt tracking; MTP speculative loops often work without it if you call Self::process after every target decode.

Examples

session.begin(0, &prompt_tokens)?;

Errors

Returns MtpSessionError::BadSeqId if seq_id is out of range.

pub fn process(&mut self, batch: &LlamaBatch) -> Result<(), MtpSessionError>

Hand the session a batch that was just decoded on the target context.

Call this after every successful target.decode(batch) so upstream can sync draft recurrent state with the target KV cache.

Examples

target.decode(&mut batch)?;
session.process(&batch)?;

Errors

Returns MtpSessionError::Process when upstream rejects the batch.

pub fn draft( &mut self, seq_id: i32, n_past: i32, id_last: LlamaToken, ) -> Result<Vec<LlamaToken>, MtpSessionError>

Generate up to n_draft_max speculative tokens.

n_past is the number of tokens already in the target KV cache for seq_id. id_last is the last token accepted on the target (usually the token you just sampled).

Examples

let drafts = session.draft(0, n_past, last_token)?;
for draft in &drafts {
    // verify each draft against target logits ...
}

Errors

Returns MtpSessionError::BadSeqId if seq_id is out of range.

pub fn accept( &mut self, seq_id: i32, n_accepted: u16, ) -> Result<(), MtpSessionError>

Inform the session how many draft tokens the target verifier accepted.

Pass 0 when every draft was rejected. Upstream rolls back draft recurrent state accordingly.

Examples

session.accept(0, n_accepted)?;

Errors

Returns MtpSessionError::BadSeqId if seq_id is out of range.

pub fn is_quiescent(&self) -> bool

Returns true when every draft proposal has been completed.

pub fn speculative_state(&self, seq_id: i32) -> Result<Vec<u8>, MtpSessionError>

Captures versioned per-sequence speculative continuation state.

Target and draft context bytes are separate and must be checkpointed at the same quiescent boundary.

Errors

Returns an error for an invalid sequence, pending proposal, incomplete native support, or excessive state.

pub fn restore_speculative_state( &mut self, seq_id: i32, state: &[u8], ) -> Result<(), MtpSessionError>

Restores versioned per-sequence speculative continuation state.

Restore the corresponding target and draft context bytes before calling this method.

Errors

Returns an error for an invalid sequence, pending proposal, excessive input, or any version/configuration/state mismatch.

pub fn clear_target_kv_cache_seq( &mut self, seq_id: Option<u32>, p0: Option<u32>, p1: Option<u32>, ) -> Result<bool, KvCacheConversionError>

Removes a target-context KV range.

Errors

Returns a conversion error when an identifier or position exceeds native i32 bounds.

pub fn clear_draft_kv_cache_seq( &mut self, seq_id: Option<u32>, p0: Option<u32>, p1: Option<u32>, ) -> Result<bool, KvCacheConversionError>

Removes a draft-context KV range.

Errors

Returns a conversion error when an identifier or position exceeds native i32 bounds.

pub fn target_state_seq_get_size_ext( &mut self, seq_id: i32, flags: u32, ) -> usize

Returns the target context’s exact sequence-state byte count.

pub fn target_state_seq_get_data_ext( &mut self, dst: &mut [u8], seq_id: i32, flags: u32, ) -> usize

Copies target context sequence state with exact native flags.

pub fn target_state_seq_set_data_ext( &mut self, src: &[u8], seq_id: i32, flags: u32, ) -> usize

Restores target context sequence state with exact native flags.

pub fn draft_state_seq_get_size_ext(&mut self, seq_id: i32, flags: u32) -> usize

Returns the draft context’s exact sequence-state byte count.

pub fn draft_state_seq_get_data_ext( &mut self, dst: &mut [u8], seq_id: i32, flags: u32, ) -> usize

Copies draft context sequence state with exact native flags.

pub fn draft_state_seq_set_data_ext( &mut self, src: &[u8], seq_id: i32, flags: u32, ) -> usize

Restores draft context sequence state with exact native flags.

Trait Implementations

impl Debug for MtpSession<'_, '_>

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Drop for MtpSession<'_, '_>

fn drop(&mut self)

Executes the destructor for this type.

fn pin_drop(self: Pin<&mut Self>)

🔬This is a nightly-only experimental API. (pin_ergonomics)

Execute the destructor for this type, but different to Drop::drop, it requires self to be pinned.