osobhandClaude Sonnet 5
4aaa36a57a
Whole-workspace rustfmt pass picked up while iterating on Mamba GPU backward work. Verified formatting-only via diff sampling; no logic changed. Co-Authored-By: Claude Sonnet 5 <[email protected]>
729 lines
27 KiB
Rust
729 lines
27 KiB
Rust
//! Lookahead Decoding (Fu et al. 2024, arXiv:2402.02057)
|
|
//!
|
|
//! Accelerates autoregressive LLM decoding via Jacobi iteration: instead of
|
|
//! generating one token per forward pass the decoder speculatively generates a
|
|
//! "lookahead window" of W tokens in parallel using an n-gram cache, then
|
|
//! verifies the best continuation against a scoring function that simulates a
|
|
//! single model forward pass. Accepted tokens advance the sequence without
|
|
//! additional forward passes.
|
|
//!
|
|
//! # Quick start
|
|
//!
|
|
//! ```
|
|
//! use rtx_inference::lookahead::{LookaheadDecoder, LookaheadConfig};
|
|
//!
|
|
//! let config = LookaheadConfig::default();
|
|
//! let mut decoder = LookaheadDecoder::new(config);
|
|
//!
|
|
//! // Warm the cache with some known text (token ids)
|
|
//! decoder.warm_cache(&[1, 2, 3, 4, 5, 6]);
|
|
//!
|
|
//! // score_fn: given the current token sequence, return the single most-likely
|
|
//! // next token (greedy — simulates one model forward pass).
|
|
//! let score_fn = |tokens: &[u32]| -> u32 {
|
|
//! tokens.last().copied().unwrap_or(0).wrapping_add(1)
|
|
//! };
|
|
//!
|
|
//! let generated = decoder.decode(&[1, 2, 3], 10, score_fn);
|
|
//! assert!(!generated.is_empty());
|
|
//! ```
|
|
|
|
use std::collections::{HashMap, VecDeque};
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// NGram
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// A single n-gram entry: a prefix of `n-1` tokens paired with the observed
|
|
/// continuation token.
|
|
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
|
pub struct NGram {
|
|
/// The `n-1` token prefix (context).
|
|
pub prefix: Vec<u32>,
|
|
/// The observed next token.
|
|
pub next: u32,
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// NGramCache
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Cache of n-gram statistics accumulated during or before generation.
|
|
///
|
|
/// Internally maps each observed prefix to a frequency table of continuation
|
|
/// tokens. When the total number of distinct prefixes exceeds `max_size` the
|
|
/// oldest prefix (FIFO) is evicted.
|
|
#[derive(Debug, Clone)]
|
|
pub struct NGramCache {
|
|
/// The *n* in n-gram: prefix length is `n-1`, context window is `n`.
|
|
n: usize,
|
|
/// prefix → { next_token → count }
|
|
table: HashMap<Vec<u32>, HashMap<u32, usize>>,
|
|
/// Maximum number of distinct prefixes to store before eviction.
|
|
max_size: usize,
|
|
/// FIFO insertion order for eviction.
|
|
insertion_order: VecDeque<Vec<u32>>,
|
|
/// Running total of all individual (prefix, next) observations.
|
|
total_entries: usize,
|
|
}
|
|
|
|
impl NGramCache {
|
|
/// Create a new empty cache.
|
|
///
|
|
/// # Arguments
|
|
/// * `n` — n-gram order (must be ≥ 2; `n == 1` means prefix length 0,
|
|
/// which degenerates to a unigram and is not useful here).
|
|
/// * `max_size` — maximum number of distinct prefixes before FIFO eviction.
|
|
#[must_use]
|
|
pub fn new(n: usize, max_size: usize) -> Self {
|
|
let n = n.max(2); // guard: n must be at least 2 so prefix length ≥ 1
|
|
Self {
|
|
n,
|
|
table: HashMap::new(),
|
|
max_size: max_size.max(1),
|
|
insertion_order: VecDeque::new(),
|
|
total_entries: 0,
|
|
}
|
|
}
|
|
|
|
/// Slide a window of size `n` over `tokens`, recording every (prefix, next)
|
|
/// pair. Does nothing if `tokens.len() < n`.
|
|
pub fn observe(&mut self, tokens: &[u32]) {
|
|
let prefix_len = self.n - 1;
|
|
if tokens.len() < self.n {
|
|
return;
|
|
}
|
|
for i in 0..=(tokens.len() - self.n) {
|
|
let prefix: Vec<u32> = tokens[i..i + prefix_len].to_vec();
|
|
let next = tokens[i + prefix_len];
|
|
self.insert(prefix, next);
|
|
}
|
|
}
|
|
|
|
/// Return the top-`k` candidate next tokens for `prefix`, sorted by
|
|
/// observed count descending. Returns an empty `Vec` if the prefix is
|
|
/// not in the cache or `prefix.len() != n-1`.
|
|
#[must_use]
|
|
pub fn top_candidates(&self, prefix: &[u32], k: usize) -> Vec<u32> {
|
|
if k == 0 {
|
|
return Vec::new();
|
|
}
|
|
let counts = match self.table.get(prefix) {
|
|
Some(m) => m,
|
|
None => return Vec::new(),
|
|
};
|
|
let mut pairs: Vec<(u32, usize)> = counts.iter().map(|(&t, &c)| (t, c)).collect();
|
|
// Stable sort descending by count, ties broken by token id (determinism).
|
|
pairs.sort_unstable_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
|
|
pairs.into_iter().take(k).map(|(t, _)| t).collect()
|
|
}
|
|
|
|
/// Number of distinct prefixes currently stored.
|
|
#[must_use]
|
|
pub fn num_prefixes(&self) -> usize {
|
|
self.table.len()
|
|
}
|
|
|
|
/// Total number of individual (prefix, next_token) observations recorded
|
|
/// (sum of all counts across all entries).
|
|
#[must_use]
|
|
pub fn total_observations(&self) -> usize {
|
|
self.total_entries
|
|
}
|
|
|
|
// ------------------------------------------------------------------
|
|
// Private helpers
|
|
// ------------------------------------------------------------------
|
|
|
|
fn insert(&mut self, prefix: Vec<u32>, next: u32) {
|
|
// If prefix is new and we're at capacity, evict the oldest.
|
|
if !self.table.contains_key(&prefix) {
|
|
if self.table.len() >= self.max_size {
|
|
self.evict_oldest();
|
|
}
|
|
self.table.insert(prefix.clone(), HashMap::new());
|
|
self.insertion_order.push_back(prefix.clone());
|
|
}
|
|
// Increment count.
|
|
let count = self
|
|
.table
|
|
.get_mut(&prefix)
|
|
.expect("just inserted above; prefix must exist")
|
|
.entry(next)
|
|
.or_insert(0);
|
|
*count += 1;
|
|
self.total_entries += 1;
|
|
}
|
|
|
|
fn evict_oldest(&mut self) {
|
|
if let Some(oldest) = self.insertion_order.pop_front() {
|
|
if let Some(inner) = self.table.remove(&oldest) {
|
|
// Deduct evicted counts from the running total.
|
|
let removed: usize = inner.values().sum();
|
|
self.total_entries = self.total_entries.saturating_sub(removed);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// LookaheadConfig
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Configuration for [`LookaheadDecoder`].
|
|
#[derive(Debug, Clone)]
|
|
pub struct LookaheadConfig {
|
|
/// `W`: number of speculative tokens to draft per step.
|
|
pub window_size: usize,
|
|
/// `n` in n-gram; prefix length is `n-1`.
|
|
pub ngram_order: usize,
|
|
/// Maximum number of distinct prefixes in the n-gram cache.
|
|
pub cache_size: usize,
|
|
/// Top-k n-gram candidates to consider at each draft position.
|
|
pub max_candidates_per_step: usize,
|
|
/// Token id that signals end-of-sequence.
|
|
pub eos_token_id: u32,
|
|
/// Hard cap on the total output length (initial + generated).
|
|
pub max_length: usize,
|
|
}
|
|
|
|
impl Default for LookaheadConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
window_size: 5,
|
|
ngram_order: 3,
|
|
cache_size: 10_000,
|
|
max_candidates_per_step: 3,
|
|
eos_token_id: 2,
|
|
max_length: 256,
|
|
}
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// LookaheadStats
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Cumulative statistics collected during [`LookaheadDecoder::decode`].
|
|
#[derive(Debug, Clone, Default)]
|
|
pub struct LookaheadStats {
|
|
/// Total tokens appended to the output sequence.
|
|
pub total_tokens_generated: usize,
|
|
/// Number of lookahead iterations (decode steps).
|
|
pub total_steps: usize,
|
|
/// Tokens accepted directly from n-gram cache predictions.
|
|
pub total_accepted_from_ngram: usize,
|
|
/// Tokens accepted from the model (either as correction or no-draft path).
|
|
pub total_accepted_from_model: usize,
|
|
/// Average tokens accepted per step (`total_tokens_generated / total_steps`).
|
|
pub avg_tokens_per_step: f32,
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// LookaheadDecoder
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// CPU-reference lookahead decoder implementing the Jacobi-iteration algorithm
|
|
/// from Fu et al. 2024 (arXiv:2402.02057).
|
|
///
|
|
/// This implementation runs entirely on the CPU and uses a caller-supplied
|
|
/// `score_fn` to simulate a single model forward pass (greedy: returns the
|
|
/// single most-likely next token). It is suitable for unit testing, benchmarking
|
|
/// the acceptance logic, and as a reference implementation for GPU ports.
|
|
///
|
|
/// # Algorithm overview (per step)
|
|
///
|
|
/// 1. **Lookahead phase** — draft up to `window_size` tokens from the n-gram
|
|
/// cache by repeatedly looking up the most-likely next token for the current
|
|
/// prefix.
|
|
/// 2. **Verification phase** — run `score_fn` greedily over the draft,
|
|
/// accepting tokens while they match and stopping at the first mismatch
|
|
/// (accepting the model's correction).
|
|
/// 3. **Cache update** — observe the newly appended tokens to warm the cache.
|
|
/// 4. **EOS check** — stop if the accepted tail contains the EOS token id.
|
|
pub struct LookaheadDecoder {
|
|
config: LookaheadConfig,
|
|
cache: NGramCache,
|
|
stats: LookaheadStats,
|
|
}
|
|
|
|
impl LookaheadDecoder {
|
|
/// Create a new decoder with the given configuration.
|
|
#[must_use]
|
|
pub fn new(config: LookaheadConfig) -> Self {
|
|
let cache = NGramCache::new(config.ngram_order, config.cache_size);
|
|
Self {
|
|
config,
|
|
cache,
|
|
stats: LookaheadStats::default(),
|
|
}
|
|
}
|
|
|
|
/// Create a decoder with default configuration.
|
|
#[must_use]
|
|
pub fn with_defaults() -> Self {
|
|
Self::new(LookaheadConfig::default())
|
|
}
|
|
|
|
/// Pre-load n-grams from a reference corpus to warm-start the cache.
|
|
///
|
|
/// Call this before [`decode`](Self::decode) when you have representative
|
|
/// text available (e.g., a system prompt or previous conversation turns).
|
|
pub fn warm_cache(&mut self, corpus: &[u32]) {
|
|
self.cache.observe(corpus);
|
|
}
|
|
|
|
/// Run lookahead decoding.
|
|
///
|
|
/// # Arguments
|
|
/// * `initial_tokens` — the prompt token ids (not included in the return value).
|
|
/// * `max_new_tokens` — maximum number of *new* tokens to generate.
|
|
/// * `score_fn` — closure that accepts the full current token sequence and
|
|
/// returns the single greedy-best next token. This simulates one model
|
|
/// forward pass.
|
|
///
|
|
/// # Returns
|
|
/// The newly generated tokens (does not include `initial_tokens`).
|
|
pub fn decode<F>(
|
|
&mut self,
|
|
initial_tokens: &[u32],
|
|
max_new_tokens: usize,
|
|
score_fn: F,
|
|
) -> Vec<u32>
|
|
where
|
|
F: Fn(&[u32]) -> u32,
|
|
{
|
|
let mut tokens: Vec<u32> = initial_tokens.to_vec();
|
|
let mut generated: Vec<u32> = Vec::new();
|
|
|
|
while generated.len() < max_new_tokens {
|
|
let remaining = max_new_tokens - generated.len();
|
|
|
|
// ----------------------------------------------------------
|
|
// 1. Lookahead phase: draft up to window_size tokens from cache
|
|
// ----------------------------------------------------------
|
|
let mut draft: Vec<u32> = Vec::with_capacity(self.config.window_size);
|
|
for _ in 0..self.config.window_size.min(remaining) {
|
|
// Build the prefix from the tail of (tokens + draft so far)
|
|
let combined_len = tokens.len() + draft.len();
|
|
let prefix_len = self.config.ngram_order - 1; // n-1
|
|
|
|
let prefix: Vec<u32> = if combined_len >= prefix_len {
|
|
// Take last prefix_len tokens from the combined sequence.
|
|
let start = combined_len - prefix_len;
|
|
if draft.len() >= prefix_len {
|
|
// The prefix falls entirely within draft.
|
|
draft[draft.len() - prefix_len..].to_vec()
|
|
} else {
|
|
// The prefix spans the end of tokens and the start of draft.
|
|
let from_tokens = prefix_len - draft.len();
|
|
let token_start = tokens.len().saturating_sub(from_tokens);
|
|
tokens[token_start..]
|
|
.iter()
|
|
.chain(draft.iter())
|
|
.cloned()
|
|
.collect()
|
|
}
|
|
} else {
|
|
// Not enough history yet — build what we can.
|
|
tokens.iter().chain(draft.iter()).cloned().collect()
|
|
};
|
|
let _ = combined_len; // silence unused warning
|
|
|
|
let candidates = self
|
|
.cache
|
|
.top_candidates(&prefix, self.config.max_candidates_per_step);
|
|
if candidates.is_empty() {
|
|
break; // no n-gram available — stop drafting
|
|
}
|
|
draft.push(candidates[0]);
|
|
}
|
|
|
|
// ----------------------------------------------------------
|
|
// 2. Verification phase
|
|
// ----------------------------------------------------------
|
|
let mut accepted_this_step: Vec<u32> = Vec::new();
|
|
|
|
if !draft.is_empty() {
|
|
let mut verify_ctx = tokens.clone();
|
|
for &draft_tok in &draft {
|
|
let model_tok = score_fn(&verify_ctx);
|
|
if model_tok == draft_tok {
|
|
accepted_this_step.push(draft_tok);
|
|
verify_ctx.push(draft_tok);
|
|
self.stats.total_accepted_from_ngram += 1;
|
|
} else {
|
|
// Accept model's correction and stop verifying.
|
|
accepted_this_step.push(model_tok);
|
|
self.stats.total_accepted_from_model += 1;
|
|
break;
|
|
}
|
|
// Don't exceed max_new_tokens in a single step either.
|
|
if accepted_this_step.len() >= remaining {
|
|
break;
|
|
}
|
|
}
|
|
} else {
|
|
// No draft available: just run the model once.
|
|
let model_tok = score_fn(&tokens);
|
|
accepted_this_step.push(model_tok);
|
|
self.stats.total_accepted_from_model += 1;
|
|
}
|
|
|
|
// ----------------------------------------------------------
|
|
// 3. Update n-gram cache with newly accepted tokens
|
|
// ----------------------------------------------------------
|
|
{
|
|
// Build a small window: last (n-1) tokens from the current
|
|
// sequence followed by the accepted tokens.
|
|
let prefix_len = self.config.ngram_order.saturating_sub(1);
|
|
let context_start = tokens.len().saturating_sub(prefix_len);
|
|
let new_tokens: Vec<u32> = tokens[context_start..]
|
|
.iter()
|
|
.chain(accepted_this_step.iter())
|
|
.cloned()
|
|
.collect();
|
|
self.cache.observe(&new_tokens);
|
|
}
|
|
|
|
// ----------------------------------------------------------
|
|
// 4. Extend sequence
|
|
// ----------------------------------------------------------
|
|
let to_add = accepted_this_step.len().min(remaining);
|
|
let accepted_slice = &accepted_this_step[..to_add];
|
|
|
|
tokens.extend_from_slice(accepted_slice);
|
|
generated.extend_from_slice(accepted_slice);
|
|
self.stats.total_tokens_generated += accepted_slice.len();
|
|
self.stats.total_steps += 1;
|
|
|
|
// ----------------------------------------------------------
|
|
// 5. EOS check
|
|
// ----------------------------------------------------------
|
|
if accepted_slice.last() == Some(&self.config.eos_token_id) {
|
|
break;
|
|
}
|
|
}
|
|
|
|
// Update derived statistic.
|
|
self.stats.avg_tokens_per_step =
|
|
self.stats.total_tokens_generated as f32 / self.stats.total_steps.max(1) as f32;
|
|
|
|
generated
|
|
}
|
|
|
|
/// Access accumulated decoding statistics.
|
|
#[must_use]
|
|
pub fn stats(&self) -> &LookaheadStats {
|
|
&self.stats
|
|
}
|
|
|
|
/// Reset all statistics counters to zero (cache is retained).
|
|
pub fn reset_stats(&mut self) {
|
|
self.stats = LookaheadStats::default();
|
|
}
|
|
|
|
/// Access the underlying n-gram cache.
|
|
#[must_use]
|
|
pub fn cache(&self) -> &NGramCache {
|
|
&self.cache
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Tests
|
|
// ---------------------------------------------------------------------------
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
// Helper: score_fn that always predicts last_token + 1 (greedy incrementer).
|
|
fn incrementer(tokens: &[u32]) -> u32 {
|
|
tokens.last().copied().unwrap_or(0).wrapping_add(1)
|
|
}
|
|
|
|
// ------------------------------------------------------------------
|
|
// NGramCache unit tests
|
|
// ------------------------------------------------------------------
|
|
|
|
#[test]
|
|
fn test_ngram_cache_observe_single() {
|
|
// observe [1,2,3] with n=3: prefix=[1,2], next=3
|
|
let mut cache = NGramCache::new(3, 100);
|
|
cache.observe(&[1, 2, 3]);
|
|
let candidates = cache.top_candidates(&[1, 2], 5);
|
|
assert_eq!(candidates, vec![3]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_ngram_cache_top_candidates_most_frequent_first() {
|
|
// observe token 42 three times and token 7 once for prefix [1,2]
|
|
let mut cache = NGramCache::new(3, 100);
|
|
cache.observe(&[1, 2, 42]);
|
|
cache.observe(&[1, 2, 42]);
|
|
cache.observe(&[1, 2, 42]);
|
|
cache.observe(&[1, 2, 7]);
|
|
let candidates = cache.top_candidates(&[1, 2], 2);
|
|
assert_eq!(candidates[0], 42, "42 should rank first (count=3)");
|
|
assert_eq!(candidates[1], 7, "7 should rank second (count=1)");
|
|
}
|
|
|
|
#[test]
|
|
fn test_ngram_cache_empty_prefix_returns_empty() {
|
|
let cache = NGramCache::new(3, 100);
|
|
let candidates = cache.top_candidates(&[99, 100], 5);
|
|
assert!(
|
|
candidates.is_empty(),
|
|
"unknown prefix must return empty vec"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_ngram_cache_observe_window_slides() {
|
|
// observe [1,2,3,4] with n=3: should store (1,2)→3 and (2,3)→4
|
|
let mut cache = NGramCache::new(3, 100);
|
|
cache.observe(&[1, 2, 3, 4]);
|
|
assert_eq!(cache.top_candidates(&[1, 2], 1), vec![3]);
|
|
assert_eq!(cache.top_candidates(&[2, 3], 1), vec![4]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_ngram_cache_top_k_limit() {
|
|
let mut cache = NGramCache::new(3, 100);
|
|
// Insert 5 different next tokens for the same prefix.
|
|
for next in 10u32..15 {
|
|
cache.observe(&[1, 2, next]);
|
|
}
|
|
let candidates = cache.top_candidates(&[1, 2], 3);
|
|
assert_eq!(candidates.len(), 3, "should return at most k=3 candidates");
|
|
}
|
|
|
|
#[test]
|
|
fn test_ngram_cache_total_observations() {
|
|
let mut cache = NGramCache::new(3, 100);
|
|
// [1,2,3] → 1 observation, [2,3,4] → 1 observation from one call
|
|
cache.observe(&[1, 2, 3, 4]);
|
|
assert_eq!(cache.total_observations(), 2);
|
|
// Add two more.
|
|
cache.observe(&[5, 6, 7]);
|
|
assert_eq!(cache.total_observations(), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_ngram_cache_num_prefixes() {
|
|
let mut cache = NGramCache::new(3, 100);
|
|
cache.observe(&[1, 2, 3]);
|
|
assert_eq!(cache.num_prefixes(), 1);
|
|
cache.observe(&[4, 5, 6]);
|
|
assert_eq!(cache.num_prefixes(), 2);
|
|
// Observing the same prefix again must not increase count.
|
|
cache.observe(&[1, 2, 9]);
|
|
assert_eq!(cache.num_prefixes(), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_warm_cache_populates_candidates() {
|
|
let mut decoder = LookaheadDecoder::new(LookaheadConfig::default());
|
|
decoder.warm_cache(&[10, 20, 30, 40, 50]);
|
|
// After warming, the prefix [10,20] should yield at least one candidate.
|
|
let candidates = decoder.cache().top_candidates(&[10, 20], 3);
|
|
assert!(
|
|
!candidates.is_empty(),
|
|
"warm_cache should populate the cache"
|
|
);
|
|
}
|
|
|
|
// ------------------------------------------------------------------
|
|
// Decoder integration tests
|
|
// ------------------------------------------------------------------
|
|
|
|
#[test]
|
|
fn test_decode_empty_cache_falls_back_to_model() {
|
|
let mut decoder = LookaheadDecoder::new(LookaheadConfig {
|
|
window_size: 5,
|
|
..LookaheadConfig::default()
|
|
});
|
|
// No warm_cache call — cache is empty, so every step must call score_fn.
|
|
let initial = vec![1u32, 2, 3];
|
|
let generated = decoder.decode(&initial, 5, incrementer);
|
|
assert_eq!(generated.len(), 5);
|
|
// All tokens should come from the model (none from ngram).
|
|
assert_eq!(decoder.stats().total_accepted_from_ngram, 0);
|
|
assert!(decoder.stats().total_accepted_from_model > 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_decode_returns_correct_length() {
|
|
let mut decoder = LookaheadDecoder::new(LookaheadConfig::default());
|
|
let generated = decoder.decode(&[1, 2, 3], 8, incrementer);
|
|
assert_eq!(generated.len(), 8, "should generate exactly max_new_tokens");
|
|
}
|
|
|
|
#[test]
|
|
fn test_decode_stops_on_eos() {
|
|
let eos: u32 = 99;
|
|
let config = LookaheadConfig {
|
|
eos_token_id: eos,
|
|
..LookaheadConfig::default()
|
|
};
|
|
let mut decoder = LookaheadDecoder::new(config);
|
|
// score_fn always returns the EOS token.
|
|
let generated = decoder.decode(&[1, 2, 3], 100, |_| eos);
|
|
// Should stop after the first token (which is EOS).
|
|
assert!(!generated.is_empty());
|
|
assert_eq!(*generated.last().unwrap(), eos);
|
|
assert!(generated.len() < 100, "should stop early on EOS");
|
|
}
|
|
|
|
#[test]
|
|
fn test_decode_with_matching_cache() {
|
|
// Warm the cache so it exactly predicts what the incrementer will generate.
|
|
// incrementer produces: 4,5,6,7,8,... starting from initial=[1,2,3].
|
|
let corpus: Vec<u32> = (1u32..=20).collect();
|
|
let config = LookaheadConfig {
|
|
window_size: 4,
|
|
ngram_order: 3,
|
|
eos_token_id: 255,
|
|
..LookaheadConfig::default()
|
|
};
|
|
let mut decoder = LookaheadDecoder::new(config);
|
|
decoder.warm_cache(&corpus);
|
|
|
|
let generated = decoder.decode(&[1, 2, 3], 10, incrementer);
|
|
assert_eq!(generated.len(), 10);
|
|
// Cache matches the model so we should see some ngram acceptances.
|
|
assert!(
|
|
decoder.stats().total_accepted_from_ngram > 0,
|
|
"warm matching cache should yield ngram acceptances"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_decode_with_mismatched_cache() {
|
|
// Cache predicts token 99 for every position; model (incrementer) disagrees.
|
|
let mut decoder = LookaheadDecoder::new(LookaheadConfig {
|
|
window_size: 4,
|
|
ngram_order: 2,
|
|
eos_token_id: 255,
|
|
..LookaheadConfig::default()
|
|
});
|
|
// Observe [1, 99] → cache predicts 99 after any prefix ending in 1.
|
|
decoder.warm_cache(&[1, 99, 99, 99, 99, 99]);
|
|
|
|
// incrementer starting from [1] will produce 2,3,4,... never 99.
|
|
let generated = decoder.decode(&[1], 5, incrementer);
|
|
assert_eq!(generated.len(), 5);
|
|
// Model should have overridden the cache prediction many times.
|
|
assert!(decoder.stats().total_accepted_from_model > 0);
|
|
}
|
|
|
|
// ------------------------------------------------------------------
|
|
// Statistics tests
|
|
// ------------------------------------------------------------------
|
|
|
|
#[test]
|
|
fn test_stats_total_tokens_matches_generated_length() {
|
|
let mut decoder = LookaheadDecoder::new(LookaheadConfig::default());
|
|
let generated = decoder.decode(&[1, 2, 3], 12, incrementer);
|
|
assert_eq!(
|
|
decoder.stats().total_tokens_generated,
|
|
generated.len(),
|
|
"stats counter must equal returned vec length"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_stats_steps_counted() {
|
|
let mut decoder = LookaheadDecoder::new(LookaheadConfig::default());
|
|
decoder.decode(&[1, 2, 3], 6, incrementer);
|
|
assert!(
|
|
decoder.stats().total_steps > 0,
|
|
"total_steps must be incremented each iteration"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_stats_avg_tokens_per_step_at_least_one() {
|
|
let mut decoder = LookaheadDecoder::new(LookaheadConfig::default());
|
|
decoder.decode(&[1, 2, 3], 10, incrementer);
|
|
assert!(
|
|
decoder.stats().avg_tokens_per_step >= 1.0,
|
|
"avg_tokens_per_step must be >= 1.0"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_stats_ngram_accepted_positive_when_cache_matches() {
|
|
let corpus: Vec<u32> = (1u32..=30).collect();
|
|
let config = LookaheadConfig {
|
|
window_size: 4,
|
|
ngram_order: 3,
|
|
eos_token_id: 200,
|
|
..LookaheadConfig::default()
|
|
};
|
|
let mut decoder = LookaheadDecoder::new(config);
|
|
decoder.warm_cache(&corpus);
|
|
decoder.decode(&[1, 2, 3], 15, incrementer);
|
|
assert!(
|
|
decoder.stats().total_accepted_from_ngram > 0,
|
|
"matching cache should produce ngram acceptances"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_reset_stats_clears_all_counters() {
|
|
let mut decoder = LookaheadDecoder::new(LookaheadConfig::default());
|
|
decoder.decode(&[1, 2, 3], 5, incrementer);
|
|
// Sanity: stats were actually incremented.
|
|
assert!(decoder.stats().total_steps > 0);
|
|
decoder.reset_stats();
|
|
let s = decoder.stats();
|
|
assert_eq!(s.total_tokens_generated, 0);
|
|
assert_eq!(s.total_steps, 0);
|
|
assert_eq!(s.total_accepted_from_ngram, 0);
|
|
assert_eq!(s.total_accepted_from_model, 0);
|
|
assert_eq!(s.avg_tokens_per_step, 0.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_decode_longer_acceptance_when_cache_warm_and_model_agrees() {
|
|
// Build a corpus that perfectly matches what incrementer will generate.
|
|
let corpus: Vec<u32> = (1u32..=50).collect();
|
|
let config = LookaheadConfig {
|
|
window_size: 5,
|
|
ngram_order: 3,
|
|
cache_size: 10_000,
|
|
max_candidates_per_step: 3,
|
|
eos_token_id: 200,
|
|
max_length: 256,
|
|
};
|
|
let mut decoder = LookaheadDecoder::new(config);
|
|
decoder.warm_cache(&corpus);
|
|
|
|
decoder.decode(&[1, 2, 3], 20, incrementer);
|
|
assert!(
|
|
decoder.stats().avg_tokens_per_step > 1.0,
|
|
"avg_tokens_per_step should exceed 1.0 when cache is warm and matches the model"
|
|
);
|
|
}
|
|
|
|
// Extra coverage for edge conditions and k=0 guard.
|
|
#[test]
|
|
fn test_ngram_cache_top_k_zero_returns_empty() {
|
|
let mut cache = NGramCache::new(3, 100);
|
|
cache.observe(&[1, 2, 3]);
|
|
assert!(cache.top_candidates(&[1, 2], 0).is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_decode_zero_new_tokens_returns_empty() {
|
|
let mut decoder = LookaheadDecoder::new(LookaheadConfig::default());
|
|
let generated = decoder.decode(&[1, 2, 3], 0, incrementer);
|
|
assert!(
|
|
generated.is_empty(),
|
|
"max_new_tokens=0 should return empty vec"
|
|
);
|
|
}
|
|
}
|