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        // `mtp_state_status` is `u32` on Unix but `i32` on MSVC; `as _` coerces
144        // the unknown value to the `u32` error field on every target.
145        unknown => Err(SpeculativeStateError::Unknown(unknown as _)),
146    }
147}
148
149pub(crate) fn validate_config(
150    n_seq: u32,
151    n_draft_max: i32,
152    n_min: i32,
153    p_min: f32,
154) -> Result<(), &'static str> {
155    if n_seq == 0 || n_seq > MAX_SPECULATIVE_SEQUENCES {
156        return Err("n_seq is outside the supported bound");
157    }
158    if !(1..=MAX_SPECULATIVE_DRAFT_TOKENS).contains(&n_draft_max) {
159        return Err("n_draft_max is outside the supported bound");
160    }
161    if n_min < 0 || n_min > n_draft_max {
162        return Err("n_min must be between zero and n_draft_max");
163    }
164    if !p_min.is_finite() || !(0.0..=1.0).contains(&p_min) {
165        return Err("p_min must be finite and between zero and one");
166    }
167    Ok(())
168}
169
170#[derive(Clone, Copy)]
171pub(crate) struct SpeculativeContextCapacity {
172    pub(crate) batch: u32,
173    pub(crate) micro_batch: u32,
174    pub(crate) recurrent_slots: u32,
175    pub(crate) recurrent_or_hybrid: bool,
176}
177
178pub(crate) fn validate_context_capacities(
179    target: SpeculativeContextCapacity,
180    draft: SpeculativeContextCapacity,
181    maximum_draft_tokens: u32,
182) -> Result<(), &'static str> {
183    let required_rows = maximum_draft_tokens
184        .checked_add(1)
185        .ok_or("n_draft_max plus one exceeds u32")?;
186    if target.batch < required_rows || draft.batch < required_rows {
187        return Err("target or draft batch capacity is smaller than n_draft_max plus one");
188    }
189    if (target.recurrent_or_hybrid && target.micro_batch < required_rows)
190        || (draft.recurrent_or_hybrid && draft.micro_batch < required_rows)
191    {
192        return Err("recurrent micro-batch capacity is smaller than n_draft_max plus one");
193    }
194    if (target.recurrent_or_hybrid && target.recurrent_slots < maximum_draft_tokens)
195        || (draft.recurrent_or_hybrid && draft.recurrent_slots < maximum_draft_tokens)
196    {
197        return Err("recurrent context capacity is smaller than n_draft_max");
198    }
199    Ok(())
200}
201
202#[cfg(test)]
203mod tests {
204    use super::*;
205
206    #[test]
207    fn input_state_is_bounded_before_native_access() {
208        assert!(validate_size(MAX_SPECULATIVE_STATE_BYTES).is_ok());
209        assert_eq!(
210            validate_size(MAX_SPECULATIVE_STATE_BYTES + 1),
211            Err(SpeculativeStateError::Excessive {
212                size: MAX_SPECULATIVE_STATE_BYTES + 1,
213                maximum: MAX_SPECULATIVE_STATE_BYTES,
214            })
215        );
216    }
217
218    #[test]
219    fn speculative_configuration_is_bounded_before_allocation() {
220        assert!(validate_config(1, 4, 0, 0.0).is_ok());
221        assert!(validate_config(0, 4, 0, 0.0).is_err());
222        assert!(validate_config(MAX_SPECULATIVE_SEQUENCES + 1, 4, 0, 0.0).is_err());
223        assert!(validate_config(1, MAX_SPECULATIVE_DRAFT_TOKENS + 1, 0, 0.0).is_err());
224        assert!(validate_config(1, 4, 5, 0.0).is_err());
225        assert!(validate_config(1, 4, 0, f32::NAN).is_err());
226        assert!(validate_config(1, 4, 0, 1.1).is_err());
227    }
228
229    #[test]
230    fn prompt_and_state_bounds_are_consistent() {
231        let maximum_prompt_bytes = MAX_SPECULATIVE_PROMPT_TOKENS
232            .checked_mul(std::mem::size_of::<i32>())
233            .unwrap();
234        assert!(maximum_prompt_bytes < MAX_SPECULATIVE_STATE_BYTES);
235    }
236
237    #[test]
238    fn speculative_decode_capacity_is_checked_before_native_allocation() {
239        let transformer = SpeculativeContextCapacity {
240            batch: 4,
241            micro_batch: 1,
242            recurrent_slots: 0,
243            recurrent_or_hybrid: false,
244        };
245        assert!(validate_context_capacities(transformer, transformer, 3).is_ok());
246        assert!(validate_context_capacities(
247            SpeculativeContextCapacity {
248                batch: 3,
249                ..transformer
250            },
251            transformer,
252            3,
253        )
254        .is_err());
255
256        let recurrent = SpeculativeContextCapacity {
257            batch: 4,
258            micro_batch: 4,
259            recurrent_slots: 3,
260            recurrent_or_hybrid: true,
261        };
262        assert!(validate_context_capacities(recurrent, recurrent, 3).is_ok());
263        assert!(validate_context_capacities(
264            SpeculativeContextCapacity {
265                micro_batch: 3,
266                ..recurrent
267            },
268            recurrent,
269            3,
270        )
271        .is_err());
272        assert!(validate_context_capacities(
273            SpeculativeContextCapacity {
274                recurrent_slots: 2,
275                ..recurrent
276            },
277            recurrent,
278            3,
279        )
280        .is_err());
281        assert!(validate_context_capacities(transformer, transformer, u32::MAX).is_err());
282    }
283}
llama_cpp_4/speculative.rs

llama_cpp_4/
speculative.rs
