feat(batch16): SOAP optimizer, lookahead decoding, SWA+SWAG

- SoapOptimizer: Adam in Shampoo eigenbasis (arXiv:2409.11321); Jacobi
  eigendecomposition for L/R Kronecker factors; projection G_hat=Q_L^T@G@Q_R,
  bias-corrected Adam, unproject U=Q_L@U_hat@Q_R^T; 1D plain Adam fallback; 19 tests
- LookaheadDecoder: NGramCache (FIFO eviction, count-sorted candidates);
  draft-then-verify loop; auto-cache update on accepted tokens; LookaheadStats
  with avg_tokens_per_step; 22 tests
- SwaTrainer+SwagBuffer: cyclic cosine LR schedule; online incremental mean
  (SwaBuffer); E[θ²]-E[θ]² diagonal variance + low-rank deviation columns;
  Box-Muller SWAG sample; 29 tests

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>

crates/production/rtx-inference/src/lookahead.rs

@@ -0,0 +1,712 @@
+//! 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");
+    }
+}