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 582b99765..eadf2c4a0 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,90 @@ 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, 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. +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); + // 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, + 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 +2652,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); } @@ -2590,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 4bcb66065..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}; @@ -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. @@ -964,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. @@ -1417,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/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 f395a6a27..bc0707709 100644 --- a/src/lib/mlxcel-core/src/generate.rs +++ b/src/lib/mlxcel-core/src/generate.rs @@ -343,6 +343,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 @@ -1662,9 +1725,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 { @@ -1756,42 +1817,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..56d90e81f --- /dev/null +++ b/src/lib/mlxcel-core/src/speculative/prompt_lookup.rs @@ -0,0 +1,921 @@ +// 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, SequenceStateBackend, 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::{ + 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; +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; +/// 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)] +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.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(()) + } +} + +/// 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 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( + "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.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 \ + 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 + } + + /// 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) { + 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, 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, + prompt_tokens: &[i32], + max_tokens: usize, + sampling: &SamplingConfig, + ) -> (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!( + !self.caches.is_empty() && can_trim_prompt_cache(&self.caches), + "prompt-lookup decoding requires trimmable external KV caches" + ); + 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(); + 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 + // 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, _) = 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); + 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, _) = 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) + } + }; + + 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..57c9aea1b --- /dev/null +++ b/src/lib/mlxcel-core/src/speculative/prompt_lookup_tests.rs @@ -0,0 +1,638 @@ +// 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]); +} + +#[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 64f466244..728591cb3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -370,6 +370,76 @@ 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, + 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, + 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, + 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, requires = "prompt_lookup")] + 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 +468,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 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