feat(batch15): Shampoo optimizer, beam search decoder, sliding window attention

- ShampooOptimizer: Kronecker-factored L/R preconditioners; Schulz iteration
  for A^{-1/4} (two-pass: inv_sqrt then inv_sqrt of sqrt); spectral-norm
  normalization; large-dim SGD fallback; 16 tests
- BeamSearchDecoder: length normalization (Wu et al. α); n-gram blocking;
  EOS suppression before min_length; DiverseBeamSearchDecoder with per-group
  diversity penalty; 20 tests
- SlidingWindowAttention: causal/bidir window; global tokens attend to all;
  O(n·W) forward_single_head + multi-head forward; AttentionStats sparsity;
  WindowMask; 20 tests

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

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

@@ -0,0 +1,964 @@
+//! Beam search decoding for autoregressive language models.
+//!
+//! Implements standard beam search (Sutskever et al. 2014) and diverse beam search
+//! (Vijayakumar et al. 2018) with Wu et al. 2016 length normalization.
+//!
+//! # Example
+//!
+//! ```rust
+//! use rtx_inference::beam_search::{BeamSearchConfig, BeamSearchDecoder};
+//!
+//! let config = BeamSearchConfig {
+//!     beam_width: 4,
+//!     max_length: 32,
+//!     eos_token_id: 2,
+//!     length_penalty_alpha: 0.6,
+//!     min_length: 1,
+//!     no_repeat_ngram_size: 0,
+//! };
+//! let decoder = BeamSearchDecoder::new(config);
+//!
+//! let vocab_size = 32_000;
+//! let results = decoder.decode(&[1], vocab_size, 2, |tokens| {
+//!     let mut lp = vec![-10.0_f32; vocab_size];
+//!     lp[5] = -0.1;
+//!     lp[2] = -0.5; // EOS
+//!     lp
+//! });
+//! assert!(!results.is_empty());
+//! ```
+
+use std::collections::HashSet;
+
+// ---------------------------------------------------------------------------
+// Beam
+// ---------------------------------------------------------------------------
+
+/// A single beam hypothesis tracking generated tokens and cumulative log-probability.
+#[derive(Debug, Clone)]
+pub struct Beam {
+    /// Token IDs generated so far (including any initial prompt tokens).
+    pub tokens: Vec<u32>,
+    /// Cumulative sum of log-probabilities for all generated tokens.
+    pub log_prob: f32,
+    /// Whether this beam has emitted the EOS token and is considered finished.
+    pub finished: bool,
+}
+
+impl Beam {
+    /// Create a new, empty beam with the given initial log-probability.
+    ///
+    /// The token list starts empty; use `tokens` field to pre-populate prompt tokens.
+    #[must_use]
+    pub fn new(log_prob: f32) -> Self {
+        Self {
+            tokens: Vec::new(),
+            log_prob,
+            finished: false,
+        }
+    }
+
+    /// Length-normalized score using the Wu et al. 2016 formula.
+    ///
+    /// `score = log_prob / length_penalty(len)`
+    /// `length_penalty(l) = ((5 + l) / 6) ^ alpha`
+    ///
+    /// With `alpha = 0` the raw log-probability is returned unchanged.
+    /// With `alpha = 1` full length normalization is applied.
+    ///
+    /// An empty token sequence returns `log_prob` directly to avoid division issues.
+    #[must_use]
+    pub fn normalized_score(&self, alpha: f32) -> f32 {
+        let len = self.tokens.len();
+        if len == 0 || alpha == 0.0 {
+            return self.log_prob;
+        }
+        let penalty = ((5.0 + len as f32) / 6.0_f32).powf(alpha);
+        self.log_prob / penalty
+    }
+}
+
+// ---------------------------------------------------------------------------
+// BeamSearchConfig
+// ---------------------------------------------------------------------------
+
+/// Configuration parameters for the beam search decoder.
+#[derive(Debug, Clone)]
+pub struct BeamSearchConfig {
+    /// Number of beams to maintain at each step.
+    pub beam_width: usize,
+    /// Maximum number of tokens to generate per beam.
+    pub max_length: usize,
+    /// Token ID that marks end-of-sequence.
+    pub eos_token_id: u32,
+    /// Length penalty exponent (Wu et al. 2016). 0 = no penalty, 1 = full normalization.
+    pub length_penalty_alpha: f32,
+    /// Suppress EOS emission before this many generated tokens.
+    pub min_length: usize,
+    /// If > 0, block any token that would complete an already-seen n-gram of this size.
+    pub no_repeat_ngram_size: usize,
+}
+
+impl Default for BeamSearchConfig {
+    fn default() -> Self {
+        Self {
+            beam_width: 4,
+            max_length: 128,
+            eos_token_id: 2,
+            length_penalty_alpha: 0.6,
+            min_length: 1,
+            no_repeat_ngram_size: 0,
+        }
+    }
+}
+
+// ---------------------------------------------------------------------------
+// BeamSearchDecoder
+// ---------------------------------------------------------------------------
+
+/// Standard beam search decoder.
+pub struct BeamSearchDecoder {
+    config: BeamSearchConfig,
+}
+
+impl BeamSearchDecoder {
+    /// Create a decoder with a custom configuration.
+    #[must_use]
+    pub fn new(config: BeamSearchConfig) -> Self {
+        Self { config }
+    }
+
+    /// Create a decoder with default configuration.
+    #[must_use]
+    pub fn default_decoder() -> Self {
+        Self::new(BeamSearchConfig::default())
+    }
+
+    /// Run beam search given a scoring function.
+    ///
+    /// # Arguments
+    ///
+    /// * `initial_tokens` — Prompt tokens prepended to every beam (not generated; used as
+    ///   context only). Their length offsets `min_length` enforcement.
+    /// * `vocab_size` — Total vocabulary size; used as the upper bound for candidate tokens.
+    /// * `num_return` — Number of best beams to return.
+    /// * `score_fn` — Callable that receives the full current token sequence (prompt +
+    ///   generated) and returns a `Vec<f32>` of length `vocab_size` containing log-probs.
+    ///
+    /// # Returns
+    ///
+    /// Up to `num_return` beams sorted by descending normalized score. If no beam has
+    /// emitted EOS the best still-active beams are returned instead.
+    pub fn decode<F>(
+        &self,
+        initial_tokens: &[u32],
+        vocab_size: usize,
+        num_return: usize,
+        score_fn: F,
+    ) -> Vec<Beam>
+    where
+        F: Fn(&[u32]) -> Vec<f32>,
+    {
+        let bw = self.config.beam_width;
+        let alpha = self.config.length_penalty_alpha;
+
+        // Initialise with a single beam carrying all prompt tokens.
+        let mut active: Vec<Beam> = vec![Beam {
+            tokens: initial_tokens.to_vec(),
+            log_prob: 0.0,
+            finished: false,
+        }];
+        let mut finished: Vec<Beam> = Vec::new();
+
+        for _step in 0..self.config.max_length {
+            if active.is_empty() {
+                break;
+            }
+
+            let mut candidates: Vec<Beam> = Vec::new();
+
+            for beam in &active {
+                let log_probs = score_fn(&beam.tokens);
+
+                // Determine which token ids are blocked by the n-gram filter.
+                let blocked = if self.config.no_repeat_ngram_size > 0 {
+                    Self::blocked_tokens_for_ngram(&beam.tokens, self.config.no_repeat_ngram_size)
+                } else {
+                    HashSet::new()
+                };
+
+                // EOS is suppressed until the beam has generated at least `min_length` tokens
+                // beyond the prompt.
+                let generated_so_far = beam.tokens.len().saturating_sub(initial_tokens.len());
+                let suppress_eos = generated_so_far < self.config.min_length;
+
+                // Request extra tokens so that after filtering (EOS suppression, n-gram
+                // blocking) we still have at least `bw` live candidates.  The +2 absorbs
+                // one EOS suppression slot and one n-gram block slot simultaneously.
+                let fetch_k = bw * 2 + 2;
+                let top = Self::top_k_indices(&log_probs, fetch_k);
+                let mut added = 0usize;
+                for (token_id, lp) in top {
+                    if added >= bw {
+                        break;
+                    }
+                    let token_id = token_id as u32;
+                    if blocked.contains(&token_id) {
+                        continue;
+                    }
+                    if suppress_eos && token_id == self.config.eos_token_id {
+                        continue;
+                    }
+
+                    let mut new_tokens = beam.tokens.clone();
+                    new_tokens.push(token_id);
+                    let new_lp = beam.log_prob + lp;
+                    let is_finished = token_id == self.config.eos_token_id;
+                    candidates.push(Beam {
+                        tokens: new_tokens,
+                        log_prob: new_lp,
+                        finished: is_finished,
+                    });
+                    added += 1;
+                }
+            }
+
+            // Sort all candidates by normalized score descending.
+            candidates.sort_by(|a, b| {
+                b.normalized_score(alpha)
+                    .partial_cmp(&a.normalized_score(alpha))
+                    .unwrap_or(std::cmp::Ordering::Equal)
+            });
+
+            // Partition into finished and active, keeping at most `bw` active beams.
+            active.clear();
+            for beam in candidates.into_iter().take(bw * 2) {
+                if beam.finished {
+                    finished.push(beam);
+                } else if active.len() < bw {
+                    active.push(beam);
+                }
+            }
+
+            // Early stopping: if we already have enough finished beams and no active beam
+            // can possibly outscore the worst finished beam, we are done.
+            if finished.len() >= num_return && !active.is_empty() {
+                let worst_finished = finished
+                    .iter()
+                    .map(|b| b.normalized_score(alpha))
+                    .fold(f32::INFINITY, f32::min);
+                let best_active = active
+                    .iter()
+                    .map(|b| b.normalized_score(alpha))
+                    .fold(f32::NEG_INFINITY, f32::max);
+                if best_active < worst_finished {
+                    break;
+                }
+            }
+        }
+
+        // Any remaining active beams are folded into finished so the caller always gets
+        // results even when no EOS was produced within `max_length`.
+        finished.extend(active);
+        finished.sort_by(|a, b| {
+            b.normalized_score(alpha)
+                .partial_cmp(&a.normalized_score(alpha))
+                .unwrap_or(std::cmp::Ordering::Equal)
+        });
+
+        finished.into_iter().take(num_return).collect()
+    }
+
+    /// Return the top-`k` `(index, log_prob)` pairs from `log_probs`, sorted descending.
+    ///
+    /// If `k` is greater than or equal to the length of `log_probs`, all entries are
+    /// returned (still sorted descending).
+    #[must_use]
+    pub fn top_k_indices(log_probs: &[f32], k: usize) -> Vec<(usize, f32)> {
+        let k = k.min(log_probs.len());
+        // Partial sort: O(n log k) via a min-heap maintained at size k.
+        // For typical vocab sizes (32k–256k) and small k (beam width ≤ 16) this is
+        // substantially faster than a full sort.
+        use std::collections::BinaryHeap;
+        use std::cmp::Reverse;
+
+        // BinaryHeap is a max-heap; we wrap in Reverse to get a min-heap so we can
+        // efficiently evict the smallest element as we scan.
+        // The heap stores (ordered_float_bits, index) where bits are from f32::to_bits
+        // reinterpreted for ordering. We use a small wrapper instead.
+        #[derive(PartialEq)]
+        struct OrdF32(f32);
+
+        impl Eq for OrdF32 {}
+
+        impl PartialOrd for OrdF32 {
+            fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+                Some(self.cmp(other))
+            }
+        }
+
+        impl Ord for OrdF32 {
+            fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+                self.0
+                    .partial_cmp(&other.0)
+                    .unwrap_or(std::cmp::Ordering::Equal)
+            }
+        }
+
+        // min-heap of size k: smallest value is at the top and will be evicted first.
+        let mut heap: BinaryHeap<Reverse<(OrdF32, usize)>> = BinaryHeap::with_capacity(k + 1);
+
+        for (idx, &lp) in log_probs.iter().enumerate() {
+            heap.push(Reverse((OrdF32(lp), idx)));
+            if heap.len() > k {
+                heap.pop();
+            }
+        }
+
+        // Drain heap into a Vec and sort descending.
+        let mut result: Vec<(usize, f32)> = heap
+            .into_iter()
+            .map(|Reverse((OrdF32(lp), idx))| (idx, lp))
+            .collect();
+        result.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
+        result
+    }
+
+    /// Compute the set of token IDs that would complete an already-seen n-gram.
+    ///
+    /// For a token sequence `tokens` and n-gram size `n`, we look at the last `n-1`
+    /// tokens as a suffix key and find all tokens `t` such that `suffix ++ [t]` appears
+    /// earlier in `tokens`. Those tokens are "blocked" because appending them would
+    /// produce a repeated n-gram.
+    ///
+    /// Returns an empty set when there are fewer than `n-1` tokens (no suffix to match).
+    #[must_use]
+    pub fn blocked_tokens_for_ngram(tokens: &[u32], n: usize) -> HashSet<u32> {
+        if n == 0 || tokens.len() < n {
+            return HashSet::new();
+        }
+
+        let suffix_len = n - 1;
+        let suffix = &tokens[tokens.len() - suffix_len..];
+
+        let mut blocked = HashSet::new();
+
+        // Slide a window of length `suffix_len` over all earlier positions.
+        // For each position i where tokens[i..i+suffix_len] == suffix, the token at
+        // position i+suffix_len would complete a repeated n-gram.
+        let end = tokens.len() - suffix_len; // exclusive upper bound for window starts
+        for i in 0..end {
+            if &tokens[i..i + suffix_len] == suffix {
+                blocked.insert(tokens[i + suffix_len]);
+            }
+        }
+
+        blocked
+    }
+}
+
+impl Default for BeamSearchDecoder {
+    fn default() -> Self {
+        Self::new(BeamSearchConfig::default())
+    }
+}
+
+// ---------------------------------------------------------------------------
+// DiverseBeamSearchDecoder
+// ---------------------------------------------------------------------------
+
+/// Diverse beam search decoder (Vijayakumar et al. 2018).
+///
+/// Divides the total beam budget into `num_groups` groups of size
+/// `config.beam_width / num_groups`. Each group runs standard beam search at each
+/// step, but tokens already chosen by earlier groups in the same step are penalised
+/// by `diversity_penalty` subtracted from their log-probability.
+///
+/// This encourages the returned hypotheses to differ from one another while still
+/// being high-quality.
+pub struct DiverseBeamSearchDecoder {
+    config: BeamSearchConfig,
+    /// Number of beam groups. Must divide `config.beam_width` evenly.
+    pub num_groups: usize,
+    /// Score penalty applied per token already chosen by a preceding group at the same step.
+    pub diversity_penalty: f32,
+}
+
+impl DiverseBeamSearchDecoder {
+    /// Create a diverse beam search decoder.
+    ///
+    /// # Panics
+    ///
+    /// Panics in debug mode if `num_groups` is 0 or exceeds `config.beam_width`.
+    #[must_use]
+    pub fn new(config: BeamSearchConfig, num_groups: usize, diversity_penalty: f32) -> Self {
+        debug_assert!(num_groups > 0, "num_groups must be at least 1");
+        debug_assert!(
+            num_groups <= config.beam_width,
+            "num_groups must not exceed beam_width"
+        );
+        Self {
+            config,
+            num_groups,
+            diversity_penalty,
+        }
+    }
+
+    /// Run diverse beam search.
+    ///
+    /// Groups run sequentially at each step. Group `g` penalises every token that was
+    /// already selected as the first new token by groups `0..g` at the current step.
+    ///
+    /// Returns up to `num_return` beams sorted by descending normalized score.
+    pub fn decode<F>(
+        &self,
+        initial_tokens: &[u32],
+        vocab_size: usize,
+        num_return: usize,
+        score_fn: F,
+    ) -> Vec<Beam>
+    where
+        F: Fn(&[u32]) -> Vec<f32>,
+    {
+        let alpha = self.config.length_penalty_alpha;
+        let num_groups = self.num_groups.max(1);
+        // Each group maintains its own set of active beams.
+        let beams_per_group = (self.config.beam_width / num_groups).max(1);
+
+        let initial_beam = Beam {
+            tokens: initial_tokens.to_vec(),
+            log_prob: 0.0,
+            finished: false,
+        };
+
+        // One Vec<Beam> per group.
+        let mut group_active: Vec<Vec<Beam>> = (0..num_groups)
+            .map(|_| vec![initial_beam.clone()])
+            .collect();
+        let mut finished: Vec<Beam> = Vec::new();
+
+        for _step in 0..self.config.max_length {
+            // Track which tokens have been committed by each group this step so that
+            // subsequent groups can be penalised.
+            let mut chosen_tokens: Vec<HashSet<u32>> = vec![HashSet::new(); num_groups];
+
+            for g in 0..num_groups {
+                let active = &group_active[g];
+                if active.is_empty() {
+                    continue;
+                }
+
+                // Tokens penalised are those committed by groups 0..g.
+                let penalised: HashSet<u32> = chosen_tokens[..g]
+                    .iter()
+                    .flat_map(|s| s.iter().copied())
+                    .collect();
+
+                let mut candidates: Vec<Beam> = Vec::new();
+
+                for beam in active {
+                    let log_probs = score_fn(&beam.tokens);
+
+                    let blocked = if self.config.no_repeat_ngram_size > 0 {
+                        BeamSearchDecoder::blocked_tokens_for_ngram(
+                            &beam.tokens,
+                            self.config.no_repeat_ngram_size,
+                        )
+                    } else {
+                        HashSet::new()
+                    };
+
+                    let generated_so_far =
+                        beam.tokens.len().saturating_sub(initial_tokens.len());
+                    let suppress_eos = generated_so_far < self.config.min_length;
+
+                    // Apply diversity penalty before selecting top-k.
+                    let adjusted: Vec<f32> = log_probs
+                        .iter()
+                        .enumerate()
+                        .map(|(i, &lp)| {
+                            if penalised.contains(&(i as u32)) {
+                                lp - self.diversity_penalty
+                            } else {
+                                lp
+                            }
+                        })
+                        .collect();
+
+                    let fetch_k = beams_per_group * 2 + 2;
+                    let top = BeamSearchDecoder::top_k_indices(&adjusted, fetch_k);
+                    let mut added = 0usize;
+                    for (token_id, lp) in top {
+                        if added >= beams_per_group {
+                            break;
+                        }
+                        let token_id = token_id as u32;
+                        if blocked.contains(&token_id) {
+                            continue;
+                        }
+                        if suppress_eos && token_id == self.config.eos_token_id {
+                            continue;
+                        }
+
+                        // Record what each group committed.
+                        chosen_tokens[g].insert(token_id);
+                        added += 1;
+
+                        // Use the un-penalised log-prob to keep beam scores comparable.
+                        let raw_lp = if token_id < vocab_size as u32 {
+                            log_probs[token_id as usize]
+                        } else {
+                            lp
+                        };
+
+                        let mut new_tokens = beam.tokens.clone();
+                        new_tokens.push(token_id);
+                        let new_lp = beam.log_prob + raw_lp;
+                        let is_finished = token_id == self.config.eos_token_id;
+                        candidates.push(Beam {
+                            tokens: new_tokens,
+                            log_prob: new_lp,
+                            finished: is_finished,
+                        });
+                    }
+                }
+
+                // Sort and prune.
+                candidates.sort_by(|a, b| {
+                    b.normalized_score(alpha)
+                        .partial_cmp(&a.normalized_score(alpha))
+                        .unwrap_or(std::cmp::Ordering::Equal)
+                });
+
+                let mut new_active: Vec<Beam> = Vec::new();
+                for beam in candidates.into_iter().take(beams_per_group * 2) {
+                    if beam.finished {
+                        finished.push(beam);
+                    } else if new_active.len() < beams_per_group {
+                        new_active.push(beam);
+                    }
+                }
+                group_active[g] = new_active;
+            }
+
+            // Check whether all groups are exhausted.
+            if group_active.iter().all(|g| g.is_empty()) {
+                break;
+            }
+        }
+
+        // Collect remaining active beams as fallback results.
+        for group in group_active {
+            finished.extend(group);
+        }
+
+        finished.sort_by(|a, b| {
+            b.normalized_score(alpha)
+                .partial_cmp(&a.normalized_score(alpha))
+                .unwrap_or(std::cmp::Ordering::Equal)
+        });
+
+        finished.into_iter().take(num_return).collect()
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // ---- Helper ---------------------------------------------------------------
+
+    fn make_score_fn(
+        vocab_size: usize,
+        eos_id: u32,
+    ) -> impl Fn(&[u32]) -> Vec<f32> {
+        move |_tokens: &[u32]| {
+            let mut lp = vec![-10.0_f32; vocab_size];
+            lp[5] = -0.1;
+            lp[eos_id as usize] = -0.5;
+            lp
+        }
+    }
+
+    // ---- Beam::normalized_score -----------------------------------------------
+
+    #[test]
+    fn test_beam_normalized_score_alpha0() {
+        // alpha=0 => score == log_prob regardless of length
+        let mut beam = Beam::new(-3.0);
+        beam.tokens = vec![1, 2, 3];
+        let score = beam.normalized_score(0.0);
+        assert!(
+            (score - (-3.0)).abs() < 1e-6,
+            "alpha=0 should return raw log_prob, got {score}"
+        );
+    }
+
+    #[test]
+    fn test_beam_normalized_score_alpha1() {
+        // alpha=1, len=1 => penalty = ((5+1)/6)^1 = 1.0 => score == log_prob
+        let mut beam = Beam::new(-6.0);
+        beam.tokens = vec![42];
+        let score = beam.normalized_score(1.0);
+        let expected = -6.0_f32 / ((5.0 + 1.0) / 6.0_f32).powf(1.0);
+        assert!(
+            (score - expected).abs() < 1e-5,
+            "Expected {expected}, got {score}"
+        );
+    }
+
+    #[test]
+    fn test_beam_normalized_score_longer_beam_wins() {
+        // Two beams with the same total log_prob; the longer one gets a better
+        // (less negative) normalised score when alpha > 0.
+        let mut short_beam = Beam::new(-6.0);
+        short_beam.tokens = vec![1];
+
+        let mut long_beam = Beam::new(-6.0);
+        long_beam.tokens = vec![1, 2, 3, 4, 5];
+
+        let alpha = 0.6;
+        let short_score = short_beam.normalized_score(alpha);
+        let long_score = long_beam.normalized_score(alpha);
+        assert!(
+            long_score > short_score,
+            "Longer beam should have higher normalised score: long={long_score}, short={short_score}"
+        );
+    }
+
+    // ---- top_k_indices --------------------------------------------------------
+
+    #[test]
+    fn test_top_k_indices_returns_k() {
+        let lp: Vec<f32> = (0..20).map(|i| -(i as f32)).collect();
+        let result = BeamSearchDecoder::top_k_indices(&lp, 5);
+        assert_eq!(result.len(), 5, "Should return exactly 5 elements");
+    }
+
+    #[test]
+    fn test_top_k_indices_descending() {
+        let lp = vec![-3.0, -1.0, -4.0, -1.5, -2.0];
+        let result = BeamSearchDecoder::top_k_indices(&lp, 3);
+        assert_eq!(result.len(), 3);
+        // Verify descending order.
+        for window in result.windows(2) {
+            assert!(
+                window[0].1 >= window[1].1,
+                "Not sorted descending: {:?}",
+                result
+            );
+        }
+        // Top element should be index 1 (log_prob -1.0).
+        assert_eq!(result[0].0, 1);
+    }
+
+    #[test]
+    fn test_top_k_larger_than_vocab() {
+        let lp = vec![-1.0, -2.0, -3.0];
+        // k=10 > vocab_size=3 => return all 3
+        let result = BeamSearchDecoder::top_k_indices(&lp, 10);
+        assert_eq!(result.len(), 3);
+    }
+
+    // ---- blocked_tokens_for_ngram ---------------------------------------------
+
+    #[test]
+    fn test_blocked_ngrams_empty() {
+        let blocked = BeamSearchDecoder::blocked_tokens_for_ngram(&[], 2);
+        assert!(blocked.is_empty());
+    }
+
+    #[test]
+    fn test_blocked_ngrams_bigram() {
+        // Sequence: [A, B, C, A] — suffix is [A], look for earlier [A] => B is blocked.
+        let tokens: Vec<u32> = vec![10, 20, 30, 10];
+        let blocked = BeamSearchDecoder::blocked_tokens_for_ngram(&tokens, 2);
+        // Earlier [10] at position 0 is followed by 20 => 20 is blocked.
+        assert!(
+            blocked.contains(&20),
+            "Token 20 should be blocked; got {blocked:?}"
+        );
+    }
+
+    #[test]
+    fn test_blocked_ngrams_trigram() {
+        // [A, B, C, A, B] — suffix is [A, B], look for earlier [A, B] at pos 0 => C(=30) blocked.
+        let tokens: Vec<u32> = vec![10, 20, 30, 10, 20];
+        let blocked = BeamSearchDecoder::blocked_tokens_for_ngram(&tokens, 3);
+        assert!(
+            blocked.contains(&30),
+            "Token 30 should be blocked; got {blocked:?}"
+        );
+    }
+
+    #[test]
+    fn test_blocked_ngrams_no_repeat() {
+        // No repeated n-gram => empty blocked set.
+        let tokens: Vec<u32> = vec![1, 2, 3, 4, 5];
+        let blocked = BeamSearchDecoder::blocked_tokens_for_ngram(&tokens, 2);
+        assert!(
+            blocked.is_empty(),
+            "No repeated bigrams, should be empty; got {blocked:?}"
+        );
+    }
+
+    // ---- decode ---------------------------------------------------------------
+
+    #[test]
+    fn test_decode_single_step() {
+        // With beam_width=1 and a deterministic scorer, we should get exactly one beam
+        // after one step (prompt=[1], max_length=1).
+        let vocab_size = 10;
+        let eos_id = 2_u32;
+        let config = BeamSearchConfig {
+            beam_width: 1,
+            max_length: 1,
+            eos_token_id: eos_id,
+            length_penalty_alpha: 0.0,
+            min_length: 0,
+            no_repeat_ngram_size: 0,
+        };
+        let decoder = BeamSearchDecoder::new(config);
+        let score_fn = |_: &[u32]| {
+            let mut lp = vec![-10.0_f32; vocab_size];
+            lp[5] = -0.1;
+            lp[eos_id as usize] = -0.5;
+            lp
+        };
+        let results = decoder.decode(&[1], vocab_size, 1, score_fn);
+        assert_eq!(results.len(), 1);
+        // With min_length=0, EOS is allowed. But token 5 has higher lp, so beam should end with 5.
+        let last = *results[0].tokens.last().unwrap();
+        assert_eq!(last, 5, "Greedy first token should be 5");
+    }
+
+    #[test]
+    fn test_decode_returns_num_return() {
+        let vocab_size = 32;
+        let config = BeamSearchConfig {
+            beam_width: 4,
+            max_length: 10,
+            eos_token_id: 2,
+            length_penalty_alpha: 0.6,
+            min_length: 1,
+            no_repeat_ngram_size: 0,
+        };
+        let decoder = BeamSearchDecoder::new(config);
+        let results = decoder.decode(&[1], vocab_size, 2, make_score_fn(vocab_size, 2));
+        assert_eq!(results.len(), 2, "Should return exactly 2 beams");
+    }
+
+    #[test]
+    fn test_decode_eos_stops_beam() {
+        // Force EOS to be the highest-scoring token so the beam finishes immediately.
+        let vocab_size = 10;
+        let eos_id = 2_u32;
+        let config = BeamSearchConfig {
+            beam_width: 1,
+            max_length: 50,
+            eos_token_id: eos_id,
+            length_penalty_alpha: 0.0,
+            min_length: 0,
+            no_repeat_ngram_size: 0,
+        };
+        let decoder = BeamSearchDecoder::new(config);
+        let score_fn = |_: &[u32]| {
+            let mut lp = vec![-10.0_f32; vocab_size];
+            lp[eos_id as usize] = -0.1; // EOS is best
+            lp
+        };
+        let results = decoder.decode(&[1], vocab_size, 1, score_fn);
+        assert_eq!(results.len(), 1);
+        assert!(results[0].finished, "Beam should be marked finished after EOS");
+        assert_eq!(
+            *results[0].tokens.last().unwrap(),
+            eos_id,
+            "Last token should be EOS"
+        );
+    }
+
+    #[test]
+    fn test_decode_min_length_suppresses_eos() {
+        // With min_length=5 the beam must generate at least 5 tokens before EOS is allowed.
+        let vocab_size = 10;
+        let eos_id = 2_u32;
+        let config = BeamSearchConfig {
+            beam_width: 1,
+            max_length: 20,
+            eos_token_id: eos_id,
+            length_penalty_alpha: 0.0,
+            min_length: 5,
+            no_repeat_ngram_size: 0,
+        };
+        let decoder = BeamSearchDecoder::new(config);
+        // EOS is always the best token, but is suppressed until min_length is met.
+        let score_fn = |_: &[u32]| {
+            let mut lp = vec![-10.0_f32; vocab_size];
+            lp[5] = -0.2;
+            lp[eos_id as usize] = -0.1; // EOS beats others, but suppressed until len>=5
+            lp
+        };
+        let results = decoder.decode(&[1], vocab_size, 1, score_fn);
+        assert!(!results.is_empty());
+        // The beam must have at least 1 (prompt) + 5 (generated) = 6 tokens before finishing.
+        let generated = results[0].tokens.len() - 1; // subtract prompt token
+        assert!(
+            generated >= 5,
+            "Should have generated at least 5 tokens before EOS, got {generated}"
+        );
+    }
+
+    #[test]
+    fn test_decode_always_returns_result() {
+        // Even if no EOS is emitted within max_length, active beams are returned.
+        let vocab_size = 10;
+        let config = BeamSearchConfig {
+            beam_width: 2,
+            max_length: 5,
+            eos_token_id: 99, // EOS id that will never be produced by the scorer
+            length_penalty_alpha: 0.0,
+            min_length: 0,
+            no_repeat_ngram_size: 0,
+        };
+        let decoder = BeamSearchDecoder::new(config);
+        let score_fn = |_: &[u32]| {
+            let mut lp = vec![-10.0_f32; vocab_size];
+            lp[5] = -0.1;
+            lp[3] = -0.2;
+            lp
+        };
+        let results = decoder.decode(&[1], vocab_size, 1, score_fn);
+        assert!(!results.is_empty(), "Should always return at least one beam");
+    }
+
+    #[test]
+    fn test_decode_beam_width_1_greedy() {
+        // beam_width=1 should behave identically to greedy decoding.
+        let vocab_size = 10;
+        let eos_id = 2_u32;
+        let config = BeamSearchConfig {
+            beam_width: 1,
+            max_length: 8,
+            eos_token_id: eos_id,
+            length_penalty_alpha: 0.0,
+            min_length: 3,
+            no_repeat_ngram_size: 0,
+        };
+        let decoder = BeamSearchDecoder::new(config);
+        // Scorer alternates between two tokens to create a deterministic sequence.
+        let score_fn = |tokens: &[u32]| {
+            let mut lp = vec![-10.0_f32; vocab_size];
+            if tokens.len() % 2 == 0 {
+                lp[7] = -0.1;
+            } else {
+                lp[3] = -0.1;
+            }
+            lp[eos_id as usize] = -0.5;
+            lp
+        };
+        let results = decoder.decode(&[1], vocab_size, 1, score_fn);
+        assert_eq!(results.len(), 1);
+        // Manually simulate greedy: prompt=[1](len=1, odd), so first pick is lp[3].
+        // Sequence: [1, 3, 7, 3, ...] until min_length=3 allows EOS but EOS < 0.5 vs -0.1 so it
+        // never fires. We just verify the length constraint held.
+        let generated = results[0].tokens.len() - 1;
+        assert!(generated >= 3 || results[0].finished);
+    }
+
+    #[test]
+    fn test_decode_prefers_higher_logprob() {
+        // With two beams and clear score separation, the higher-lp beam should rank first.
+        let vocab_size = 10;
+        let eos_id = 2_u32;
+        let config = BeamSearchConfig {
+            beam_width: 2,
+            max_length: 2,
+            eos_token_id: eos_id,
+            length_penalty_alpha: 0.0,
+            min_length: 0,
+            no_repeat_ngram_size: 0,
+        };
+        let decoder = BeamSearchDecoder::new(config);
+        // Always give token 5 a much better score.
+        let score_fn = |_: &[u32]| {
+            let mut lp = vec![-20.0_f32; vocab_size];
+            lp[5] = -0.01; // very good
+            lp[3] = -5.0;  // much worse
+            lp[eos_id as usize] = -0.5;
+            lp
+        };
+        let results = decoder.decode(&[1], vocab_size, 2, score_fn);
+        // The top beam should contain token 5.
+        assert!(
+            results[0].tokens.contains(&5),
+            "Top beam should contain token 5"
+        );
+    }
+
+    #[test]
+    fn test_default_config() {
+        let cfg = BeamSearchConfig::default();
+        assert_eq!(cfg.beam_width, 4);
+        assert_eq!(cfg.max_length, 128);
+        assert_eq!(cfg.eos_token_id, 2);
+        assert!((cfg.length_penalty_alpha - 0.6).abs() < 1e-6);
+        assert_eq!(cfg.min_length, 1);
+        assert_eq!(cfg.no_repeat_ngram_size, 0);
+    }
+
+    // ---- DiverseBeamSearchDecoder --------------------------------------------
+
+    #[test]
+    fn test_diverse_decode_returns_results() {
+        let vocab_size = 20;
+        let config = BeamSearchConfig {
+            beam_width: 4,
+            max_length: 8,
+            eos_token_id: 2,
+            length_penalty_alpha: 0.6,
+            min_length: 1,
+            no_repeat_ngram_size: 0,
+        };
+        let decoder = DiverseBeamSearchDecoder::new(config, 2, 1.0);
+        let results = decoder.decode(&[1], vocab_size, 2, make_score_fn(vocab_size, 2));
+        assert!(
+            !results.is_empty(),
+            "Diverse beam search should return at least one beam"
+        );
+    }
+
+    #[test]
+    fn test_diverse_decode_group1_equals_standard() {
+        // With num_groups=1 and diversity_penalty=0, diverse beam search should be
+        // equivalent to standard beam search.
+        let vocab_size = 20;
+        let config = BeamSearchConfig {
+            beam_width: 4,
+            max_length: 6,
+            eos_token_id: 2,
+            length_penalty_alpha: 0.0,
+            min_length: 1,
+            no_repeat_ngram_size: 0,
+        };
+        let std_decoder = BeamSearchDecoder::new(config.clone());
+        let div_decoder = DiverseBeamSearchDecoder::new(config, 1, 0.0);
+
+        let sf = make_score_fn(vocab_size, 2);
+        let sf2 = make_score_fn(vocab_size, 2);
+
+        let std_results = std_decoder.decode(&[1], vocab_size, 1, sf);
+        let div_results = div_decoder.decode(&[1], vocab_size, 1, sf2);
+
+        assert!(!std_results.is_empty());
+        assert!(!div_results.is_empty());
+        // Both should produce the same top beam sequence.
+        assert_eq!(
+            std_results[0].tokens, div_results[0].tokens,
+            "1-group diverse should match standard beam search"
+        );
+    }
+}