llama_cpp_4/
speculative.rs

1//! Shared exact-state contracts for speculative draft sessions.
2
3use std::ptr::NonNull;
4
5/// Maximum speculative continuation bytes copied through the safe API.
6pub const MAX_SPECULATIVE_STATE_BYTES: usize = 64 * 1024 * 1024;
7/// Maximum concurrent sequence slots allocated by one speculative session.
8pub const MAX_SPECULATIVE_SEQUENCES: u32 = 4_096;
9/// Maximum tokens requested from one speculative draft boundary.
10pub const MAX_SPECULATIVE_DRAFT_TOKENS: i32 = 4_096;
11/// Maximum prompt tokens copied into one speculative session.
12pub const MAX_SPECULATIVE_PROMPT_TOKENS: usize = 1_048_576;
13
14/// Failure while capturing or restoring versioned speculative state.
15#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
16pub enum SpeculativeStateError {
17    /// A required native session or output pointer was null.
18    #[error("speculative state operation received a null native value")]
19    Null,
20    /// The sequence id is outside the configured session range.
21    #[error("speculative state sequence id is outside the configured range")]
22    BadSequence,
23    /// A draft proposal has not yet been completed with `accept`.
24    #[error("speculative state is available only at a quiescent boundary")]
25    NotQuiescent,
26    /// Native state changed size between the size and copy calls.
27    #[error("speculative state size changed from {expected} to {actual} bytes")]
28    SizeChanged {
29        /// Size returned by the initial query.
30        expected: usize,
31        /// Size required or written by the copy call.
32        actual: usize,
33    },
34    /// Native reported that the supplied buffer was too small without a size.
35    #[error("the speculative state buffer was too small")]
36    BufferTooSmall,
37    /// The active native implementation did not expose complete state.
38    #[error("the active speculative implementation did not expose complete state")]
39    Unavailable,
40    /// The supplied state failed its exact version/configuration checks.
41    #[error("speculative state is invalid for this session")]
42    Invalid,
43    /// Native state size arithmetic overflowed.
44    #[error("speculative state size overflowed")]
45    Overflow,
46    /// A C++ exception was contained by the speculative-state shim.
47    #[error("the native speculative-state operation raised an exception")]
48    Exception,
49    /// State exceeded the safe allocation/input bound.
50    #[error("speculative state is {size} bytes, exceeding the {maximum}-byte bound")]
51    Excessive {
52        /// Requested or supplied byte count.
53        size: usize,
54        /// Inclusive safe bound.
55        maximum: usize,
56    },
57    /// Native returned a status unknown to this binding revision.
58    #[error("unknown speculative state status {0}")]
59    Unknown(u32),
60}
61
62pub(crate) fn capture_state(
63    raw: NonNull<llama_cpp_sys_4::mtp_session>,
64    seq_id: i32,
65) -> Result<Vec<u8>, SpeculativeStateError> {
66    let mut size = 0_usize;
67    // SAFETY: `raw` is owned by the calling safe session and `size` is a live
68    // output for the duration of the synchronous call.
69    let status =
70        unsafe { llama_cpp_sys_4::mtp_session_state_size(raw.as_ptr(), seq_id, &raw mut size) };
71    status_result(status)?;
72    validate_size(size)?;
73
74    let mut state = vec![0_u8; size];
75    let mut written = 0_usize;
76    // SAFETY: the vector has exactly `size` writable bytes and the session
77    // remains exclusively borrowed by the caller.
78    let status = unsafe {
79        llama_cpp_sys_4::mtp_session_state_get(
80            raw.as_ptr(),
81            seq_id,
82            state.as_mut_ptr(),
83            state.len(),
84            &raw mut written,
85        )
86    };
87    if status == llama_cpp_sys_4::MTP_STATE_STATUS_BUFFER_SMALL {
88        return Err(SpeculativeStateError::SizeChanged {
89            expected: size,
90            actual: written,
91        });
92    }
93    status_result(status)?;
94    if written != size {
95        return Err(SpeculativeStateError::SizeChanged {
96            expected: size,
97            actual: written,
98        });
99    }
100    Ok(state)
101}
102
103pub(crate) fn restore_state(
104    raw: NonNull<llama_cpp_sys_4::mtp_session>,
105    seq_id: i32,
106    state: &[u8],
107) -> Result<(), SpeculativeStateError> {
108    validate_size(state.len())?;
109    if state.is_empty() {
110        return Err(SpeculativeStateError::Invalid);
111    }
112    // SAFETY: `state` remains live and immutable for the synchronous call; the
113    // owning safe session provides exclusive native access.
114    let status = unsafe {
115        llama_cpp_sys_4::mtp_session_state_set(raw.as_ptr(), seq_id, state.as_ptr(), state.len())
116    };
117    status_result(status)
118}
119
120fn validate_size(size: usize) -> Result<(), SpeculativeStateError> {
121    if size > MAX_SPECULATIVE_STATE_BYTES {
122        return Err(SpeculativeStateError::Excessive {
123            size,
124            maximum: MAX_SPECULATIVE_STATE_BYTES,
125        });
126    }
127    Ok(())
128}
129
130fn status_result(status: llama_cpp_sys_4::mtp_state_status) -> Result<(), SpeculativeStateError> {
131    match status {
132        llama_cpp_sys_4::MTP_STATE_STATUS_OK => Ok(()),
133        llama_cpp_sys_4::MTP_STATE_STATUS_NULL => Err(SpeculativeStateError::Null),
134        llama_cpp_sys_4::MTP_STATE_STATUS_BAD_SEQUENCE => Err(SpeculativeStateError::BadSequence),
135        llama_cpp_sys_4::MTP_STATE_STATUS_NOT_QUIESCENT => Err(SpeculativeStateError::NotQuiescent),
136        llama_cpp_sys_4::MTP_STATE_STATUS_BUFFER_SMALL => {
137            Err(SpeculativeStateError::BufferTooSmall)
138        }
139        llama_cpp_sys_4::MTP_STATE_STATUS_UNAVAILABLE => Err(SpeculativeStateError::Unavailable),
140        llama_cpp_sys_4::MTP_STATE_STATUS_INVALID => Err(SpeculativeStateError::Invalid),
141        llama_cpp_sys_4::MTP_STATE_STATUS_OVERFLOW => Err(SpeculativeStateError::Overflow),
142        llama_cpp_sys_4::MTP_STATE_STATUS_EXCEPTION => Err(SpeculativeStateError::Exception),
143        unknown => Err(SpeculativeStateError::Unknown(unknown)),
144    }
145}
146
147pub(crate) fn validate_config(
148    n_seq: u32,
149    n_draft_max: i32,
150    n_min: i32,
151    p_min: f32,
152) -> Result<(), &'static str> {
153    if n_seq == 0 || n_seq > MAX_SPECULATIVE_SEQUENCES {
154        return Err("n_seq is outside the supported bound");
155    }
156    if !(1..=MAX_SPECULATIVE_DRAFT_TOKENS).contains(&n_draft_max) {
157        return Err("n_draft_max is outside the supported bound");
158    }
159    if n_min < 0 || n_min > n_draft_max {
160        return Err("n_min must be between zero and n_draft_max");
161    }
162    if !p_min.is_finite() || !(0.0..=1.0).contains(&p_min) {
163        return Err("p_min must be finite and between zero and one");
164    }
165    Ok(())
166}
167
168#[derive(Clone, Copy)]
169pub(crate) struct SpeculativeContextCapacity {
170    pub(crate) batch: u32,
171    pub(crate) micro_batch: u32,
172    pub(crate) recurrent_slots: u32,
173    pub(crate) recurrent_or_hybrid: bool,
174}
175
176pub(crate) fn validate_context_capacities(
177    target: SpeculativeContextCapacity,
178    draft: SpeculativeContextCapacity,
179    maximum_draft_tokens: u32,
180) -> Result<(), &'static str> {
181    let required_rows = maximum_draft_tokens
182        .checked_add(1)
183        .ok_or("n_draft_max plus one exceeds u32")?;
184    if target.batch < required_rows || draft.batch < required_rows {
185        return Err("target or draft batch capacity is smaller than n_draft_max plus one");
186    }
187    if (target.recurrent_or_hybrid && target.micro_batch < required_rows)
188        || (draft.recurrent_or_hybrid && draft.micro_batch < required_rows)
189    {
190        return Err("recurrent micro-batch capacity is smaller than n_draft_max plus one");
191    }
192    if (target.recurrent_or_hybrid && target.recurrent_slots < maximum_draft_tokens)
193        || (draft.recurrent_or_hybrid && draft.recurrent_slots < maximum_draft_tokens)
194    {
195        return Err("recurrent context capacity is smaller than n_draft_max");
196    }
197    Ok(())
198}
199
200#[cfg(test)]
201mod tests {
202    use super::*;
203
204    #[test]
205    fn input_state_is_bounded_before_native_access() {
206        assert!(validate_size(MAX_SPECULATIVE_STATE_BYTES).is_ok());
207        assert_eq!(
208            validate_size(MAX_SPECULATIVE_STATE_BYTES + 1),
209            Err(SpeculativeStateError::Excessive {
210                size: MAX_SPECULATIVE_STATE_BYTES + 1,
211                maximum: MAX_SPECULATIVE_STATE_BYTES,
212            })
213        );
214    }
215
216    #[test]
217    fn speculative_configuration_is_bounded_before_allocation() {
218        assert!(validate_config(1, 4, 0, 0.0).is_ok());
219        assert!(validate_config(0, 4, 0, 0.0).is_err());
220        assert!(validate_config(MAX_SPECULATIVE_SEQUENCES + 1, 4, 0, 0.0).is_err());
221        assert!(validate_config(1, MAX_SPECULATIVE_DRAFT_TOKENS + 1, 0, 0.0).is_err());
222        assert!(validate_config(1, 4, 5, 0.0).is_err());
223        assert!(validate_config(1, 4, 0, f32::NAN).is_err());
224        assert!(validate_config(1, 4, 0, 1.1).is_err());
225    }
226
227    #[test]
228    fn prompt_and_state_bounds_are_consistent() {
229        let maximum_prompt_bytes = MAX_SPECULATIVE_PROMPT_TOKENS
230            .checked_mul(std::mem::size_of::<i32>())
231            .unwrap();
232        assert!(maximum_prompt_bytes < MAX_SPECULATIVE_STATE_BYTES);
233    }
234
235    #[test]
236    fn speculative_decode_capacity_is_checked_before_native_allocation() {
237        let transformer = SpeculativeContextCapacity {
238            batch: 4,
239            micro_batch: 1,
240            recurrent_slots: 0,
241            recurrent_or_hybrid: false,
242        };
243        assert!(validate_context_capacities(transformer, transformer, 3).is_ok());
244        assert!(validate_context_capacities(
245            SpeculativeContextCapacity {
246                batch: 3,
247                ..transformer
248            },
249            transformer,
250            3,
251        )
252        .is_err());
253
254        let recurrent = SpeculativeContextCapacity {
255            batch: 4,
256            micro_batch: 4,
257            recurrent_slots: 3,
258            recurrent_or_hybrid: true,
259        };
260        assert!(validate_context_capacities(recurrent, recurrent, 3).is_ok());
261        assert!(validate_context_capacities(
262            SpeculativeContextCapacity {
263                micro_batch: 3,
264                ..recurrent
265            },
266            recurrent,
267            3,
268        )
269        .is_err());
270        assert!(validate_context_capacities(
271            SpeculativeContextCapacity {
272                recurrent_slots: 2,
273                ..recurrent
274            },
275            recurrent,
276            3,
277        )
278        .is_err());
279        assert!(validate_context_capacities(transformer, transformer, u32::MAX).is_err());
280    }
281}
llama_cpp_4/speculative.rs

llama_cpp_4/
speculative.rs
