From f70526df0d45c3659effd9b30e8b2db41700954a Mon Sep 17 00:00:00 2001 From: daemyung-lablup Date: Wed, 30 Sep 2026 13:28:12 +0900 Subject: [PATCH 1/2] feat(speculative): add prompt-lookup decoding to mlxcel generate Replies that edit or reformat the prompt (fix typos, add a JSON field, rename an identifier) repeat long runs of the prompt word for word, and speculative decoding can emit those runs several tokens per target forward without any draft model. mlxcel had no drafter-free speculation: the classic, MTP, and DFlash paths all need a second model or head. Add PromptLookupGenerator in mlxcel-core and wire it to `mlxcel generate --prompt-lookup`. Each round matches the tail of the sequence against earlier context and verifies up to seven proposed tokens in one target forward, trimming the rejected tail off the KV caches. - An incremental n-gram index answers each lookup with one hash probe per n-gram length instead of a scan over the whole context; a test checks it against the reference scan at every context length. - DraftGovernor shortens the block while few tokens land and pauses proposals after three misses in a row, doubling the pause up to 32 rounds. The default minimum n-gram is 2, since one-token matches slowed prose by a quarter. - Rounds without a proposal are pipelined the way plain decoding is: the next step is submitted from the lazy sample before the host reads it. Waiting on each token cost 12% of decode time on an M4 Pro. Pipelining starts after two plain rounds so short gaps inside edits keep their proposals. - The prompt is prefilled through the plain-decode routine, extracted as prefill_prompt_last_logits, so the caches and the first token match plain decoding exactly. KV caches come from make_caches with the configured --kv-cache-mode and Boundary-V policy, --seed is honored, and the repetition-loop guard runs after every emitted token. - Models whose state a cache trim cannot roll back are refused with the reason (recurrent or hybrid layers, model-owned sliding caches, or state that opts out of padded prefill). Multimodal prompts and interactive chat are refused. - The CLI rate line times the whole call, prefill included, as the plain path does; --profile prints the measured prefill/decode split. The acceptance summary goes to stderr. Measured on an M4 Pro, greedy, 400-token cap, whole-call timing: Qwen3-1.7B-4bit function edit 1.93x, JSON edit 1.73x, summary 1.13x, free writing 1.01x over plain decode; Qwen3-4B-4bit 1.73x, 1.70x, 1.52x, 0.95x. A Qwen3-0.6B classic drafter reached at most 0.99x on 1.7B and 1.40x on 4B on the same prompts. Edit and summary replies are byte-identical to plain decoding at greedy; a multi-token verify can still flip a near-tie, the same jitter class as the draft-model path. Validated with cargo test (mlxcel-core 179 including new index, governor, and loop-guard tests; mlxcel commands::generate 89) and clippy with no new warnings. --- src/commands/generate.rs | 77 +- src/commands/generate_tests.rs | 1 + src/commands/run.rs | 1 + src/lib.rs | 4 + src/lib/mlxcel-core/src/generate.rs | 104 ++- src/lib/mlxcel-core/src/speculative/mod.rs | 4 + .../src/speculative/prompt_lookup.rs | 816 ++++++++++++++++++ .../src/speculative/prompt_lookup_tests.rs | 313 +++++++ src/main.rs | 58 ++ 9 files changed, 1337 insertions(+), 41 deletions(-) create mode 100644 src/lib/mlxcel-core/src/speculative/prompt_lookup.rs create mode 100644 src/lib/mlxcel-core/src/speculative/prompt_lookup_tests.rs diff --git a/src/commands/generate.rs b/src/commands/generate.rs index 582b99765..6015463ee 100644 --- a/src/commands/generate.rs +++ b/src/commands/generate.rs @@ -444,7 +444,7 @@ fn validate_pipeline_parallel_args(args: &GenerateArgs) -> Result<()> { // Single-adapter only; multi-adapter stacking and runtime hot-swap // remain out of scope for v1. ensure!( - args.model.draft_model.is_none(), + args.model.draft_model.is_none() && !args.prompt_lookup.prompt_lookup, "CLI pipeline parallelism does not support speculative decoding yet" ); ensure!( @@ -1730,7 +1730,17 @@ pub(super) fn run_generation_mode( // models whose suppressed set is empty. token_bias.suppress_tokens(&model.output_suppressed_token_ids()); - let output = if let Some(ref draft_model_path) = args.model.draft_model { + let output = if args.prompt_lookup.prompt_lookup { + run_prompt_lookup( + model, + args, + prompt_tokens, + sampling_config, + vlm_embeddings.is_some(), + kv_cache_mode, + token_bias, + )? + } else if let Some(ref draft_model_path) = args.model.draft_model { // resolve the effective DrafterKind from // (a) the explicit `--draft-kind` CLI flag, OR // (b) the drafter's `config.json::model_type` auto-detection. @@ -1962,6 +1972,65 @@ pub(super) fn run_generation_mode( Ok(output) } +/// Offline prompt-lookup speculative decoding (`--prompt-lookup`). +/// +/// Refuses the model families whose state cannot be rewound after a rejected +/// block, and multimodal prompts, whose embeddings the generator cannot take. +/// Runs one short warmup generation first, as `generate_standard` does, so the +/// timed run does not pay for kernel compilation. The warmup decodes a few +/// rounds rather than one token so the verify widths get compiled too. +/// +/// The returned stats time the whole call, prefill included, exactly as +/// `generate_standard` does, so the printed tok/s compares like for like. +fn run_prompt_lookup( + model: &mlxcel::LoadedModel, + args: &GenerateArgs, + prompt_tokens: &[i32], + sampling_config: &SamplingConfig, + has_multimodal_input: bool, + kv_cache_mode: KVCacheMode, + token_bias: TokenBiasMap, +) -> Result<(Vec, GenerationStats)> { + ensure!( + !has_multimodal_input, + "--prompt-lookup supports text-only prompts; drop --image/--audio/--video" + ); + if let Some(reason) = mlxcel::prompt_lookup_unsupported_reason(model) { + anyhow::bail!("--prompt-lookup cannot run on this model: {reason}"); + } + let config = args.prompt_lookup.config(); + config + .validate() + .map_err(|err| anyhow!("--prompt-lookup: {err}"))?; + + let mut generator = mlxcel::PromptLookupGenerator::new(config) + .with_kv_cache_mode(kv_cache_mode) + .with_token_bias(token_bias); + let warmup_tokens = args.generation.max_tokens.min(16); + let _ = generator.generate(model, prompt_tokens, warmup_tokens, sampling_config); + let start_time = Instant::now(); + let (tokens, measured) = generator.generate( + model, + prompt_tokens, + args.generation.max_tokens, + sampling_config, + ); + let total_time = start_time.elapsed(); + // Diagnostics go to stderr on their own line: stdout still holds the + // echoed prompt without a trailing newline, and the reply follows it. + eprintln!(); + eprintln!("{}", generator.stats().summary_line(tokens.len())); + // `--profile` prints the prefill / decode split the generator measured, + // as the plain path does; otherwise the one-line rate covers the whole + // call, prefill included, like `generate_standard`. + let stats = if args.generation.profile { + measured + } else { + generation_stats_from_duration(prompt_tokens.len(), tokens.len(), total_time) + }; + Ok((tokens, stats)) +} + /// Routing gate for the offline MTP speculative path (issue #166). /// /// Returns `true` only when the operator explicitly passed `--draft-kind mtp` @@ -2558,6 +2627,10 @@ pub(crate) fn run_generate(mut args: GenerateArgs) -> Result<()> { if args.generation.layout_detections.is_some() { args.generation.prompt = Some(String::new()); } else { + ensure!( + !args.prompt_lookup.prompt_lookup, + "--prompt-lookup requires a one-shot -p/--prompt run (not interactive chat)" + ); let opts = chat_options_from_args(&args)?; return crate::commands::run_chat(opts); } diff --git a/src/commands/generate_tests.rs b/src/commands/generate_tests.rs index 4bcb66065..fd79e511b 100644 --- a/src/commands/generate_tests.rs +++ b/src/commands/generate_tests.rs @@ -638,6 +638,7 @@ fn sample_generate_args(model_path: PathBuf) -> crate::GenerateArgs { }, lang_bias: mlxcel::lang_bias::LangBiasCliArgs::default(), speculative: mlxcel::cli::speculative_args::SpeculativeArgs::default(), + prompt_lookup: crate::PromptLookupOptions::default(), // Default to None so existing tests stay on // the bit-exact baseline load path; tests that need surgery // override this field explicitly. diff --git a/src/commands/run.rs b/src/commands/run.rs index 6713e2167..fa7786542 100644 --- a/src/commands/run.rs +++ b/src/commands/run.rs @@ -143,6 +143,7 @@ impl RunArgs { tensor_parallel: crate::TensorParallelOptions::default(), lang_bias: mlxcel::lang_bias::LangBiasCliArgs::default(), speculative: mlxcel::cli::speculative_args::SpeculativeArgs::default(), + prompt_lookup: crate::PromptLookupOptions::default(), #[cfg(feature = "surgery")] surgery: None, } diff --git a/src/lib.rs b/src/lib.rs index 635ade59c..985d29c1f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -79,6 +79,10 @@ pub use mlxcel_core::generate::{ SamplingConfig, }; pub use mlxcel_core::speculative::SpeculativeGenerator; +pub use mlxcel_core::speculative::prompt_lookup::{ + PromptLookupConfig, PromptLookupGenerator, prompt_lookup_unsupported_reason, + supports_prompt_lookup, +}; #[cfg(feature = "xla-diagnostics")] pub use multimodal::host_preprocessor::LlavaHostReferenceCapture; #[cfg(feature = "xla-iree")] diff --git a/src/lib/mlxcel-core/src/generate.rs b/src/lib/mlxcel-core/src/generate.rs index 7acac9935..21d814e3a 100644 --- a/src/lib/mlxcel-core/src/generate.rs +++ b/src/lib/mlxcel-core/src/generate.rs @@ -338,6 +338,69 @@ fn chunked_prefill_last_logits( logits.expect("chunked_prefill_last_logits requires a non-empty prompt") } +/// Prefill the whole prompt and return the `[1, 1, vocab]` logits of its last +/// position, the way [`CxxGenerator::generate_streaming`] does. +/// +/// Picks cache-level chunked prefill when `MLXCEL_PREFILL_CHUNK` applies, +/// tile-aligned padded prefill on M5+ hardware, and one `forward_last_logits` +/// over the prompt otherwise. Speculative generators that must stay +/// byte-identical to plain decoding at greedy sampling call this instead of +/// their own prefill: splitting the prompt differently changes the fp16 +/// rounding of the cached keys and values, and with it the first token. +/// +/// Used by: `CxxGenerator::generate_streaming`, `PromptLookupGenerator` +pub(crate) fn prefill_prompt_last_logits( + model: &M, + caches: &mut [KVCache], + prompt_tokens: &[i32], +) -> UniquePtr { + let actual_len = prompt_tokens.len(); + let prefill_chunk = effective_prefill_chunk( + prefill_chunk_len(), + model.supports_chunked_prefill(), + actual_len, + ); + if let Some(chunk) = prefill_chunk { + // Cache-level chunked prefill (MLXCEL_PREFILL_CHUNK, default 2048). + chunked_prefill_last_logits(model, caches, prompt_tokens, chunk) + } else if should_align_prefill() && model.supports_padded_prefill() { + let padded_len = align_to_na_tile(actual_len); + let (padded_tokens, mask_opt) = pad_tokens_for_prefill( + prompt_tokens, + padded_len, + model.supports_maskless_padded_prefill(), + ); + let input = ffi::from_slice_i32(&padded_tokens, &[1, padded_len as i32]); + // Last *real* token position; `forward_last_logits` slices there, + // replacing the previous forward + `logits_at_position` pair. + let raw_logits = model.forward_last_logits( + &input, + caches, + mask_opt.as_ref().map(|m| m.as_ref().unwrap()), + actual_len.saturating_sub(1), + ); + // Trim padding positions from all KV caches so decode uses the + // correct cache offset (actual_len, not padded_len). + if padded_len > actual_len { + trim_caches_to_actual_len(caches, actual_len, padded_len); + model.trim_internal_caches((padded_len - actual_len) as i32); + } + raw_logits + } else { + let input = ffi::from_slice_i32(prompt_tokens, &[1, actual_len as i32]); + model.forward_last_logits(&input, caches, None, actual_len.saturating_sub(1)) + } +} + +/// Per-layer KV cache modes for `n_layers` caches under the nominal `mode`, +/// with the Boundary-V upgrade (`MLXCEL_KV_BOUNDARY_V_LAYERS`) applied. +/// +/// Used by: `CxxGenerator`, `PromptLookupGenerator` +pub(crate) fn resolve_kv_cache_layer_modes(mode: KVCacheMode, n_layers: usize) -> Vec { + let requested = crate::cache::turbo::boundary_v_layers_from_env(); + crate::cache::turbo::resolve_layer_modes(mode, n_layers, requested) +} + /// Trait for language models that can be used for generation pub trait LanguageModel { /// Forward pass through the model @@ -1657,9 +1720,7 @@ impl CxxGenerator { /// including `reset_with_model` boundary cases. /// fn resolved_layer_modes_for(&self, n_layers: usize) -> Vec { - let nominal = self.kv_cache_mode; - let requested = crate::cache::turbo::boundary_v_layers_from_env(); - crate::cache::turbo::resolve_layer_modes(nominal, n_layers, requested) + resolve_kv_cache_layer_modes(self.kv_cache_mode, n_layers) } fn apply_kv_cache_mode_with_boundary_policy(&mut self) -> Vec { @@ -1751,42 +1812,7 @@ impl CxxGenerator { // Prefill: process all prompt tokens at once. // On M5+ hardware pad the sequence to a 32-token tile boundary for // optimal Neural Accelerator throughput. - let actual_len = prompt_tokens.len(); - let prefill_chunk = effective_prefill_chunk( - prefill_chunk_len(), - model.supports_chunked_prefill(), - actual_len, - ); - let logits = if let Some(chunk) = prefill_chunk { - // Cache-level chunked prefill (MLXCEL_PREFILL_CHUNK, default 2048). - chunked_prefill_last_logits(model, &mut self.caches, prompt_tokens, chunk) - } else if should_align_prefill() && model.supports_padded_prefill() { - let padded_len = align_to_na_tile(actual_len); - let (padded_tokens, mask_opt) = pad_tokens_for_prefill( - prompt_tokens, - padded_len, - model.supports_maskless_padded_prefill(), - ); - let input = ffi::from_slice_i32(&padded_tokens, &[1, padded_len as i32]); - // Last *real* token position; `forward_last_logits` slices there, - // replacing the previous forward + `logits_at_position` pair. - let raw_logits = model.forward_last_logits( - &input, - &mut self.caches, - mask_opt.as_ref().map(|m| m.as_ref().unwrap()), - actual_len.saturating_sub(1), - ); - // Trim padding positions from all KV caches so decode uses the - // correct cache offset (actual_len, not padded_len). - if padded_len > actual_len { - trim_caches_to_actual_len(&mut self.caches, actual_len, padded_len); - model.trim_internal_caches((padded_len - actual_len) as i32); - } - raw_logits - } else { - let input = ffi::from_slice_i32(prompt_tokens, &[1, actual_len as i32]); - model.forward_last_logits(&input, &mut self.caches, None, actual_len.saturating_sub(1)) - }; + let logits = prefill_prompt_last_logits(model, &mut self.caches, prompt_tokens); if trace_dtype { ffi::eval(&logits); diff --git a/src/lib/mlxcel-core/src/speculative/mod.rs b/src/lib/mlxcel-core/src/speculative/mod.rs index e97bfde0f..acc0d6d12 100644 --- a/src/lib/mlxcel-core/src/speculative/mod.rs +++ b/src/lib/mlxcel-core/src/speculative/mod.rs @@ -57,8 +57,12 @@ //! `trim_caches`). See [`mtp::MtpGenerator`]. //! - [`stochastic_accept`] — the distribution-preserving acceptance rule and //! its residual resample, shared by every speculative verify path. +//! - [`prompt_lookup`] — drafter-free speculation that proposes the tokens +//! following an earlier occurrence of the sequence's tail. See +//! [`prompt_lookup::PromptLookupGenerator`]. pub mod mtp; +pub mod prompt_lookup; pub mod stochastic_accept; use crate::cache::can_trim_prompt_cache; diff --git a/src/lib/mlxcel-core/src/speculative/prompt_lookup.rs b/src/lib/mlxcel-core/src/speculative/prompt_lookup.rs new file mode 100644 index 000000000..8daf8932a --- /dev/null +++ b/src/lib/mlxcel-core/src/speculative/prompt_lookup.rs @@ -0,0 +1,816 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Prompt-lookup (n-gram) speculative decoding. +//! +//! A drafter that needs no second model. When the last few tokens of the +//! sequence also appear earlier in the context (prompt plus everything emitted +//! so far), the tokens that followed that earlier occurrence are proposed as +//! the draft, and the target verifies the whole block in one forward, exactly +//! as [`super::SpeculativeGenerator`] verifies a draft model's proposals. +//! +//! It pays off on replies that copy long runs of the prompt: editing a +//! function, fixing typos, adding a field to a JSON document, turning a CSV +//! into a table. Those replies repeat the prompt word for word for most of +//! their length, so most proposals are accepted and a verify forward (whose +//! cost is dominated by reading the weights once) emits several tokens. +//! +//! ## Round shape +//! +//! Each round forwards `[current_token, d_0, ..., d_{k-1}]` through the target, +//! keeps the longest prefix of the proposals the target agrees with, emits one +//! more token from the target (the replacement at the first disagreement, or +//! the bonus token after a fully accepted block), and trims the rejected tail +//! off the KV caches. With no match the round is `k = 0`: a plain one-token +//! decode step. +//! +//! ## Pipelining +//! +//! A round without a proposal is pipelined the way plain decoding is: its +//! one-token forward is submitted from the previous round's still-lazy sample, +//! before the host reads that sample, so the GPU goes straight from one step to +//! the next. Waiting on each token before encoding the next step cost 12% of +//! decode time on an M4 Pro. Pipelining starts only after a few rounds in a +//! row found no proposal, because a pipelined step cannot carry one. When +//! lookup finds a proposal while a step is in flight, the host reads the +//! in-flight token and verifies from there synchronously; that switch costs no +//! extra forward. +//! +//! ## Adaptive drafting +//! +//! A verify block is not free: on an M4 Pro an eight-token block costs about +//! two and a half one-token steps. On prose, where lookup matches a common +//! token and the target rarely agrees, proposing every round made decoding a +//! quarter slower than plain decode. [`DraftGovernor`] keeps the block short +//! while recent rounds land few tokens and, after repeated misses, stops +//! proposing for a few rounds (doubling up to a cap) before trying again. +//! Edits keep landing five or six tokens a round, so they keep the full block. +//! +//! ## Acceptance +//! +//! A lookup proposal is deterministic, so its proposal distribution `q` is one +//! hot on the proposed token. Under that `q` the acceptance-optimal rule of +//! modified rejection sampling accepts with probability `p(d)` and resamples +//! from `p` restricted to tokens other than `d`, which is exactly what the +//! sampler-match rule does: draw `t ~ p` and accept iff `t == d`. So one rule +//! serves every sampler configuration and the emitted stream is distributed as +//! the target's own. +//! +//! ## Greedy exactness +//! +//! The prompt is prefilled by the same routine plain decoding uses +//! ([`crate::generate::prefill_prompt_last_logits`]), so the caches and the +//! first token match plain decoding exactly, and a round without a proposal +//! is the same one-token forward. A round with proposals runs a multi-token +//! forward, whose kernels may round differently from the one-token path; where +//! the target's top two logits are that close, greedy output can take the +//! other token from there on. Measured on Qwen3-1.7B-4bit (M4 Pro), edit +//! replies stayed identical to plain decoding. This is the same jitter class as +//! the draft-model path; nothing here gates on it the way the MTP exactness +//! probe does. +//! +//! ## Cache contract +//! +//! Rollback is `KVCache::trim`, so the target must keep all of its sequence +//! state in the external KV caches. Recurrent and hybrid families (Mamba, +//! GatedDeltaNet, NemotronH, ...) fold every consumed token into a fixed-size +//! state that cannot be trimmed; callers must refuse them, see +//! [`prompt_lookup_unsupported_reason`]. + +use crate::cache::{KVCacheMode, can_trim_prompt_cache}; +use crate::drafter::dflash::materialize_argmax_i32_vec; +use crate::ffi; +use crate::ffi::MlxThreadLocalStream; +use crate::generate::{ + GenerationStats, LanguageModel, SamplingConfig, prefill_prompt_last_logits, + resolve_kv_cache_layer_modes, +}; +use crate::generation_policy::{initial_token_history, merged_eos_token_ids, seed_rng_if_needed}; +use crate::layers::KVCache; +use crate::loop_detection::detect_repetition_loop; +use crate::sampling::{TokenBiasMap, sample_token_optimized}; +use crate::streams::{install_thread_local_default_stream, shared_thread_local_generation_stream}; +use cxx::UniquePtr; +use std::borrow::Cow; +use std::collections::HashMap; +use std::time::Instant; + +use super::trim_caches; + +/// Default longest n-gram matched against the context. +pub const DEFAULT_NGRAM_MAX: usize = 3; +/// Default shortest n-gram matched against the context. +/// +/// A one-token match is usually a common token (a space, `the`, a comma), and +/// what followed it earlier rarely comes next. Measured on Qwen3-1.7B-4bit, a +/// minimum of 1 slowed prose by a quarter while a minimum of 2 kept edits just +/// as fast. +pub const DEFAULT_NGRAM_MIN: usize = 2; +/// Default maximum number of tokens proposed per round (verify block of 8). +pub const DEFAULT_MAX_DRAFT: usize = 7; + +/// Tunables for [`PromptLookupGenerator`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PromptLookupConfig { + /// Longest suffix n-gram tried first. + pub ngram_max: usize, + /// Shortest suffix n-gram tried before giving up on a round. + pub ngram_min: usize, + /// Maximum proposals per round. The verify block is one wider. + pub max_draft: usize, + /// Shorten or pause proposals while they stop landing, see + /// [`DraftGovernor`]. `false` proposes up to `max_draft` every round. + pub adaptive: bool, +} + +impl Default for PromptLookupConfig { + fn default() -> Self { + Self { + ngram_max: DEFAULT_NGRAM_MAX, + ngram_min: DEFAULT_NGRAM_MIN, + max_draft: DEFAULT_MAX_DRAFT, + adaptive: true, + } + } +} + +impl PromptLookupConfig { + /// Reject configurations that could never propose anything. + pub fn validate(&self) -> Result<(), String> { + if self.ngram_min == 0 { + return Err("prompt-lookup ngram-min must be at least 1".to_string()); + } + if self.ngram_max < self.ngram_min { + return Err(format!( + "prompt-lookup ngram-max ({}) must be >= ngram-min ({})", + self.ngram_max, self.ngram_min + )); + } + if self.max_draft == 0 { + return Err("prompt-lookup max-draft must be at least 1".to_string()); + } + Ok(()) + } +} + +/// Why `model` cannot run under [`PromptLookupGenerator`], or `None` when it +/// can. +/// +/// Verification rolls rejected tokens back with `KVCache::trim`, which is only +/// sound when every piece of per-sequence state lives in the external caches. +/// `supports_batching` is the existing marker for that, and +/// `supports_padded_prefill` is `false` for the families whose state trailing +/// tokens corrupt past what a trim undoes, which is the same hazard a rejected +/// draft block poses. +pub fn prompt_lookup_unsupported_reason(model: &M) -> Option<&'static str> { + if !model.supports_batching() { + return Some( + "this model keeps sequence state outside the external KV caches (recurrent or \ + hybrid layers, or model-owned sliding-window caches), so a rejected draft \ + cannot be rolled back", + ); + } + if !model.supports_padded_prefill() { + return Some( + "this model's state is corrupted by trailing tokens that a cache trim cannot \ + undo (the same reason it opts out of padded prefill), so a rejected draft \ + cannot be rolled back", + ); + } + None +} + +/// Whether `model` can run under [`PromptLookupGenerator`]. See +/// [`prompt_lookup_unsupported_reason`] for why not. +pub fn supports_prompt_lookup(model: &M) -> bool { + prompt_lookup_unsupported_reason(model).is_none() +} + +/// Smallest pause, in rounds, after proposals stop landing. +const MIN_COOLDOWN: usize = 4; +/// Largest pause. Kept short so an edit that writes a new passage and then +/// resumes copying is back on full blocks within a line or two. +const MAX_COOLDOWN: usize = 32; +/// Consecutive drafted rounds that land nothing before a pause. +const MISSES_BEFORE_COOLDOWN: usize = 3; +/// Consecutive rounds without a proposal before the decode loop pipelines. +/// +/// A pipelined step is submitted before the host knows the token it follows, +/// so it cannot carry a proposal: the first proposal after a pipelined run +/// waits one step. Edits that miss for a token or two at each changed field +/// (a JSON id, a renamed identifier) would pay that step at every field, so +/// only a run of plain rounds switches to pipelining. +const PLAIN_ROUNDS_BEFORE_PIPELINE: usize = 2; + +/// Decides, round by round, how many lookup tokens to propose. +/// +/// Tracks an exponential moving average of the proposals each drafted round +/// landed and caps the next block at twice that plus two, so a run of full +/// matches keeps the full block and a run of one-token landings verifies +/// five-wide blocks instead of eight. Lookup landings are all-or-nothing more +/// often than not (a copy either resumes or it does not), so the cap leaves +/// room past the average. After [`MISSES_BEFORE_COOLDOWN`] drafted rounds in +/// a row land nothing, it stops proposing for a cooldown that doubles on every +/// repeat up to [`MAX_COOLDOWN`] and resets once a round lands a token. +#[derive(Debug, Clone)] +pub(crate) struct DraftGovernor { + max_draft: usize, + adaptive: bool, + accepted_ema: f64, + misses: usize, + cooldown: usize, + next_cooldown: usize, +} + +impl DraftGovernor { + pub(crate) fn new(config: &PromptLookupConfig) -> Self { + Self { + max_draft: config.max_draft, + adaptive: config.adaptive, + // Optimistic start: the first match in an edit should get the + // whole block. + accepted_ema: config.max_draft as f64, + misses: 0, + cooldown: 0, + next_cooldown: MIN_COOLDOWN, + } + } + + /// Proposals allowed this round; `0` means run a plain decode step. A + /// paused round counts down the pause. + pub(crate) fn budget(&mut self) -> usize { + if !self.adaptive { + return self.max_draft; + } + if self.cooldown > 0 { + self.cooldown -= 1; + return 0; + } + ((2.0 * self.accepted_ema).round() as usize + 2).clamp(1, self.max_draft) + } + + /// Record a drafted round that landed `accepted` proposals. + pub(crate) fn record(&mut self, accepted: usize) { + if !self.adaptive { + return; + } + self.accepted_ema = 0.5 * self.accepted_ema + 0.5 * accepted as f64; + if accepted > 0 { + self.misses = 0; + self.next_cooldown = MIN_COOLDOWN; + return; + } + self.misses += 1; + if self.misses >= MISSES_BEFORE_COOLDOWN { + self.misses = 0; + self.cooldown = self.next_cooldown; + self.next_cooldown = (self.next_cooldown * 2).min(MAX_COOLDOWN); + } + } +} + +/// Propose up to `max_draft` tokens by matching the tail of `context` against +/// an earlier position in `context`. +/// +/// Tries suffix lengths from `ngram_max` down to `ngram_min` and, for each, +/// takes the most recent earlier occurrence that has at least one token after +/// it. The most recent occurrence is preferred because a reply that copies the +/// prompt advances through it front to back, so the latest match is the one +/// the copy is currently walking. +/// +/// Returns an empty vector when nothing matches. +pub fn find_draft(context: &[i32], config: &PromptLookupConfig) -> Vec { + let len = context.len(); + if config.max_draft == 0 || config.ngram_min == 0 { + return Vec::new(); + } + let ngram_max = config.ngram_max.min(len.saturating_sub(1)); + if ngram_max < config.ngram_min { + return Vec::new(); + } + + for n in (config.ngram_min..=ngram_max).rev() { + let pattern = &context[len - n..]; + // `start` is where a candidate occurrence begins. It must end before + // the suffix itself does (`start + n < len`) so at least one token + // follows it; walk from the latest such position backwards. + let mut start = len - n; + while start > 0 { + start -= 1; + if &context[start..start + n] == pattern { + let from = start + n; + let to = (from + config.max_draft).min(len); + return context[from..to].to_vec(); + } + } + } + Vec::new() +} + +/// Incremental index answering [`find_draft`] without scanning the context. +/// +/// For each n-gram length in the configured range it maps an n-gram to the +/// start of its most recent occurrence that already has a token after it, +/// which is exactly the occurrence [`find_draft`] settles on. It is built once +/// from the prompt and extended by the tokens emitted since the last call, so +/// a lookup costs one hash probe per n-gram length however long the context +/// grows, where a scan costs a pass over the whole context every round. +#[derive(Debug)] +pub(crate) struct NgramIndex { + ngram_min: usize, + /// `latest[n - ngram_min]`: n-gram to the start of its latest occurrence + /// that has a follower. + latest: Vec, usize>>, + /// Context length already indexed. + indexed_len: usize, +} + +impl NgramIndex { + pub(crate) fn new(config: &PromptLookupConfig) -> Self { + let lengths = (config.ngram_min..=config.ngram_max).count(); + Self { + ngram_min: config.ngram_min, + latest: (0..lengths).map(|_| HashMap::new()).collect(), + indexed_len: 0, + } + } + + /// Index every n-gram that gained a follower since the last call. + /// `context` must extend the context of the previous call. + pub(crate) fn extend(&mut self, context: &[i32]) { + // An n-gram ending at `end` has a follower once `end + 1 < len`. The + // previous call indexed every end up to `indexed_len - 2`. + for end in self.indexed_len.saturating_sub(1)..context.len().saturating_sub(1) { + for (offset, map) in self.latest.iter_mut().enumerate() { + let n = self.ngram_min + offset; + if end + 1 < n { + break; + } + let start = end + 1 - n; + map.insert(context[start..=end].into(), start); + } + } + self.indexed_len = context.len(); + } + + /// Same proposal as [`find_draft`] on the indexed `context`. + pub(crate) fn find(&self, context: &[i32], config: &PromptLookupConfig) -> Vec { + debug_assert_eq!(self.indexed_len, context.len(), "index is stale"); + let len = context.len(); + if config.max_draft == 0 || config.ngram_min == 0 { + return Vec::new(); + } + let ngram_max = config.ngram_max.min(len.saturating_sub(1)); + if ngram_max < config.ngram_min { + return Vec::new(); + } + for n in (config.ngram_min..=ngram_max).rev() { + if let Some(&start) = self.latest[n - self.ngram_min].get(&context[len - n..]) { + let from = start + n; + let to = (from + config.max_draft).min(len); + return context[from..to].to_vec(); + } + } + Vec::new() + } +} + +/// Acceptance accounting for one [`PromptLookupGenerator::generate`] call. +#[derive(Debug, Clone, Copy, Default, PartialEq)] +pub struct PromptLookupStats { + /// Target forwards in the decode loop (one per round). + pub rounds: usize, + /// Rounds whose lookup found a match and verified a draft block. + pub drafted_rounds: usize, + /// One-token rounds run without looking up a proposal, because the + /// [`DraftGovernor`] paused proposals or `max_tokens` left no room. + pub paused_rounds: usize, + /// Tokens proposed across all rounds. + pub proposed_draft_tokens: usize, + /// Proposals the target accepted. + pub accepted_draft_tokens: usize, +} + +impl PromptLookupStats { + /// Accepted proposals per proposed token. + pub fn acceptance_rate(&self) -> f64 { + ratio(self.accepted_draft_tokens, self.proposed_draft_tokens) + } + + /// Emitted tokens per target forward in the decode loop. `1.0` is plain + /// decoding; this is the multiplier lookup buys before verify overhead. + pub fn tokens_per_forward(&self, generated_tokens: usize) -> f64 { + // The first token comes out of the prefill forward, not a round. + ratio(generated_tokens.saturating_sub(1), self.rounds) + } + + /// One-line summary for the CLI, which installs no tracing subscriber. + pub fn summary_line(&self, generated_tokens: usize) -> String { + format!( + "[Prompt lookup] rounds={} drafted_rounds={} paused_rounds={} proposed={} \ + accepted={} acceptance_rate={:.4} tokens_per_forward={:.4}", + self.rounds, + self.drafted_rounds, + self.paused_rounds, + self.proposed_draft_tokens, + self.accepted_draft_tokens, + self.acceptance_rate(), + self.tokens_per_forward(generated_tokens), + ) + } +} + +fn ratio(num: usize, den: usize) -> f64 { + if den == 0 { + 0.0 + } else { + num as f64 / den as f64 + } +} + +/// Speculative generator whose drafter is [`find_draft`], answered from an +/// [`NgramIndex`]. +pub struct PromptLookupGenerator { + caches: Vec, + config: PromptLookupConfig, + kv_cache_mode: KVCacheMode, + /// Prompt followed by every emitted token. The lookup searches this. + context: Vec, + generated_tokens: Vec, + generation_stream: Option>, + token_bias: TokenBiasMap, + stats: PromptLookupStats, +} + +impl PromptLookupGenerator { + /// Create a generator. Caches are built from the model on every + /// [`Self::generate`] call. + pub fn new(config: PromptLookupConfig) -> Self { + Self { + caches: Vec::new(), + config, + kv_cache_mode: KVCacheMode::Fp16, + context: Vec::new(), + generated_tokens: Vec::new(), + generation_stream: shared_thread_local_generation_stream(), + token_bias: TokenBiasMap::default(), + stats: PromptLookupStats::default(), + } + } + + /// Store the KV caches in `mode`, with the same Boundary-V policy plain + /// decoding applies. + pub fn with_kv_cache_mode(mut self, mode: KVCacheMode) -> Self { + self.kv_cache_mode = mode; + self + } + + /// Attach a pre-resolved token bias, applied to every target sample. + pub fn with_token_bias(mut self, bias: TokenBiasMap) -> Self { + self.token_bias = bias; + self + } + + /// Acceptance accounting for the most recent [`Self::generate`] call. + pub fn stats(&self) -> PromptLookupStats { + self.stats + } + + /// Fresh caches from the model, in the configured mode, as + /// `CxxGenerator::reset_with_model` builds them. + fn reset(&mut self, model: &M) { + model.reset_runtime_state(); + self.caches = model.make_caches(); + let modes = resolve_kv_cache_layer_modes(self.kv_cache_mode, self.caches.len()); + for (cache, mode) in self.caches.iter_mut().zip(modes) { + cache.mode = mode; + } + model.set_kv_cache_layer_modes(resolve_kv_cache_layer_modes( + self.kv_cache_mode, + model.num_layers(), + )); + self.context.clear(); + self.generated_tokens.clear(); + self.stats = PromptLookupStats::default(); + } + + /// Same composition rule as `SpeculativeGenerator::compose_target_sampling`: + /// the caller's own bias wins, otherwise the cached one is injected. + fn compose_sampling<'a>(&self, sampling: &'a SamplingConfig) -> Cow<'a, SamplingConfig> { + if self.token_bias.is_empty() || !sampling.token_bias.is_empty() { + Cow::Borrowed(sampling) + } else { + let mut cloned = sampling.clone(); + cloned.token_bias = self.token_bias.clone(); + Cow::Owned(cloned) + } + } + + /// Generate up to `max_tokens` tokens after `prompt_tokens`. + /// + /// # Panics + /// + /// Panics on an empty prompt or a non-trimmable cache. Callers must check + /// [`supports_prompt_lookup`] first. + pub fn generate( + &mut self, + model: &M, + prompt_tokens: &[i32], + max_tokens: usize, + sampling: &SamplingConfig, + ) -> (Vec, GenerationStats) { + self.reset(model); + + assert!( + can_trim_prompt_cache(&self.caches), + "prompt-lookup decoding requires a trimmable prompt cache" + ); + assert!( + !prompt_tokens.is_empty(), + "prompt-lookup generate requires at least one prompt token" + ); + + let sampling_cow = self.compose_sampling(sampling); + let sampling: &SamplingConfig = sampling_cow.as_ref(); + seed_rng_if_needed(sampling); + install_thread_local_default_stream(self.generation_stream.as_ref()); + + let eos_tokens = merged_eos_token_ids(model.eos_token_ids(), &sampling.stop_token_ids); + let needs_history = sampling.needs_token_history(); + let mut token_history = initial_token_history(prompt_tokens, needs_history); + // Pure argmax with nothing that depends on the emitted prefix: every + // verify position can be decided from one batched argmax and a single + // host read. Anything else goes position by position through the + // sampler, as the draft-model path does. + let batched_argmax = + sampling.is_greedy_path() && !needs_history && sampling.token_bias.is_empty(); + + self.context.extend_from_slice(prompt_tokens); + + // PREFILL through the plain-decode routine, so the caches and the + // first token are exactly what plain decoding produces. + let prefill_start = Instant::now(); + let logits = prefill_prompt_last_logits(model, &mut self.caches, prompt_tokens); + ffi::clear_memory_cache(); + let (first_arr, _) = sample_token_optimized(&logits, sampling, &token_history); + ffi::async_eval(&first_arr); + // Index the prompt while the GPU finishes the prefill. + let mut index = NgramIndex::new(&self.config); + index.extend(&self.context); + let first_token = ffi::item_i32(&first_arr); + let prefill_time = prefill_start.elapsed(); + + if eos_tokens.contains(&first_token) || max_tokens == 0 { + return self.finish(prompt_tokens.len(), prefill_time, Default::default()); + } + let mut done = self.emit( + first_token, + needs_history.then_some(&mut token_history), + max_tokens, + sampling, + ); + if max_tokens > 1 { + for cache in &mut self.caches { + cache.prepare_turbo4_delegated_for_decode(); + } + } + + // DECODE. + let decode_start = Instant::now(); + let mut current_token = first_token; + let mut governor = DraftGovernor::new(&self.config); + // Rounds without a proposal are pipelined the way plain decoding is: + // the next step is encoded and submitted from the still-lazy sample + // before the host reads it, so the GPU does not idle while the host + // reads a token and builds the next graph (measured 12% of decode time + // on an M4 Pro). A sampler that reads the emitted history needs each + // token on the host first, so those runs stay synchronous. + let pipeline = !needs_history && std::env::var("MLXCEL_FORCE_SYNC").is_err(); + // `Some(next)` while a one-token forward of `current_token` is in + // flight: the cache already holds `current_token` and `next` is the + // lazy sample of the token after it. `None` means the cache holds + // every emitted token except `current_token`. + let mut in_flight: Option> = None; + // Rounds in a row, this one included, whose lookup proposed nothing. + let mut plain_streak = 0usize; + + while !done { + // Never propose past `max_tokens`: every accepted proposal is + // emitted, and the round also emits one target token. + let remaining = max_tokens - self.generated_tokens.len(); + let budget = governor.budget().min(remaining.saturating_sub(1)); + let draft = if budget == 0 { + Vec::new() + } else { + index.extend(&self.context); + let mut draft = index.find(&self.context, &self.config); + draft.truncate(budget); + draft + }; + + if draft.is_empty() { + plain_streak += 1; + } else { + plain_streak = 0; + } + if let Some(next) = in_flight.take() { + // Keep the pipeline full unless lookup has something to verify + // or `next` is the last token allowed. + if draft.is_empty() && remaining > 1 { + let next_input = ffi::reshape_token_for_forward(&next); + let logits = model.forward(&next_input, &mut self.caches, None); + let (after_next, _) = sample_token_optimized(&logits, sampling, &token_history); + ffi::async_eval(&after_next); + in_flight = Some(after_next); + self.count_plain_round(budget); + } + let token = ffi::item_i32(&next); + if eos_tokens.contains(&token) { + break; + } + done = self.emit(token, None, max_tokens, sampling); + current_token = token; + // With nothing in flight now, the next round verifies from + // `current_token` synchronously; the proposal found this round + // was for the token before it and is looked up again there. + continue; + } + + if draft.is_empty() + && pipeline + && remaining > 1 + && plain_streak > PLAIN_ROUNDS_BEFORE_PIPELINE + { + let input = ffi::from_slice_i32(&[current_token], &[1, 1]); + let logits = model.forward(&input, &mut self.caches, None); + let (next, _) = sample_token_optimized(&logits, sampling, &token_history); + ffi::async_eval(&next); + in_flight = Some(next); + self.count_plain_round(budget); + continue; + } + + let mut verify_tokens = Vec::with_capacity(draft.len() + 1); + verify_tokens.push(current_token); + verify_tokens.extend_from_slice(&draft); + let verify_len = verify_tokens.len(); + let verify_input = ffi::from_slice_i32(&verify_tokens, &[1, verify_len as i32]); + let logits = model.forward(&verify_input, &mut self.caches, None); + + // Target token at every verify position, in order. Position `i` + // is the target's choice for the slot draft[i] claims; position + // `draft.len()` is the bonus token after a fully accepted block. + let vocab = ffi::array_shape(&logits)[2]; + let greedy_targets = batched_argmax.then(|| { + let argmax = ffi::argmax_last_axis(&logits); + ffi::eval(&argmax); + materialize_argmax_i32_vec(&argmax, verify_len) + }); + + let mut accepted = 0usize; + for pos in 0..verify_len { + let target = match &greedy_targets { + Some(targets) => targets[pos], + None => { + let pos_logits = + ffi::slice(&logits, &[0, pos as i32, 0], &[1, pos as i32 + 1, vocab]); + let (arr, _) = + sample_token_optimized(&pos_logits, sampling, &token_history); + ffi::eval(&arr); + ffi::item_i32(&arr) + } + }; + + if eos_tokens.contains(&target) { + done = true; + break; + } + done = self.emit( + target, + needs_history.then_some(&mut token_history), + max_tokens, + sampling, + ); + current_token = target; + if done { + break; + } + // `target` is emitted either way; it only extends the round + // when it confirms the proposal in the same slot. + if pos < draft.len() && target == draft[pos] { + accepted += 1; + } else { + break; + } + } + + if draft.is_empty() { + self.count_plain_round(budget); + } else { + self.stats.rounds += 1; + self.stats.drafted_rounds += 1; + governor.record(accepted); + } + self.stats.proposed_draft_tokens += draft.len(); + self.stats.accepted_draft_tokens += accepted; + + // The verify forward appended `current, d_0..d_{k-1}`. The cache + // must end holding every emitted token except the next round's + // `current_token`, i.e. through `d_{a-1}`: trim `k - a`. + trim_caches(&mut self.caches, (draft.len() - accepted) as i32); + + if crate::memory::should_clear_cache_at( + self.generated_tokens.len(), + crate::memory::cache_clear_interval(), + ) { + ffi::clear_memory_cache(); + } + } + + let decode_time = decode_start.elapsed(); + tracing::info!( + rounds = self.stats.rounds, + drafted_rounds = self.stats.drafted_rounds, + proposed_draft_tokens = self.stats.proposed_draft_tokens, + accepted_draft_tokens = self.stats.accepted_draft_tokens, + generated_tokens = self.generated_tokens.len(), + "prompt-lookup decode finished" + ); + self.finish(prompt_tokens.len(), prefill_time, decode_time) + } + + /// Count a one-token forward. `budget == 0` means the governor paused + /// proposals (or `max_tokens` left room for none) rather than lookup + /// finding no match. + fn count_plain_round(&mut self, budget: usize) { + self.stats.rounds += 1; + if budget == 0 { + self.stats.paused_rounds += 1; + } + } + + /// Emit `token` and report whether generation stops after it: at + /// `max_tokens`, or when the repetition-loop guard fires, which plain + /// decoding checks after every token too. + fn emit( + &mut self, + token: i32, + history: Option<&mut Vec>, + max_tokens: usize, + sampling: &SamplingConfig, + ) -> bool { + self.generated_tokens.push(token); + self.context.push(token); + if let Some(history) = history { + history.push(token); + } + self.generated_tokens.len() >= max_tokens + || detect_repetition_loop(&self.generated_tokens, &sampling.loop_detection) + } + + fn finish( + &self, + prompt_count: usize, + prefill_time: std::time::Duration, + decode_time: std::time::Duration, + ) -> (Vec, GenerationStats) { + let prefill_ms = prefill_time.as_secs_f64() * 1000.0; + let decode_ms = decode_time.as_secs_f64() * 1000.0; + let gen_count = self.generated_tokens.len(); + // The first token comes out of the prefill forward; the decode rate + // counts only what the decode loop produced in `decode_time`. + let decode_count = gen_count.saturating_sub(1); + let stats = GenerationStats { + prompt_tokens: prompt_count, + generated_tokens: gen_count, + prefill_time_ms: prefill_ms, + decode_time_ms: decode_ms, + prefill_tok_per_sec: if prefill_ms > 0.0 { + prompt_count as f64 / (prefill_ms / 1000.0) + } else { + 0.0 + }, + decode_tok_per_sec: if decode_ms > 0.0 { + decode_count as f64 / (decode_ms / 1000.0) + } else { + 0.0 + }, + }; + (self.generated_tokens.clone(), stats) + } +} + +#[cfg(test)] +#[path = "prompt_lookup_tests.rs"] +mod tests; diff --git a/src/lib/mlxcel-core/src/speculative/prompt_lookup_tests.rs b/src/lib/mlxcel-core/src/speculative/prompt_lookup_tests.rs new file mode 100644 index 000000000..df0297365 --- /dev/null +++ b/src/lib/mlxcel-core/src/speculative/prompt_lookup_tests.rs @@ -0,0 +1,313 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::*; + +fn cfg(ngram_max: usize, ngram_min: usize, max_draft: usize) -> PromptLookupConfig { + PromptLookupConfig { + ngram_max, + ngram_min, + max_draft, + adaptive: true, + } +} + +#[test] +fn proposes_the_tokens_after_an_earlier_occurrence() { + // "... 1 2 3 4 5 6 ... 1 2" -> propose "3 4 5". + let context = [9, 1, 2, 3, 4, 5, 6, 8, 1, 2]; + assert_eq!(find_draft(&context, &cfg(2, 1, 3)), vec![3, 4, 5]); +} + +#[test] +fn draft_is_capped_by_max_draft_and_context_end() { + let context = [1, 2, 3, 4, 1, 2]; + assert_eq!(find_draft(&context, &cfg(2, 1, 7)), vec![3, 4, 1, 2]); + assert_eq!(find_draft(&context, &cfg(2, 1, 1)), vec![3]); +} + +#[test] +fn prefers_the_longest_matching_suffix() { + // Suffix [5, 2]: the 2-gram matches at index 3, the 1-gram [2] matches + // later at index 6. The longer match must win. + let context = [0, 0, 0, 5, 2, 7, 2, 9, 5, 2]; + assert_eq!(find_draft(&context, &cfg(2, 1, 1)), vec![7]); +} + +#[test] +fn prefers_the_most_recent_occurrence_of_a_given_length() { + let context = [1, 2, 10, 1, 2, 20, 1, 2]; + assert_eq!(find_draft(&context, &cfg(2, 2, 1)), vec![20]); +} + +#[test] +fn never_matches_the_suffix_against_itself() { + // The only occurrence of [4] is the suffix itself. + assert!(find_draft(&[1, 2, 3, 4], &cfg(1, 1, 3)).is_empty()); +} + +#[test] +fn falls_back_to_shorter_ngrams() { + // [3, 4] never occurred before, [4] did. + let context = [4, 7, 8, 3, 4]; + assert_eq!(find_draft(&context, &cfg(2, 1, 2)), vec![7, 8]); + assert!(find_draft(&context, &cfg(2, 2, 2)).is_empty()); +} + +#[test] +fn short_and_empty_contexts_propose_nothing() { + let c = cfg(3, 1, 7); + assert!(find_draft(&[], &c).is_empty()); + assert!(find_draft(&[1], &c).is_empty()); + assert_eq!(find_draft(&[1, 1], &c), vec![1]); +} + +#[test] +fn repeated_run_proposes_the_run() { + // A run of one token: the latest occurrence of [7, 7] is one back, so + // exactly one token is available after it. + let context = [7, 7, 7, 7]; + assert_eq!(find_draft(&context, &cfg(2, 1, 5)), vec![7]); +} + +#[test] +fn config_validation() { + assert!(PromptLookupConfig::default().validate().is_ok()); + assert!(cfg(3, 0, 7).validate().is_err()); + assert!(cfg(1, 2, 7).validate().is_err()); + assert!(cfg(3, 1, 0).validate().is_err()); +} + +#[test] +fn tokens_per_forward_excludes_the_prefill_token() { + let stats = PromptLookupStats { + rounds: 4, + drafted_rounds: 4, + paused_rounds: 0, + proposed_draft_tokens: 28, + accepted_draft_tokens: 20, + }; + // 25 generated: 1 from prefill, 24 over 4 rounds. + assert!((stats.tokens_per_forward(25) - 6.0).abs() < 1e-9); + assert!((stats.acceptance_rate() - 20.0 / 28.0).abs() < 1e-9); + assert_eq!(PromptLookupStats::default().tokens_per_forward(1), 0.0); +} + +#[test] +fn governor_starts_with_the_full_block() { + let mut g = DraftGovernor::new(&cfg(3, 2, 7)); + assert_eq!(g.budget(), 7); +} + +#[test] +fn governor_keeps_the_full_block_while_blocks_land() { + let mut g = DraftGovernor::new(&cfg(3, 2, 7)); + for _ in 0..10 { + g.record(6); + } + assert_eq!(g.budget(), 7); +} + +#[test] +fn governor_shortens_the_block_when_few_tokens_land() { + let mut g = DraftGovernor::new(&cfg(3, 2, 7)); + for _ in 0..10 { + g.record(1); + } + // EMA converges on 1: twice that plus two. + assert_eq!(g.budget(), 4); +} + +/// Record one miss streak long enough to start a pause. +fn miss_streak(g: &mut DraftGovernor) { + for _ in 0..MISSES_BEFORE_COOLDOWN { + g.record(0); + } +} + +/// Count the paused rounds, ending on the first round that proposes again. +fn pause_len(g: &mut DraftGovernor) -> usize { + let mut paused = 0; + while g.budget() == 0 { + paused += 1; + } + paused +} + +#[test] +fn governor_pauses_after_repeated_misses_and_backs_off() { + let mut g = DraftGovernor::new(&cfg(3, 2, 7)); + for _ in 1..MISSES_BEFORE_COOLDOWN { + g.record(0); + } + assert!(g.budget() > 0, "a shorter run of misses is not a streak"); + g.record(0); + assert_eq!(pause_len(&mut g), MIN_COOLDOWN); + // A second streak pauses twice as long. + miss_streak(&mut g); + assert_eq!(pause_len(&mut g), 2 * MIN_COOLDOWN); +} + +#[test] +fn governor_backoff_resets_once_a_round_lands() { + let mut g = DraftGovernor::new(&cfg(3, 2, 7)); + for _ in 0..3 { + miss_streak(&mut g); + pause_len(&mut g); + } + g.record(2); + miss_streak(&mut g); + assert_eq!(pause_len(&mut g), MIN_COOLDOWN); +} + +#[test] +fn governor_backoff_is_capped() { + let mut g = DraftGovernor::new(&cfg(3, 2, 7)); + let mut last = 0; + for _ in 0..10 { + miss_streak(&mut g); + last = pause_len(&mut g); + } + assert_eq!(last, MAX_COOLDOWN); +} + +#[test] +fn non_adaptive_governor_always_proposes_max_draft() { + let mut c = cfg(3, 2, 7); + c.adaptive = false; + let mut g = DraftGovernor::new(&c); + for _ in 0..10 { + g.record(0); + assert_eq!(g.budget(), 7); + } +} + +#[test] +fn default_ngram_min_skips_single_token_matches() { + // Only a 1-gram ([4]) recurs; the default minimum of 2 proposes nothing. + let context = [4, 7, 8, 3, 4]; + assert!(find_draft(&context, &PromptLookupConfig::default()).is_empty()); +} + +/// Deterministic pseudo-random tokens from a small alphabet, so n-grams recur. +fn lcg_tokens(seed: u64, len: usize, alphabet: u64) -> Vec { + let mut state = seed; + (0..len) + .map(|_| { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + ((state >> 33) % alphabet) as i32 + }) + .collect() +} + +#[test] +fn index_matches_the_scan_at_every_length() { + for (seed, alphabet) in [(1, 3), (2, 5), (3, 11), (4, 40)] { + for config in [cfg(3, 1, 7), cfg(3, 2, 7), cfg(5, 2, 4), cfg(1, 1, 3)] { + let tokens = lcg_tokens(seed, 300, alphabet); + let mut index = NgramIndex::new(&config); + // Grow the context one token at a time, as decoding does. + for len in 0..=tokens.len() { + let context = &tokens[..len]; + index.extend(context); + assert_eq!( + index.find(context, &config), + find_draft(context, &config), + "seed {seed} alphabet {alphabet} {config:?} len {len}" + ); + } + } + } +} + +#[test] +fn index_matches_the_scan_when_extended_in_blocks() { + // A verify round emits several tokens before the next lookup. + let config = cfg(3, 2, 7); + let tokens = lcg_tokens(7, 400, 6); + let mut index = NgramIndex::new(&config); + let mut len = 50; + while len <= tokens.len() { + let context = &tokens[..len]; + index.extend(context); + assert_eq!(index.find(context, &config), find_draft(context, &config)); + len += 1 + len % 7; + } +} + +/// Target that always predicts token 3 at every position, so a reply is a run +/// of 3s: lookup proposes full blocks and the loop guard has a loop to catch. +struct ConstantModel; + +impl LanguageModel for ConstantModel { + fn forward( + &self, + input_ids: &ffi::MlxArray, + _caches: &mut [KVCache], + _mask: Option<&ffi::MlxArray>, + ) -> UniquePtr { + let seq_len = ffi::array_shape(input_ids)[1] as usize; + let mut logits = Vec::with_capacity(seq_len * 8); + for _ in 0..seq_len { + logits.extend_from_slice(&[0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0]); + } + ffi::from_slice_f32(&logits, &[1, seq_len as i32, 8]) + } + + fn make_caches(&self) -> Vec { + vec![KVCache::new()] + } + + fn num_layers(&self) -> usize { + 1 + } + + fn eos_token_ids(&self) -> Vec { + vec![7] + } +} + +#[test] +fn loop_guard_stops_mid_block_where_plain_decoding_stops() { + let sampling = SamplingConfig { + loop_detection: crate::loop_detection::LoopDetectionConfig::new(1, 1, 8), + ..SamplingConfig::greedy() + }; + let prompt = [1, 2, 3, 3]; + let mut plain = crate::generate::CxxGenerator::new(1); + let expected = plain.generate(&ConstantModel, &prompt, 64, &sampling); + assert_eq!( + expected, + vec![3; 8], + "plain decoding stops on the 8th repeat" + ); + + let mut generator = PromptLookupGenerator::new(PromptLookupConfig::default()); + let (tokens, _) = generator.generate(&ConstantModel, &prompt, 64, &sampling); + assert_eq!(tokens, expected); + assert!( + generator.stats().accepted_draft_tokens > 0, + "the run verified lookup blocks" + ); +} + +#[test] +fn without_the_loop_guard_lookup_runs_to_max_tokens() { + let mut generator = PromptLookupGenerator::new(PromptLookupConfig::default()); + let (tokens, _) = + generator.generate(&ConstantModel, &[1, 2, 3, 3], 64, &SamplingConfig::greedy()); + assert_eq!(tokens, vec![3; 64]); +} diff --git a/src/main.rs b/src/main.rs index 0ccab5369..d4cf585d3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -370,6 +370,61 @@ pub(crate) struct RmArgs { pub(crate) revision: Option, } +/// Prompt-lookup (n-gram) speculative decoding options +#[derive(Args, Debug, Clone)] +#[command(next_help_heading = "Prompt Lookup Options")] +pub(crate) struct PromptLookupOptions { + /// Speculate with prompt lookup: propose the tokens that followed an + /// earlier occurrence of the sequence's last few tokens, and verify them + /// in one target forward. No draft model is needed. Pays off on replies + /// that copy the prompt (edits, rewrites, reformatting); on other replies + /// proposals are shortened or paused so decoding stays near plain speed. + /// Text-only, attention-cache models. + #[arg(long, conflicts_with = "draft_model")] + pub(crate) prompt_lookup: bool, + + /// Longest suffix n-gram matched against the context + #[arg(long, default_value_t = mlxcel::PromptLookupConfig::default().ngram_max, value_name = "N")] + pub(crate) prompt_lookup_ngram_max: usize, + + /// Shortest suffix n-gram matched against the context + #[arg(long, default_value_t = mlxcel::PromptLookupConfig::default().ngram_min, value_name = "N")] + pub(crate) prompt_lookup_ngram_min: usize, + + /// Maximum tokens proposed per round (the verify block is one wider) + #[arg(long, default_value_t = mlxcel::PromptLookupConfig::default().max_draft, value_name = "N")] + pub(crate) prompt_lookup_max_draft: usize, + + /// Propose up to the maximum every round, instead of shortening or pausing + /// proposals while they stop landing (for A/B measurement) + #[arg(long)] + pub(crate) prompt_lookup_no_adaptive: bool, +} + +impl Default for PromptLookupOptions { + fn default() -> Self { + let config = mlxcel::PromptLookupConfig::default(); + Self { + prompt_lookup: false, + prompt_lookup_ngram_max: config.ngram_max, + prompt_lookup_ngram_min: config.ngram_min, + prompt_lookup_max_draft: config.max_draft, + prompt_lookup_no_adaptive: !config.adaptive, + } + } +} + +impl PromptLookupOptions { + pub(crate) fn config(&self) -> mlxcel::PromptLookupConfig { + mlxcel::PromptLookupConfig { + ngram_max: self.prompt_lookup_ngram_max, + ngram_min: self.prompt_lookup_ngram_min, + max_draft: self.prompt_lookup_max_draft, + adaptive: !self.prompt_lookup_no_adaptive, + } + } +} + #[derive(Args, Debug)] pub(crate) struct GenerateArgs { #[command(flatten)] @@ -398,6 +453,9 @@ pub(crate) struct GenerateArgs { #[command(flatten)] pub(crate) speculative: SpeculativeArgs, + #[command(flatten)] + pub(crate) prompt_lookup: PromptLookupOptions, + // Axis A weight-load surgery configuration. // The closed-repo references stay in this non-doc comment so the // user-facing `--help` block does not advertise tracker URLs (see From a1fa9e704f81d0ab3cddb88dacbcdc2dfd715069 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 2 Oct 2026 11:49:48 +0900 Subject: [PATCH 2/2] fix(speculative): harden prompt-lookup rollback gate, limits and CLI Prompt lookup rolls rejected proposals back with `KVCache::trim`, so a target must keep its whole sequence state in the external caches. The gate inferred that from `supports_batching` and `supports_padded_prefill`; Gemma 4 and Llama 4 batch yet own their caches, and were refused only through the padded-prefill flag, which answers a different question (#1335). The gate now refuses a `ModelOwned` layout directly, and the generator asserts non-empty external caches (`can_trim_prompt_cache` is vacuously true on an empty slice). - Warm every verify width before the timed CLI run. The warmup generation reached only the widths its own proposals used, so on prose the timed run paid first-use kernel costs mid-decode: a 72-token Qwen3-1.7B reply on GB10 went from 0.65x to 1.07x of plain decoding. - Bound `ngram-max` at 16 and `max-draft` at 64; the n-gram index grows with the square of the range times the context. - Sample history-reading configs through the incremental `SamplerState`; refuse mirostat and adaptive-p, whose state the loop does not carry. - Tuning flags require `--prompt-lookup`; a bad configuration or multimodal input is refused before the checkpoint is resolved. - Rollback tests drive a cache-dependent target through accepted, rejected and pipelined rounds; six deliberate mutations of the trim and in-flight logic each fail them. The penalty case uses a sequential reference because `CxxGenerator` samples against a history one token stale (#2090). - Document the acceptance rule in `docs/speculative-acceptance.md`. Refs #2074 --- docs/speculative-acceptance.md | 8 + src/commands/generate.rs | 33 +- src/commands/generate_tests.rs | 45 ++- .../src/speculative/prompt_lookup.rs | 125 ++++++- .../src/speculative/prompt_lookup_tests.rs | 325 ++++++++++++++++++ src/main.rs | 23 +- src/main_tests.rs | 93 ++++- 7 files changed, 633 insertions(+), 19 deletions(-) diff --git a/docs/speculative-acceptance.md b/docs/speculative-acceptance.md index 53a72eca9..573c4ed3d 100644 --- a/docs/speculative-acceptance.md +++ b/docs/speculative-acceptance.md @@ -111,6 +111,7 @@ real and provable. | Path | `temperature == 0` | `temperature > 0` default | `temperature > 0` opt-in | |------|--------------------|---------------------------|--------------------------| | `SpeculativeGenerator` (classic; `mlxcel generate --draft-model`) | greedy argmax (lossless) | **sampler-match** (lossless) | modified rejection sampling (lossless, acceptance-optimal) | +| `PromptLookupGenerator` (`mlxcel generate --prompt-lookup`) | greedy argmax (lossless) | **sampler-match** (lossless, and acceptance-optimal here) | none needed | | Gemma 4 MTP round loop | argmax (lossless here) | argmax-against-argmax (**biased**) | not wired | | DFlash round loop (Qwen 3.5 DFlash drafter) | argmax (lossless here) | argmax-against-argmax (**biased**) | not wired | | DFlash round loop (LFM2 DSpark drafter) | argmax (lossless here, probe-gated) | declines to classic decode | not wired | @@ -125,6 +126,13 @@ Those are the paths where this feature would buy correctness rather than only acceptance rate, and where the acceptance trade against argmax is a real decision rather than a free win. +Prompt lookup proposes deterministically, so its `q` is one-hot on the drafted +token `d`. The two acceptance probabilities then coincide at `p(d)` +(`sum_x p(x) q(x) = sum_x min(p(x), q(x)) = p(d)`), and on a rejection both +rules emit a draw from `p` conditioned on not being `d`. Sampler-match is +therefore already the acceptance-optimal rule for that path, and there is +nothing to opt into. + ## RNG dependency Modified rejection sampling draws randomness the previous rules did not: diff --git a/src/commands/generate.rs b/src/commands/generate.rs index 6015463ee..eadf2c4a0 100644 --- a/src/commands/generate.rs +++ b/src/commands/generate.rs @@ -1972,13 +1972,34 @@ pub(super) fn run_generation_mode( Ok(output) } +/// Refuse a `--prompt-lookup` run that cannot start, before the checkpoint is +/// resolved or loaded: an invalid lookup configuration, or multimodal input, +/// whose embeddings the generator cannot take. A no-op without the flag. +pub(super) fn validate_prompt_lookup_args(args: &GenerateArgs) -> Result<()> { + if !args.prompt_lookup.prompt_lookup { + return Ok(()); + } + args.prompt_lookup + .config() + .validate() + .map_err(|err| anyhow!("--prompt-lookup: {err}"))?; + ensure!( + args.generation.image.is_empty() + && args.generation.audio.is_none() + && args.generation.video.is_empty(), + "--prompt-lookup supports text-only prompts; drop --image/--audio/--video" + ); + Ok(()) +} + /// Offline prompt-lookup speculative decoding (`--prompt-lookup`). /// /// Refuses the model families whose state cannot be rewound after a rejected /// block, and multimodal prompts, whose embeddings the generator cannot take. /// Runs one short warmup generation first, as `generate_standard` does, so the -/// timed run does not pay for kernel compilation. The warmup decodes a few -/// rounds rather than one token so the verify widths get compiled too. +/// timed run does not pay for kernel compilation, then one forward at every +/// verify width so the widths the warmup's own proposals missed are compiled +/// too. /// /// The returned stats time the whole call, prefill included, exactly as /// `generate_standard` does, so the printed tok/s compares like for like. @@ -2008,6 +2029,10 @@ fn run_prompt_lookup( .with_token_bias(token_bias); let warmup_tokens = args.generation.max_tokens.min(16); let _ = generator.generate(model, prompt_tokens, warmup_tokens, sampling_config); + // The warmup generation compiles only the verify widths its own proposals + // used; a reply with nothing to copy uses none, and the timed run would + // then pay each width's first-use cost mid-decode. + generator.warm_up_verify_widths(model, prompt_tokens); let start_time = Instant::now(); let (tokens, measured) = generator.generate( model, @@ -2663,6 +2688,10 @@ fn run_generate_once(mut args: GenerateArgs) -> Result<()> { #[cfg(feature = "surgery")] install_surgery_pipeline_from_cli(&args)?; + // Model-independent too, so a bad `--prompt-lookup` combination fails + // before the resolver below can auto-download the checkpoint. + validate_prompt_lookup_args(&args)?; + // Resolve `-m` into a concrete model directory (epic #92, issue #94) // before any consumer reads it. An existing path is used verbatim // (byte-identical to the pre-#94 local-path behavior); an `owner/name` diff --git a/src/commands/generate_tests.rs b/src/commands/generate_tests.rs index fd79e511b..a20424bbe 100644 --- a/src/commands/generate_tests.rs +++ b/src/commands/generate_tests.rs @@ -21,7 +21,7 @@ use super::{ generated_suffix, generation_stats_from_duration, memory_preflight_ctx_len, reject_dflash_drafter_offline, resolve_cli_pipeline_assignments, resolve_cli_prompt, should_route_offline_mtp, strip_trailing_eos, validate_muse_glimmer_cli_unsupported_options, - validate_pipeline_parallel_args, validate_tensor_parallel_args, + validate_pipeline_parallel_args, validate_prompt_lookup_args, validate_tensor_parallel_args, validate_xla_cli_image_cardinality, validate_xla_output_audio, }; use mlxcel::server::chat_template::{ChatMessage, ChatTemplateProcessor, flatten_template_text}; @@ -965,6 +965,9 @@ fn validate_pipeline_parallel_args_rejects_incompatible_modes() { args.model.draft_model = Some(PathBuf::from("draft")); assert!(validate_pipeline_parallel_args(&args).is_err()); args.model.draft_model = None; + args.prompt_lookup.prompt_lookup = true; + assert!(validate_pipeline_parallel_args(&args).is_err()); + args.prompt_lookup.prompt_lookup = false; // Tensor parallelism + PP is now accepted (2D PP × TP composition landed). Positive coverage for the 2D path lives in // `validate_pipeline_parallel_args_accepts_2d_pp_tp` below. @@ -1418,3 +1421,43 @@ fn cli_video_frames_are_private_and_removed_with_the_run() { } assert!(!dir.exists(), "the frame directory outlived the run"); } + +#[test] +fn validate_prompt_lookup_args_is_a_no_op_without_the_flag() { + let mut args = sample_generate_args(temp_model_dir("pl-disabled")); + args.generation.image = vec![PathBuf::from("page.png")]; + args.prompt_lookup.prompt_lookup_ngram_min = 0; + assert!(validate_prompt_lookup_args(&args).is_ok()); + fs::remove_dir_all(args.model.model).unwrap(); +} + +#[test] +fn validate_prompt_lookup_args_refuses_before_the_model_loads() { + let mut args = sample_generate_args(temp_model_dir("pl-refusals")); + args.prompt_lookup.prompt_lookup = true; + assert!(validate_prompt_lookup_args(&args).is_ok()); + + args.prompt_lookup.prompt_lookup_ngram_min = 0; + let err = validate_prompt_lookup_args(&args).unwrap_err().to_string(); + assert!(err.contains("ngram-min"), "{err}"); + args.prompt_lookup.prompt_lookup_ngram_min = 2; + + args.prompt_lookup.prompt_lookup_ngram_max = + mlxcel_core::speculative::prompt_lookup::NGRAM_MAX_LIMIT + 1; + let err = validate_prompt_lookup_args(&args).unwrap_err().to_string(); + assert!(err.contains("at most"), "{err}"); + args.prompt_lookup.prompt_lookup_ngram_max = 3; + + for media in ["image", "audio", "video"] { + let mut args = sample_generate_args(args.model.model.clone()); + args.prompt_lookup.prompt_lookup = true; + match media { + "image" => args.generation.image = vec![PathBuf::from("page.png")], + "audio" => args.generation.audio = Some(PathBuf::from("clip.wav")), + _ => args.generation.video = vec![PathBuf::from("clip.mp4")], + } + let err = validate_prompt_lookup_args(&args).unwrap_err().to_string(); + assert!(err.contains("text-only"), "{media}: {err}"); + } + fs::remove_dir_all(args.model.model).unwrap(); +} diff --git a/src/lib/mlxcel-core/src/speculative/prompt_lookup.rs b/src/lib/mlxcel-core/src/speculative/prompt_lookup.rs index 8daf8932a..56d90e81f 100644 --- a/src/lib/mlxcel-core/src/speculative/prompt_lookup.rs +++ b/src/lib/mlxcel-core/src/speculative/prompt_lookup.rs @@ -88,7 +88,7 @@ //! state that cannot be trimmed; callers must refuse them, see //! [`prompt_lookup_unsupported_reason`]. -use crate::cache::{KVCacheMode, can_trim_prompt_cache}; +use crate::cache::{KVCacheMode, SequenceStateBackend, can_trim_prompt_cache}; use crate::drafter::dflash::materialize_argmax_i32_vec; use crate::ffi; use crate::ffi::MlxThreadLocalStream; @@ -99,7 +99,9 @@ use crate::generate::{ use crate::generation_policy::{initial_token_history, merged_eos_token_ids, seed_rng_if_needed}; use crate::layers::KVCache; use crate::loop_detection::detect_repetition_loop; -use crate::sampling::{TokenBiasMap, sample_token_optimized}; +use crate::sampling::{ + SamplerState, TokenBiasMap, sample_token_optimized, sample_token_optimized_with_state, +}; use crate::streams::{install_thread_local_default_stream, shared_thread_local_generation_stream}; use cxx::UniquePtr; use std::borrow::Cow; @@ -119,6 +121,17 @@ pub const DEFAULT_NGRAM_MAX: usize = 3; pub const DEFAULT_NGRAM_MIN: usize = 2; /// Default maximum number of tokens proposed per round (verify block of 8). pub const DEFAULT_MAX_DRAFT: usize = 7; +/// Largest accepted `ngram_max`. +/// +/// [`NgramIndex`] keeps one map per n-gram length in `ngram_min..=ngram_max`, +/// and each map holds a boxed `n`-token key for every context position, so its +/// host memory grows with the square of the range times the context length. +/// Without a bound, a mistyped `--prompt-lookup-ngram-max 1000` over a long +/// document asks for hundreds of gigabytes before the first decode step. +pub const NGRAM_MAX_LIMIT: usize = 16; +/// Largest accepted `max_draft`. The verify block is one wider, and every +/// proposal is a forward position the target computes. +pub const MAX_DRAFT_LIMIT: usize = 64; /// Tunables for [`PromptLookupGenerator`]. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -157,9 +170,21 @@ impl PromptLookupConfig { self.ngram_max, self.ngram_min )); } + if self.ngram_max > NGRAM_MAX_LIMIT { + return Err(format!( + "prompt-lookup ngram-max ({}) must be at most {NGRAM_MAX_LIMIT}", + self.ngram_max + )); + } if self.max_draft == 0 { return Err("prompt-lookup max-draft must be at least 1".to_string()); } + if self.max_draft > MAX_DRAFT_LIMIT { + return Err(format!( + "prompt-lookup max-draft ({}) must be at most {MAX_DRAFT_LIMIT}", + self.max_draft + )); + } Ok(()) } } @@ -169,10 +194,19 @@ impl PromptLookupConfig { /// /// Verification rolls rejected tokens back with `KVCache::trim`, which is only /// sound when every piece of per-sequence state lives in the external caches. -/// `supports_batching` is the existing marker for that, and +/// `supports_batching` is the existing marker for that, and a +/// [`SequenceStateBackend::ModelOwned`] layout is the direct statement of the +/// opposite: such a family keeps its caches inside the model, the external +/// caches are empty placeholders, and trimming them rolls nothing back. /// `supports_padded_prefill` is `false` for the families whose state trailing /// tokens corrupt past what a trim undoes, which is the same hazard a rejected /// draft block poses. +/// +/// The layout check keeps the refusal from resting on the other two flags: +/// Gemma 4 and Llama 4 report `supports_batching() == true` and were refused +/// only because they opt out of padded prefill, a flag that answers a +/// different question (the scheduler's post-pad trim, issue #1335) and may +/// change once that is solved. pub fn prompt_lookup_unsupported_reason(model: &M) -> Option<&'static str> { if !model.supports_batching() { return Some( @@ -181,6 +215,12 @@ pub fn prompt_lookup_unsupported_reason(model: &M) -> Option<& cannot be rolled back", ); } + if model.sequence_state_layout().backend == SequenceStateBackend::ModelOwned { + return Some( + "this model owns its attention caches internally rather than in the external KV \ + caches the generator trims, so a rejected draft cannot be rolled back", + ); + } if !model.supports_padded_prefill() { return Some( "this model's state is corrupted by trailing tokens that a cache trim cannot \ @@ -487,6 +527,43 @@ impl PromptLookupGenerator { self.stats } + /// Prefill `prompt_tokens` and run one verify forward at every block + /// width [`Self::generate`] can use (2 through `max_draft + 1`), so a + /// later call does not pay first-use costs in the middle of its decode. + /// + /// The first forward at a width this process has not run yet builds or + /// traces kernels for it, which on GB10 cost more than the rest of a + /// short reply: a 72-token Qwen3-1.7B answer with two drafted rounds ran + /// at 0.65x plain decoding without this warmup and at 1.07x with it. A + /// warmup generation only reaches the widths its own proposals happen to + /// use, which on a reply with nothing to copy is none of them. + /// + /// # Panics + /// + /// Same contract as [`Self::generate`] for the model and the prompt. + pub fn warm_up_verify_widths(&mut self, model: &M, prompt_tokens: &[i32]) { + self.reset(model); + assert!( + !self.caches.is_empty() && can_trim_prompt_cache(&self.caches), + "prompt-lookup decoding requires trimmable external KV caches" + ); + let last = *prompt_tokens + .last() + .expect("prompt-lookup warmup requires at least one prompt token"); + install_thread_local_default_stream(self.generation_stream.as_ref()); + let logits = prefill_prompt_last_logits(model, &mut self.caches, prompt_tokens); + ffi::eval(&logits); + for width in 2..=self.config.max_draft + 1 { + let tokens = vec![last; width]; + let input = ffi::from_slice_i32(&tokens, &[1, width as i32]); + let logits = model.forward(&input, &mut self.caches, None); + let argmax = ffi::argmax_last_axis(&logits); + ffi::eval(&argmax); + trim_caches(&mut self.caches, width as i32); + } + ffi::clear_memory_cache(); + } + /// Fresh caches from the model, in the configured mode, as /// `CxxGenerator::reset_with_model` builds them. fn reset(&mut self, model: &M) { @@ -521,8 +598,12 @@ impl PromptLookupGenerator { /// /// # Panics /// - /// Panics on an empty prompt or a non-trimmable cache. Callers must check - /// [`supports_prompt_lookup`] first. + /// Panics on an empty prompt, on a model with no external KV caches or a + /// non-trimmable one (callers must check [`supports_prompt_lookup`] + /// first), and on a sampler that carries feedback state across tokens + /// (mirostat, adaptive-p), which this loop does not thread through its + /// rounds; the server's speculative burst declines the same samplers. + /// `mlxcel generate` cannot configure either one. pub fn generate( &mut self, model: &M, @@ -532,9 +613,11 @@ impl PromptLookupGenerator { ) -> (Vec, GenerationStats) { self.reset(model); + // `can_trim_prompt_cache` is vacuously true on an empty slice, which + // is exactly what a family that keeps its caches internally returns. assert!( - can_trim_prompt_cache(&self.caches), - "prompt-lookup decoding requires a trimmable prompt cache" + !self.caches.is_empty() && can_trim_prompt_cache(&self.caches), + "prompt-lookup decoding requires trimmable external KV caches" ); assert!( !prompt_tokens.is_empty(), @@ -543,12 +626,22 @@ impl PromptLookupGenerator { let sampling_cow = self.compose_sampling(sampling); let sampling: &SamplingConfig = sampling_cow.as_ref(); + assert!( + !sampling.needs_sampler_feedback_state(), + "prompt-lookup decoding does not carry mirostat or adaptive-p sampler state" + ); seed_rng_if_needed(sampling); install_thread_local_default_stream(self.generation_stream.as_ref()); let eos_tokens = merged_eos_token_ids(model.eos_token_ids(), &sampling.stop_token_ids); let needs_history = sampling.needs_token_history(); let mut token_history = initial_token_history(prompt_tokens, needs_history); + // Incremental penalty state, as `CxxGenerator` keeps it: byte-identical + // logits to rebuilding the penalties from `token_history` on every + // sample, without the per-sample pass over the whole history. History + // only grows here (rejected proposals are never sampled into it), so + // the state's append-only fast path always applies. + let mut sampler_state: Option = None; // Pure argmax with nothing that depends on the emitted prefix: every // verify position can be decided from one batched argmax and a single // host read. Anything else goes position by position through the @@ -563,7 +656,11 @@ impl PromptLookupGenerator { let prefill_start = Instant::now(); let logits = prefill_prompt_last_logits(model, &mut self.caches, prompt_tokens); ffi::clear_memory_cache(); - let (first_arr, _) = sample_token_optimized(&logits, sampling, &token_history); + let (first_arr, _) = if needs_history { + sample_token_optimized_with_state(&logits, sampling, &token_history, &mut sampler_state) + } else { + sample_token_optimized(&logits, sampling, &token_history) + }; ffi::async_eval(&first_arr); // Index the prompt while the GPU finishes the prefill. let mut index = NgramIndex::new(&self.config); @@ -685,8 +782,16 @@ impl PromptLookupGenerator { None => { let pos_logits = ffi::slice(&logits, &[0, pos as i32, 0], &[1, pos as i32 + 1, vocab]); - let (arr, _) = - sample_token_optimized(&pos_logits, sampling, &token_history); + let (arr, _) = if needs_history { + sample_token_optimized_with_state( + &pos_logits, + sampling, + &token_history, + &mut sampler_state, + ) + } else { + sample_token_optimized(&pos_logits, sampling, &token_history) + }; ffi::eval(&arr); ffi::item_i32(&arr) } diff --git a/src/lib/mlxcel-core/src/speculative/prompt_lookup_tests.rs b/src/lib/mlxcel-core/src/speculative/prompt_lookup_tests.rs index df0297365..57c9aea1b 100644 --- a/src/lib/mlxcel-core/src/speculative/prompt_lookup_tests.rs +++ b/src/lib/mlxcel-core/src/speculative/prompt_lookup_tests.rs @@ -311,3 +311,328 @@ fn without_the_loop_guard_lookup_runs_to_max_tokens() { generator.generate(&ConstantModel, &[1, 2, 3, 3], 64, &SamplingConfig::greedy()); assert_eq!(tokens, vec![3; 64]); } + +#[test] +fn config_validation_bounds_the_index_and_the_verify_block() { + assert!(cfg(NGRAM_MAX_LIMIT, 2, MAX_DRAFT_LIMIT).validate().is_ok()); + let err = cfg(NGRAM_MAX_LIMIT + 1, 2, 7).validate().unwrap_err(); + assert!(err.contains("ngram-max"), "{err}"); + let err = cfg(3, 2, MAX_DRAFT_LIMIT + 1).validate().unwrap_err(); + assert!(err.contains("max-draft"), "{err}"); +} + +/// Vocabulary of [`InductionModel`]. +const INDUCTION_VOCAB: usize = 6; + +/// Target whose prediction depends on everything in its KV cache, so a +/// rollback that leaves a rejected proposal behind, or trims one token too +/// many, changes what it predicts from then on. +/// +/// Each forward appends its input tokens to the cache as key values and, at +/// every new position, predicts the token that followed an earlier occurrence +/// of the token at that position (a one-token induction head over the cached +/// sequence), or `token + 1` when there is none. The following token gets a +/// runner-up logit, so a repetition penalty can flip the choice. +/// +/// With `earliest: false` it copies from the most recent occurrence, the one +/// lookup also prefers, so most proposals land and the governor keeps full +/// blocks. With `earliest: true` it copies from the first occurrence, so +/// proposals keep missing, the governor pauses, and the decode loop pipelines +/// runs of plain rounds and switches back to verifying when a pause ends. +struct InductionModel { + earliest: bool, +} + +const INDUCTION_MODELS: [InductionModel; 2] = [ + InductionModel { earliest: false }, + InductionModel { earliest: true }, +]; + +impl LanguageModel for InductionModel { + fn forward( + &self, + input_ids: &ffi::MlxArray, + caches: &mut [KVCache], + _mask: Option<&ffi::MlxArray>, + ) -> UniquePtr { + let seq_len = ffi::array_shape(input_ids)[1]; + let as_f16 = ffi::astype(input_ids, crate::dtype::FLOAT16); + let keys = ffi::reshape(&as_f16, &[1, 1, seq_len, 1]); + let values = ffi::reshape(&as_f16, &[1, 1, seq_len, 1]); + let (window, _) = caches[0].update_and_fetch(keys, values); + let sequence: Vec = crate::utils::array_to_vec_f32(&window) + .into_iter() + .map(|v| v as i32) + .collect(); + let new_len = seq_len as usize; + let mut logits = vec![0.0f32; new_len * INDUCTION_VOCAB]; + for (row, pos) in (sequence.len() - new_len..sequence.len()).enumerate() { + let token = sequence[pos]; + let mut earlier = (0..pos).filter(|&i| sequence[i] == token); + let source = if self.earliest { + earlier.next() + } else { + earlier.next_back() + }; + let next = + source.map_or((token + 1) % INDUCTION_VOCAB as i32, |i| sequence[i + 1]) as usize; + logits[row * INDUCTION_VOCAB + next] = 8.0; + logits[row * INDUCTION_VOCAB + (next + 1) % INDUCTION_VOCAB] = 7.0; + } + ffi::from_slice_f32(&logits, &[1, seq_len, INDUCTION_VOCAB as i32]) + } + + fn make_caches(&self) -> Vec { + vec![KVCache::new()] + } + + fn num_layers(&self) -> usize { + 1 + } + + fn eos_token_ids(&self) -> Vec { + Vec::new() + } +} + +/// Plain decoding without pipelining: one forward per token, each sampled +/// against a history that already holds every emitted token, which is what a +/// history-reading sampler is defined against (mlx-lm's `generate_step` hands +/// its logits processors the token it just consumed). +/// +/// `CxxGenerator` cannot be the reference once the sampler reads history: its +/// pipelined loop builds the next step's sample before it pushes the token it +/// has just read, so its penalties see a history one token stale. +fn sequential_reference( + model: &M, + prompt: &[i32], + max_tokens: usize, + sampling: &SamplingConfig, +) -> Vec { + let mut caches = model.make_caches(); + let mut history = prompt.to_vec(); + let mut input = prompt.to_vec(); + let mut emitted = Vec::with_capacity(max_tokens); + while emitted.len() < max_tokens { + let ids = ffi::from_slice_i32(&input, &[1, input.len() as i32]); + let logits = model.forward(&ids, &mut caches, None); + let (token, _) = sample_token_optimized(&logits, sampling, &history); + let token = ffi::item_i32(&token); + emitted.push(token); + history.push(token); + input = vec![token]; + } + emitted +} + +/// Run prompt lookup on both [`INDUCTION_MODELS`] over several prompts and +/// configurations, assert every reply equals `reference(model, prompt)`, and +/// return the summed acceptance accounting per model. +fn induction_parity( + sampling: &SamplingConfig, + reference: impl Fn(&InductionModel, &[i32]) -> Vec, +) -> [PromptLookupStats; 2] { + let mut totals = [PromptLookupStats::default(); 2]; + for (model, total) in INDUCTION_MODELS.iter().zip(&mut totals) { + for seed in 1..=6 { + let prompt = lcg_tokens(seed, 48, INDUCTION_VOCAB as u64); + let expected = reference(model, &prompt); + assert_eq!(expected.len(), 96, "seed {seed}: the reference ran short"); + for config in [ + PromptLookupConfig::default(), + cfg(3, 1, 7), + cfg(4, 2, 3), + // Six-token matches are rare in a random prompt and appear + // once the reply repeats itself, so the loop first pipelines a + // long run of plain rounds and then switches to verifying with + // a step still in flight. + cfg(6, 6, 7), + PromptLookupConfig { + adaptive: false, + ..PromptLookupConfig::default() + }, + ] { + let mut generator = PromptLookupGenerator::new(config); + let (tokens, _) = generator.generate(model, &prompt, 96, sampling); + assert_eq!( + tokens, expected, + "earliest={} seed {seed} {config:?}", + model.earliest + ); + let stats = generator.stats(); + total.rounds += stats.rounds; + total.drafted_rounds += stats.drafted_rounds; + total.paused_rounds += stats.paused_rounds; + total.proposed_draft_tokens += stats.proposed_draft_tokens; + total.accepted_draft_tokens += stats.accepted_draft_tokens; + } + } + } + totals +} + +/// Both acceptance regimes showed up: blocks that land and blocks that are +/// trimmed, and governor pauses long enough for the loop to pipeline. +fn assert_both_regimes(totals: &[PromptLookupStats; 2]) { + for stats in totals { + assert!( + stats.accepted_draft_tokens > 0, + "no proposal landed: {stats:?}" + ); + assert!( + stats.proposed_draft_tokens > stats.accepted_draft_tokens, + "no proposal was rejected, so no trim was exercised: {stats:?}" + ); + } + let missing = &totals[1]; + assert!( + missing.paused_rounds > 0, + "the missing model never paused the governor: {missing:?}" + ); +} + +/// The rollback contract: after every verify round the caches must hold +/// exactly the emitted tokens, so a cache-dependent target produces plain +/// decoding's reply token for token, through accepted blocks, rejected tails, +/// and the pipelined plain rounds between them. +#[test] +fn rollback_keeps_a_cache_dependent_target_identical_to_plain_decoding() { + let sampling = SamplingConfig::greedy(); + let totals = induction_parity(&sampling, |model, prompt| { + let plain = crate::generate::CxxGenerator::new(1).generate(model, prompt, 96, &sampling); + assert_eq!( + plain, + sequential_reference(model, prompt, 96, &sampling), + "greedy plain decoding is the sequential reference" + ); + plain + }); + assert_both_regimes(&totals); +} + +/// Same contract on the per-position sampler path: a repetition penalty +/// reads the emitted history, so every verify position samples on its own +/// through the incremental sampler state, and each must see every token +/// emitted before it, including the ones accepted earlier in the same block. +#[test] +fn rollback_matches_plain_decoding_under_a_repetition_penalty() { + let sampling = SamplingConfig { + repetition_penalty: 1.3, + penalty_last_n: 3, + ..SamplingConfig::greedy() + }; + assert!(sampling.needs_token_history()); + let totals = induction_parity(&sampling, |model, prompt| { + sequential_reference(model, prompt, 96, &sampling) + }); + assert_both_regimes(&totals); +} + +/// A dense-cache target that nonetheless owns its sequence state, the shape +/// Gemma 4 and Llama 4 have: it batches and accepts padded prefill, so only +/// the layout says the external caches are placeholders. +struct ModelOwnedStateModel; + +impl LanguageModel for ModelOwnedStateModel { + fn forward( + &self, + input_ids: &ffi::MlxArray, + caches: &mut [KVCache], + mask: Option<&ffi::MlxArray>, + ) -> UniquePtr { + ConstantModel.forward(input_ids, caches, mask) + } + + fn make_caches(&self) -> Vec { + Vec::new() + } + + fn num_layers(&self) -> usize { + 1 + } + + fn eos_token_ids(&self) -> Vec { + vec![7] + } + + fn sequence_state_layout(&self) -> crate::cache::SequenceStateLayout { + crate::cache::SequenceStateLayout::model_owned(1) + } +} + +#[test] +fn model_owned_sequence_state_is_refused_even_when_other_flags_allow_it() { + let model = ModelOwnedStateModel; + assert!(model.supports_batching() && model.supports_padded_prefill()); + let reason = prompt_lookup_unsupported_reason(&model).expect("must be refused"); + assert!(reason.contains("internally"), "{reason}"); + assert!(supports_prompt_lookup(&ConstantModel)); +} + +#[test] +#[should_panic(expected = "trimmable external KV caches")] +fn generate_refuses_a_model_without_external_caches() { + let mut generator = PromptLookupGenerator::new(PromptLookupConfig::default()); + let _ = generator.generate( + &ModelOwnedStateModel, + &[1, 2, 3], + 8, + &SamplingConfig::greedy(), + ); +} + +#[test] +#[should_panic(expected = "mirostat or adaptive-p")] +fn generate_refuses_a_sampler_with_feedback_state() { + let sampling = SamplingConfig { + mirostat: 2, + ..SamplingConfig::with_temperature(0.8) + }; + let mut generator = PromptLookupGenerator::new(PromptLookupConfig::default()); + let _ = generator.generate(&ConstantModel, &[1, 2, 3], 8, &sampling); +} + +/// Records the width of every forward and keeps the tokens in its cache, so +/// the warmup's trims are observable. +struct WidthRecordingModel { + widths: std::cell::RefCell>, +} + +impl LanguageModel for WidthRecordingModel { + fn forward( + &self, + input_ids: &ffi::MlxArray, + caches: &mut [KVCache], + _mask: Option<&ffi::MlxArray>, + ) -> UniquePtr { + let seq_len = ffi::array_shape(input_ids)[1]; + self.widths.borrow_mut().push(seq_len); + let kv = ffi::ones(&[1, 1, seq_len, 1], crate::dtype::FLOAT32); + caches[0].update(kv, ffi::ones(&[1, 1, seq_len, 1], crate::dtype::FLOAT32)); + ConstantModel.forward(input_ids, &mut [], None) + } + + fn make_caches(&self) -> Vec { + vec![KVCache::new()] + } + + fn num_layers(&self) -> usize { + 1 + } + + fn eos_token_ids(&self) -> Vec { + vec![7] + } +} + +#[test] +fn warmup_runs_every_verify_width_and_rolls_each_back() { + let model = WidthRecordingModel { + widths: std::cell::RefCell::new(Vec::new()), + }; + let prompt = [1, 2, 3, 4, 5]; + let mut generator = PromptLookupGenerator::new(cfg(3, 2, 4)); + generator.warm_up_verify_widths(&model, &prompt); + assert_eq!(*model.widths.borrow(), vec![5, 2, 3, 4, 5]); + assert_eq!(generator.caches[0].offset, prompt.len() as i32); +} diff --git a/src/main.rs b/src/main.rs index 03d49af27..728591cb3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -384,20 +384,35 @@ pub(crate) struct PromptLookupOptions { pub(crate) prompt_lookup: bool, /// Longest suffix n-gram matched against the context - #[arg(long, default_value_t = mlxcel::PromptLookupConfig::default().ngram_max, value_name = "N")] + #[arg( + long, + requires = "prompt_lookup", + default_value_t = mlxcel::PromptLookupConfig::default().ngram_max, + value_name = "N" + )] pub(crate) prompt_lookup_ngram_max: usize, /// Shortest suffix n-gram matched against the context - #[arg(long, default_value_t = mlxcel::PromptLookupConfig::default().ngram_min, value_name = "N")] + #[arg( + long, + requires = "prompt_lookup", + default_value_t = mlxcel::PromptLookupConfig::default().ngram_min, + value_name = "N" + )] pub(crate) prompt_lookup_ngram_min: usize, /// Maximum tokens proposed per round (the verify block is one wider) - #[arg(long, default_value_t = mlxcel::PromptLookupConfig::default().max_draft, value_name = "N")] + #[arg( + long, + requires = "prompt_lookup", + default_value_t = mlxcel::PromptLookupConfig::default().max_draft, + value_name = "N" + )] pub(crate) prompt_lookup_max_draft: usize, /// Propose up to the maximum every round, instead of shortening or pausing /// proposals while they stop landing (for A/B measurement) - #[arg(long)] + #[arg(long, requires = "prompt_lookup")] pub(crate) prompt_lookup_no_adaptive: bool, } diff --git a/src/main_tests.rs b/src/main_tests.rs index 20a254353..95b84a959 100644 --- a/src/main_tests.rs +++ b/src/main_tests.rs @@ -17,8 +17,8 @@ use std::ffi::OsString; use std::sync::{Mutex, OnceLock}; use super::{ - Cli, Commands, FAMILY_ORDER, PipelineParallelOptions, TensorParallelOptions, - write_supported_models, write_supported_models_json, + Cli, Commands, FAMILY_ORDER, PipelineParallelOptions, PromptLookupOptions, + TensorParallelOptions, write_supported_models, write_supported_models_json, }; static CLI_ENV_LOCK: OnceLock> = OnceLock::new(); @@ -96,6 +96,95 @@ fn run_defaults_match_clap_defaults() { args.pipeline_parallel.pp_micro_batch_size, pp_default.pp_micro_batch_size ); + + let pl_default = PromptLookupOptions::default(); + assert_eq!(args.prompt_lookup.prompt_lookup, pl_default.prompt_lookup); + assert_eq!(args.prompt_lookup.config(), pl_default.config()); +} + +fn try_parse_cli(args: &'static [&'static str]) -> Result { + std::thread::Builder::new() + .stack_size(8 * 1024 * 1024) + .spawn(move || Cli::try_parse_from(args)) + .expect("spawn CLI parser thread") + .join() + .expect("CLI parser thread panicked") +} + +#[test] +fn prompt_lookup_flags_parse_into_the_generator_config() { + let cli = parse_cli(&[ + "mlxcel", + "generate", + "-m", + "models/foo", + "-p", + "hi", + "--prompt-lookup", + "--prompt-lookup-ngram-max", + "4", + "--prompt-lookup-ngram-min", + "1", + "--prompt-lookup-max-draft", + "5", + "--prompt-lookup-no-adaptive", + ]); + let Commands::Generate(args) = cli.command else { + panic!("expected generate command"); + }; + assert!(args.prompt_lookup.prompt_lookup); + assert_eq!( + args.prompt_lookup.config(), + mlxcel::PromptLookupConfig { + ngram_max: 4, + ngram_min: 1, + max_draft: 5, + adaptive: false, + } + ); +} + +/// A tuning flag without `--prompt-lookup` would otherwise be accepted and +/// silently do nothing. +#[test] +fn prompt_lookup_tuning_flags_require_prompt_lookup() { + for flag in [ + &["--prompt-lookup-ngram-max", "4"][..], + &["--prompt-lookup-ngram-min", "1"][..], + &["--prompt-lookup-max-draft", "5"][..], + &["--prompt-lookup-no-adaptive"][..], + ] { + let mut args = vec!["mlxcel", "generate", "-m", "models/foo", "-p", "hi"]; + args.extend_from_slice(flag); + let args: &'static [&'static str] = Vec::leak(args); + let err = try_parse_cli(args).expect_err("tuning flag alone must be refused"); + assert_eq!( + err.kind(), + clap::error::ErrorKind::MissingRequiredArgument, + "{flag:?}: {err}" + ); + } +} + +#[test] +fn prompt_lookup_conflicts_with_a_draft_model() { + let err = try_parse_cli(&[ + "mlxcel", + "generate", + "-m", + "models/foo", + "-p", + "hi", + "--prompt-lookup", + "--draft-model", + "models/draft", + ]) + .expect_err("--prompt-lookup with --draft-model must be refused"); + assert_eq!( + err.kind(), + clap::error::ErrorKind::ArgumentConflict, + "{err}" + ); } /// Issue #26: the rendered `mlxcel arch` output must mention every model