ik_llama_cpp_2/speculative.rs
1//! Multi-Token Prediction (MTP / NextN) support.
2//!
3//! ik_llama.cpp's MTP is an **embedded** mechanism: a single model carrying
4//! NextN prediction layers (loaded with [`crate::LlamaModelParams::with_mtp`]),
5//! driven by a low-level op-type state machine plus a `common_speculative` driver.
6//!
7//! * [`MtpOpType`] + [`crate::LlamaContext::set_mtp_op_type`] expose the raw
8//! op-type primitive (available without `common`).
9//! * [`MtpSpeculative`] (requires the **`common`** feature) is the full
10//! draft/accept driver, backed by a C++ glue over ik's `common_speculative_*`
11//! (`ik_llama_rs_mtp_*`). It reproduces the native `--spec-type mtp` flow.
12
13use ik_llama_cpp_sys as sys;
14
15/// MTP decode-graph operation mode (maps to ik's `llama_mtp_op_type`).
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum MtpOpType {
18 /// Normal decode (no MTP graph).
19 None,
20 /// Prompt warmup — populate MTP KV + hidden state (no logits).
21 Warmup,
22 /// Advance MTP state after accepted tokens.
23 UpdateAccepted,
24 /// Generate draft tokens from the NextN head.
25 DraftGen,
26}
27
28impl MtpOpType {
29 /// Raw `llama_mtp_op_type` value.
30 #[must_use]
31 pub fn to_raw(self) -> sys::llama_mtp_op_type {
32 match self {
33 MtpOpType::None => sys::MTP_OP_NONE,
34 MtpOpType::Warmup => sys::MTP_OP_WARMUP,
35 MtpOpType::UpdateAccepted => sys::MTP_OP_UPDATE_ACCEPTED,
36 MtpOpType::DraftGen => sys::MTP_OP_DRAFT_GEN,
37 }
38 }
39}
40
41/// Parameters for MTP speculative decoding. Mirrors the shape of llama-cpp-2's
42/// `MtpSpeculativeParams` (plus `temp`/`mtp_heads` for ik's driver).
43#[derive(Debug, Clone, Copy)]
44pub struct MtpSpeculativeParams {
45 /// Maximum draft tokens proposed per step. Defaults to `2`, the throughput
46 /// sweet spot measured on a 1-NextN-layer model (IQ4_K, ~1.2x over the plain
47 /// engine); higher values draft deeper but acceptance falls off for a
48 /// single-head NextN model. `n_max=1` = one NextN head.
49 pub n_max: i32,
50 /// Minimum draft tokens per step.
51 pub n_min: i32,
52 /// Minimum draft acceptance probability (0.0 = never early-stop drafting).
53 pub p_min: f32,
54 /// MTP heads to use (1 = default; >1 experimental).
55 pub mtp_heads: i32,
56 /// Target-sampler temperature (<= 0 = greedy).
57 pub temp: f32,
58}
59
60impl Default for MtpSpeculativeParams {
61 fn default() -> Self {
62 Self {
63 n_max: 2,
64 n_min: 0,
65 p_min: 0.0,
66 mtp_heads: 1,
67 temp: 0.0,
68 }
69 }
70}
71
72#[cfg(feature = "common")]
73pub use driver::{MtpSpeculative, MtpStep};
74
75#[cfg(feature = "common")]
76mod driver {
77 use super::MtpSpeculativeParams;
78 use ik_llama_cpp_sys as sys;
79 use std::ptr::NonNull;
80
81 use crate::context::LlamaContext;
82 use crate::model::LlamaModel;
83 use crate::token::LlamaToken;
84 use crate::LlamaError;
85
86 /// Result of one [`MtpSpeculative::step`].
87 #[derive(Debug, Clone)]
88 pub struct MtpStep {
89 /// Tokens emitted this step — append to the output; each token is emitted
90 /// exactly once across steps.
91 pub tokens: Vec<LlamaToken>,
92 /// Number of draft tokens accepted this step (0..=n_max).
93 pub n_accepted: i32,
94 }
95
96 /// Full MTP speculative driver (requires the `common` feature).
97 ///
98 /// Wraps ik's `common_speculative_*` via the `ik_llama_rs_mtp_*` C++ glue.
99 /// Two ways to drive it after `new` → `begin(prompt)`:
100 ///
101 /// * **`step()`** — the convenience path: the glue owns the whole
102 /// draft→verify→**sample**→commit cycle (greedy / `temp`), returning the
103 /// committed tokens. Use this when you don't need to constrain sampling.
104 /// * **`draft()` + `commit()`** — the caller-driven path: `draft` returns the
105 /// NextN candidates without sampling; you build the verify batch
106 /// `[id_last] + drafts` (logits on every position), decode it on
107 /// [`target_context_mut`](Self::target_context_mut), pick each committed
108 /// token yourself (e.g. through a [`crate::LlamaGrammar`] + sampler over
109 /// [`target_context`](Self::target_context)'s logits), then `commit` them.
110 /// This is what lets MTP speculation run **together with** grammar-
111 /// constrained / custom sampling — the crate never samples on this path.
112 ///
113 /// Either way ik owns all KV bookkeeping (the companion cache and the target
114 /// rollback happen inside the glue); the caller must not clear or roll back
115 /// either cache itself.
116 #[derive(Debug)]
117 pub struct MtpSpeculative<'model> {
118 // NOTE: `raw` (the C driver) references `ctx`, so it must be freed first.
119 // Fields drop in declaration order, but `Drop::drop` (which frees `raw`)
120 // runs before any field is dropped, so `ctx` is still alive then. Keep
121 // `raw` declared before `ctx` regardless, for clarity.
122 raw: NonNull<sys::ik_llama_rs_mtp>,
123 ctx: LlamaContext<'model>,
124 params: MtpSpeculativeParams,
125 }
126
127 impl<'model> MtpSpeculative<'model> {
128 /// Initialize the driver over a NextN model + its (mtp-enabled) context.
129 ///
130 /// Takes ownership of `ctx` (so the driver is self-contained and can be
131 /// stored in a struct); read/decode it back out via
132 /// [`target_context`](Self::target_context) /
133 /// [`target_context_mut`](Self::target_context_mut).
134 ///
135 /// `model` must have been loaded with `.with_mtp(true)` and `ctx` created
136 /// with `.with_mtp(true)`. Fails with [`LlamaError::MtpInit`] for a model
137 /// with 0 NextN layers (not an MTP model). Recurrent / openPangu targets
138 /// (which includes the Qwen3.5 NextN family) are supported: the driver
139 /// takes an internal rollback checkpoint before each verify so rejected
140 /// drafts restore the recurrent state correctly.
141 pub fn new(
142 model: &LlamaModel,
143 ctx: LlamaContext<'model>,
144 params: MtpSpeculativeParams,
145 ) -> Result<Self, LlamaError> {
146 let cparams = ctx.raw_params; // the exact params ctx was built with
147 // SAFETY: valid model/ctx; cparams matches ctx; glue validates preconditions.
148 let raw = unsafe {
149 sys::ik_llama_rs_mtp_init(
150 model.model.as_ptr(),
151 ctx.context.as_ptr(),
152 &cparams,
153 params.n_max,
154 params.n_min,
155 params.p_min,
156 params.mtp_heads,
157 params.temp,
158 )
159 };
160 NonNull::new(raw)
161 .map(|raw| Self { raw, ctx, params })
162 .ok_or(LlamaError::MtpInit)
163 }
164
165 /// Prompt warmup + begin. Runs the prompt through the MTP path and prepares
166 /// the first draft. Returns `n_past`. Call once before [`Self::step`].
167 pub fn begin(&mut self, prompt: &[LlamaToken]) -> Result<usize, LlamaError> {
168 let toks: Vec<sys::llama_token> = prompt.iter().map(|t| t.0).collect();
169 // SAFETY: valid handle + token buffer.
170 let n =
171 unsafe { sys::ik_llama_rs_mtp_begin(self.raw.as_ptr(), toks.as_ptr(), toks.len()) };
172 if n < 0 {
173 return Err(LlamaError::MtpBegin(n as i32));
174 }
175 Ok(n as usize)
176 }
177
178 /// Run one draft→verify→accept→commit cycle. Returns the tokens emitted
179 /// this step (to append to the output) and how many drafts were accepted.
180 pub fn step(&mut self) -> Result<MtpStep, LlamaError> {
181 let cap = self.params.n_max.max(1) as usize + 2;
182 let mut out = vec![0 as sys::llama_token; cap];
183 let mut n_out: usize = 0;
184 let mut n_accepted: i32 = 0;
185 // SAFETY: out has `cap` slots; n_out/n_accepted are valid out-params.
186 let st = unsafe {
187 sys::ik_llama_rs_mtp_step(
188 self.raw.as_ptr(),
189 out.as_mut_ptr(),
190 cap,
191 &mut n_out,
192 &mut n_accepted,
193 )
194 };
195 if st as i32 != 0 {
196 // LLAMA_RS_STATUS_OK == 0
197 return Err(LlamaError::MtpStep(st as i32));
198 }
199 out.truncate(n_out);
200 Ok(MtpStep {
201 tokens: out.into_iter().map(LlamaToken).collect(),
202 n_accepted,
203 })
204 }
205
206 /// Propose up to `n_max` draft tokens following `id_last` (the last
207 /// committed token) at position `n_past`, **without sampling**.
208 ///
209 /// This is the caller-driven entry point for running MTP speculation
210 /// alongside a grammar / custom sampler (see the module example). After
211 /// `draft`, the caller must:
212 /// 1. build the verify batch `[id_last] + drafts` (one sequence, logits on
213 /// **every** position), in that order;
214 /// 2. decode it on [`target_context_mut`](Self::target_context_mut);
215 /// 3. pick the committed tokens with its own grammar/sampler, reading
216 /// output index `i` via [`target_context`](Self::target_context) for
217 /// each verify position `i`;
218 /// 4. call [`commit`](Self::commit).
219 ///
220 /// Use one round strictly as `draft` → verify-decode → `commit`: `commit`
221 /// takes its anchor position from the `n_past` you pass here, so do not
222 /// `draft` twice before committing, and build the verify batch with
223 /// `id_last` at exactly `n_past`. Do not mix this path with
224 /// [`step`](Self::step) on the same driver — they own the loop state
225 /// differently.
226 ///
227 /// # Errors
228 ///
229 /// [`LlamaError::MtpStep`] if the C glue rejects the draft.
230 pub fn draft(
231 &mut self,
232 n_past: i32,
233 id_last: LlamaToken,
234 ) -> Result<Vec<LlamaToken>, LlamaError> {
235 let cap = self.params.n_max.max(1) as usize;
236 let mut out = vec![0 as sys::llama_token; cap];
237 let mut n_out: usize = 0;
238 // SAFETY: `out` has `cap` slots; `n_out` is a valid out-param.
239 let st = unsafe {
240 sys::ik_llama_rs_mtp_draft(
241 self.raw.as_ptr(),
242 n_past,
243 id_last.0,
244 out.as_mut_ptr(),
245 cap,
246 &mut n_out,
247 )
248 };
249 if st as i32 != 0 {
250 return Err(LlamaError::MtpStep(st as i32));
251 }
252 out.truncate(n_out);
253 Ok(out.into_iter().map(LlamaToken).collect())
254 }
255
256 /// Finalize one speculative round after the caller has decoded the verify
257 /// batch and chosen its committed tokens.
258 ///
259 /// `id_last` is the anchor token that led the verify batch (the same one
260 /// passed to [`draft`](Self::draft)); `committed` is the accepted-draft
261 /// prefix followed by exactly **one** correction/bonus token (so
262 /// `committed.len() == n_accepted + 1`); `n_draft` is the number of tokens
263 /// the matching [`draft`](Self::draft) returned.
264 ///
265 /// Advances the MTP companion by the accepted count and rolls the target
266 /// KV cache back to the committed prefix. The caller must **not** clear or
267 /// roll back either KV cache itself — ik owns that bookkeeping.
268 ///
269 /// # Errors
270 ///
271 /// [`LlamaError::MtpStep`] if `committed` is empty, longer than
272 /// `n_draft + 1` (more tokens than the verify batch produced — which would
273 /// desync the KV cache), or the glue rejects the commit.
274 pub fn commit(
275 &mut self,
276 id_last: LlamaToken,
277 committed: &[LlamaToken],
278 n_draft: usize,
279 ) -> Result<(), LlamaError> {
280 // `committed` is the accepted-draft prefix (<= n_draft) plus exactly
281 // one correction/bonus token, so it can never exceed `n_draft + 1`.
282 // Rejecting a longer slice stops a caller from advancing `n_past` past
283 // positions the verify batch never decoded (a silent KV desync).
284 if committed.is_empty() || committed.len() > n_draft + 1 {
285 return Err(LlamaError::MtpStep(-1));
286 }
287 let ids: Vec<sys::llama_token> = committed.iter().map(|t| t.0).collect();
288 // SAFETY: `ids` is a valid `ids.len()`-long token array for the call.
289 let st = unsafe {
290 sys::ik_llama_rs_mtp_commit(
291 self.raw.as_ptr(),
292 id_last.0,
293 ids.as_ptr(),
294 ids.len(),
295 n_draft,
296 )
297 };
298 if st as i32 != 0 {
299 return Err(LlamaError::MtpStep(st as i32));
300 }
301 Ok(())
302 }
303
304 /// The configured parameters.
305 #[must_use]
306 pub fn params(&self) -> MtpSpeculativeParams {
307 self.params
308 }
309
310 /// Shared access to the target context — read per-position logits here for
311 /// a grammar-gated commit (e.g. [`LlamaContext::token_data_array_ith`]).
312 #[must_use]
313 pub fn target_context(&self) -> &LlamaContext<'model> {
314 &self.ctx
315 }
316
317 /// Mutable access to the target context — decode the verify batch here.
318 pub fn target_context_mut(&mut self) -> &mut LlamaContext<'model> {
319 &mut self.ctx
320 }
321
322 /// Mutable access to the wrapped context (only between full generations).
323 ///
324 /// Alias of [`target_context_mut`](Self::target_context_mut), kept for
325 /// back-compat.
326 pub fn context_mut(&mut self) -> &mut LlamaContext<'model> {
327 &mut self.ctx
328 }
329 }
330
331 impl Drop for MtpSpeculative<'_> {
332 fn drop(&mut self) {
333 // SAFETY: frees the common_speculative + common_sampler owned by the
334 // glue. Runs before `self.ctx` is dropped, so the glue's reference to
335 // the target context is still valid here.
336 unsafe { sys::ik_llama_rs_mtp_free(self.raw.as_ptr()) };
337 }
338 }
339}