ferrox_models/sampling/

rng.rs

//! The seeded generator every draw in a generation comes off, and the
//! entry points that use it.
//!
//! Split out of `sampling.rs` so the chain runner beside it stays about
//! the chain. The RNG is one concept and a small one, but it is the
//! concept the whole run's reproducibility rests on: a request that
//! passed `seed: 42` must draw the same tokens on every machine, so
//! every draw -- the token, speculative decoding's accept coin, and
//! XTC's -- has to come off this one stream in this one order.

use super::penalties::apply_history_penalties;
use super::{filtered_distribution, greedy_choice, SamplingParams};
use crate::penalty_window::PenaltyWindow;

/// Sets logits a caller wants to forbid to `-inf`, in place, before the
/// sampler looks at them.
///
/// Two callers today, and they COMPOSE rather than exclude each other --
/// a masked logit stays masked, so the order they run in cannot matter:
/// JSON-object mode's character-class filter
/// (`ferrox_server::json_mode`), and grammar-constrained decoding
/// ([`crate::grammar_sampler::GrammarSampler::mask_logits`]).
///
/// The signature returns nothing because the callback runs from inside
/// the sampler, which has no error to return one through. A mask that
/// CAN fail -- a grammar that dead-ends leaves every logit at `-inf`,
/// and sampling from that is how an "impossible" request becomes
/// arbitrary text with a 200 -- records its refusal in the closure's own
/// captured state, and the decode loop reads it after the sample and
/// throws the token away. `ferrox_server::sample_step::sample_next` is
/// the one place that pairing lives.
pub type LogitMask<'a> = &'a mut dyn FnMut(&mut [f32]);

/// A small, seedable xorshift64* generator. Not cryptographically
/// secure -- sampling doesn't need that -- but reproducible given a
/// seed, which greedy argmax already was for free.
pub struct Sampler {
    state: u64,
}

impl Sampler {
    pub fn new(seed: u64) -> Self {
        // xorshift64* requires a nonzero seed.
        Sampler {
            state: if seed == 0 { 0x9E3779B97F4A7C15 } else { seed },
        }
    }

    fn next_u64(&mut self) -> u64 {
        self.state ^= self.state << 13;
        self.state ^= self.state >> 7;
        self.state ^= self.state << 17;
        // The `*` in xorshift64*. Without it this is plain xorshift64,
        // whose state IS its output, and a small seed's first output is
        // therefore still small: for every seed below ~4000 the first
        // draw landed in the bottom eighth of [0, 1), so a request that
        // asked for `seed: 42` always got its first token from the
        // bottom of the CDF. The multiply is what decorrelates the
        // output from a low-entropy state; see
        // `low_seeds_do_not_bias_the_first_draw`.
        self.state.wrapping_mul(0x2545F491_4F6CDD1D)
    }

    /// Uniform float in [0.0, 1.0).
    fn next_f32(&mut self) -> f32 {
        (self.next_u64() >> 40) as f32 / (1u64 << 24) as f32
    }

    /// The one uniform draw XTC needs for this token, or `None` when XTC
    /// cannot fire for these parameters.
    ///
    /// XTC is the only sampler in the chain that is itself stochastic
    /// (`llama_sample_xtc_apply` draws from its own `std::mt19937`,
    /// `src/llama-sampler.cpp:2146`), which is why the chain below takes
    /// the roll as an argument instead of owning an RNG: the chain is
    /// also what `sampling_distribution` runs for speculative
    /// verification, and a filter that drew its own randomness there
    /// would make "the distribution the sampler draws from" a different
    /// distribution every time it was asked.
    ///
    /// **The draw is skipped when XTC cannot fire**, and that is not an
    /// optimisation. Every draw advances the seeded stream, so drawing
    /// unconditionally would shift every subsequent token of every
    /// existing seeded generation, on every run that never asked for
    /// XTC. [`SamplingParams::xtc_can_fire`] is the single predicate
    /// this and [`Candidates::xtc`] share.
    pub fn xtc_roll(&mut self, params: &SamplingParams) -> Option<f32> {
        if params.xtc_can_fire() {
            Some(self.next_f32())
        } else {
            None
        }
    }

    /// Samples one token id from `logits`, given `params` and the
    /// [`PenaltyWindow`] the penalties look back over. Falls back to
    /// plain greedy argmax when `params.temperature <= 0.0`.
    ///
    /// `history` is a window and not a slice on purpose: it carries the
    /// PROMPT as well as the generated tokens, which is what llama.cpp
    /// penalises over. See [`crate::penalty_window`].
    ///
    /// A length-1 `logits` vector is treated as a precomputed greedy token
    /// id (`logits[0] as usize`) — used by the Metal dense-stack path that
    /// returns GPU argmax instead of downloading the full vocab.
    pub fn sample(
        &mut self,
        logits: &[f32],
        params: &SamplingParams,
        history: PenaltyWindow<'_>,
    ) -> usize {
        self.sample_with_mask(logits, params, history, None)
    }

    /// Like [`Self::sample`], but optionally zeroes disallowed logits via
    /// `mask` before argmax / nucleus sampling (used for JSON-object mode).
    pub fn sample_with_mask(
        &mut self,
        logits: &[f32],
        params: &SamplingParams,
        history: PenaltyWindow<'_>,
        mut mask: Option<LogitMask<'_>>,
    ) -> usize {
        let xtc_roll = self.xtc_roll(params);
        if params.temperature <= 0.0 && mask.is_none() {
            if logits.len() == 1 {
                return logits[0] as usize;
            }
            let mut scores = logits.to_vec();
            apply_history_penalties(&mut scores, params, history);
            return greedy_choice(scores, params, history, xtc_roll);
        }

        let mut scores: Vec<f32> = logits.to_vec();
        apply_history_penalties(&mut scores, params, history);

        if let Some(m) = mask.as_mut() {
            m(&mut scores);
        }

        if params.temperature <= 0.0 {
            if scores.len() == 1 {
                return scores[0] as usize;
            }
            return greedy_choice(scores, params, history, xtc_roll);
        }

        let probs = filtered_distribution(scores, params, history, xtc_roll);
        self.sample_from(&probs)
    }

    /// A uniform draw in `[0.0, 1.0)`.
    ///
    /// Exposed because speculative decoding's accept test is a coin
    /// flip against `p_target(x) / p_draft(x)` rather than a draw from
    /// a distribution, and it must come off the same seeded stream as
    /// every other draw in the run or a "reproducible given a seed"
    /// generation stops being reproducible.
    pub fn uniform(&mut self) -> f32 {
        self.next_f32()
    }

    /// Draws one index from an already-normalised distribution.
    ///
    /// Split out of [`Self::sample_with_mask`] so speculative decoding
    /// can sample from a distribution it had to compute anyway (the
    /// rejection rule needs `p_target` itself, not just a draw from it)
    /// and still go through *exactly* the same draw as ordinary
    /// sampling. Two separate copies of this loop would be two chances
    /// to be subtly non-lossless.
    pub fn sample_from(&mut self, probs: &[f32]) -> usize {
        let draw = self.next_f32();
        let mut cumulative = 0.0f32;
        for (i, &p) in probs.iter().enumerate() {
            cumulative += p;
            if draw < cumulative {
                return i;
            }
        }
        // Floating-point rounding may leave `draw` fractionally above
        // the final cumulative sum; the last nonzero-probability token
        // is the correct fallback, not index 0.
        probs
            .iter()
            .enumerate()
            .rev()
            .find(|&(_, &p)| p > 0.0)
            .map(|(i, _)| i)
            .unwrap_or(0)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::sampling::{sampling_distribution, spread_logits};

    ///
    /// This is the reason [`Sampler::xtc_roll`] is conditional rather
    /// than unconditional. Delete the `xtc_can_fire` guard there and
    /// this goes red on the first token.
    #[test]
    fn a_chain_without_xtc_does_not_consume_a_draw_for_it() {
        let params = SamplingParams {
            temperature: 1.0,
            ..SamplingParams::default()
        };
        let logits = spread_logits(32);
        assert!(
            Sampler::new(99).xtc_roll(&params).is_none(),
            "the guard must refuse the draw, not merely ignore it"
        );

        // ONE draw per token, taken by hand off a generator that never
        // heard of XTC. An unconditional roll makes `sample` consume
        // two values per token, so the very first token comes off the
        // SECOND draw and this diverges immediately.
        let mut sampled_by_chain = Sampler::new(99);
        let mut by_hand = Sampler::new(99);
        for step in 0..16 {
            let sampled = sampled_by_chain.sample(&logits, &params, PenaltyWindow::new(&[], &[]));
            let probs = sampling_distribution(&logits, &params, PenaltyWindow::new(&[], &[]), None);
            assert_eq!(
                sampled,
                by_hand.sample_from(&probs),
                "token {step} came off a different position in the stream"
            );
        }
        // And the two generators are still in lockstep afterwards.
        assert_eq!(sampled_by_chain.uniform(), by_hand.uniform());
    }

    /// **The flag does something.** A chain that runs the temperature
    /// before top-p keeps a different candidate set than the default,
    /// which is the whole reason the order is worth exposing -- and the
    /// reason getting it wrong is a silent quality regression rather
    /// than an error.
    ///
    /// A hot temperature flattens the distribution, so a top-p applied
    /// after it sums smaller probabilities and reaches `p` later,
    /// keeping MORE candidates.

    #[test]
    fn temperature_zero_accepts_precomputed_argmax_singleton() {
        let mut sampler = Sampler::new(1);
        let params = SamplingParams::default();
        assert_eq!(
            sampler.sample(&[42.0], &params, PenaltyWindow::new(&[], &[])),
            42
        );
        // Non-greedy must not treat a singleton as a token id.
        let sampled = SamplingParams {
            temperature: 0.8,
            ..SamplingParams::default()
        };
        // Softmax of a single logit → only token 0 is eligible.
        assert_eq!(
            sampler.sample(&[42.0], &sampled, PenaltyWindow::new(&[], &[])),
            0
        );
    }

    #[test]
    fn temperature_zero_is_deterministic_greedy_argmax() {
        let logits = vec![0.1, 0.9, 0.3, -0.2];
        let params = SamplingParams::default();
        let mut sampler = Sampler::new(42);
        assert_eq!(
            sampler.sample(&logits, &params, PenaltyWindow::new(&[], &[])),
            1
        );
        // Must be deterministic regardless of RNG state advancing.
        assert_eq!(
            sampler.sample(&logits, &params, PenaltyWindow::new(&[], &[])),
            1
        );
    }

    #[test]
    fn high_temperature_can_pick_a_non_argmax_token_over_many_draws() {
        let logits = vec![1.0, 1.0, 1.0, 1.0];
        let params = SamplingParams {
            temperature: 1.0,
            ..SamplingParams::default()
        };
        let mut sampler = Sampler::new(7);
        let mut seen = std::collections::HashSet::new();
        for _ in 0..200 {
            seen.insert(sampler.sample(&logits, &params, PenaltyWindow::new(&[], &[])));
        }
        assert!(
            seen.len() > 1,
            "uniform logits at temperature=1.0 must produce more than one distinct token across 200 draws"
        );
    }

    #[test]
    fn top_k_one_is_equivalent_to_greedy() {
        let logits = vec![0.1, 0.9, 0.3, -0.2];
        let params = SamplingParams {
            temperature: 1.0,
            top_k: 1,
            ..SamplingParams::default()
        };
        let mut sampler = Sampler::new(123);
        for _ in 0..20 {
            assert_eq!(
                sampler.sample(&logits, &params, PenaltyWindow::new(&[], &[])),
                1
            );
        }
    }

    #[test]
    fn top_p_near_zero_is_equivalent_to_greedy() {
        let logits = vec![0.1, 5.0, 0.3, -0.2];
        let params = SamplingParams {
            temperature: 1.0,
            top_p: 0.001,
            ..SamplingParams::default()
        };
        let mut sampler = Sampler::new(9);
        for _ in 0..20 {
            assert_eq!(
                sampler.sample(&logits, &params, PenaltyWindow::new(&[], &[])),
                1
            );
        }
    }

    #[test]
    fn presence_and_frequency_penalties_reduce_seen_token_logits() {
        let logits = vec![0.0, 5.0, 0.0];
        let params = SamplingParams {
            temperature: 1.0,
            presence_penalty: 10.0,
            frequency_penalty: 0.0,
            ..SamplingParams::default()
        };
        let mut sampler = Sampler::new(1);
        let mut counts = [0usize; 3];
        for _ in 0..500 {
            counts[sampler.sample(&logits, &params, PenaltyWindow::new(&[], &[1]))] += 1;
        }
        assert!(
            counts[1] < 250,
            "presence_penalty should discourage token 1; counts={counts:?}"
        );

        let params = SamplingParams {
            temperature: 1.0,
            presence_penalty: 0.0,
            frequency_penalty: 10.0,
            ..SamplingParams::default()
        };
        let mut sampler = Sampler::new(2);
        counts = [0; 3];
        for _ in 0..500 {
            counts[sampler.sample(&logits, &params, PenaltyWindow::new(&[], &[1, 1, 1]))] += 1;
        }
        assert!(
            counts[1] < 250,
            "frequency_penalty should discourage repeated token 1; counts={counts:?}"
        );
    }

    #[test]
    fn repetition_penalty_reduces_probability_of_recently_seen_token() {
        let logits = vec![0.0, 5.0, 0.0];
        let params = SamplingParams {
            temperature: 1.0,
            repetition_penalty: 1000.0,
            ..SamplingParams::default()
        };
        let mut sampler = Sampler::new(3);
        let mut counts = [0usize; 3];
        for _ in 0..500 {
            counts[sampler.sample(&logits, &params, PenaltyWindow::new(&[], &[1]))] += 1;
        }
        assert!(
            counts[1] < 250,
            "heavily penalizing token 1 (already in history) should make it far less likely than its raw logit alone would suggest; got counts={counts:?}"
        );
    }

    #[test]
    fn low_seeds_do_not_bias_the_first_draw() {
        // Every generation seeds a fresh `Sampler` (the server does it
        // per request, from the caller's `seed`), so the FIRST draw off
        // a freshly seeded generator is the one users actually see.
        // Plain xorshift64 returns its own state, so seeds 1..4000 all
        // produced a first draw in the bottom eighth of [0, 1) -- the
        // first sampled token of every seeded request came off the
        // bottom of the CDF.
        let vocab = 8;
        let logits = vec![0.0f32; vocab];
        let params = SamplingParams {
            temperature: 1.0,
            ..SamplingParams::default()
        };
        let seeds = 4_000u64;
        let mut counts = vec![0usize; vocab];
        for seed in 1..=seeds {
            counts[Sampler::new(seed).sample(&logits, &params, PenaltyWindow::new(&[], &[]))] += 1;
        }
        let expected = seeds as f64 / vocab as f64;
        for (token, &c) in counts.iter().enumerate() {
            assert!(
                (c as f64 - expected).abs() < expected * 0.25,
                "uniform logits: token {token} came up {c} times across {seeds} seeds, \
                 expected about {expected:.0} (counts={counts:?})"
            );
        }
    }

    #[test]
    fn the_published_distribution_is_the_one_sample_actually_draws_from() {
        // `sampling_distribution` is load-bearing for lossless
        // speculative verification: if it disagreed with what `sample`
        // draws from, every accept/reject decision would be measured
        // against the wrong target. Check them against each other
        // empirically, with filters on so the two code paths have
        // something to disagree about.
        let logits = vec![0.4, 2.0, -1.0, 1.2, 0.9, -0.3];
        let params = SamplingParams {
            temperature: 0.8,
            top_p: 0.9,
            top_k: 4,
            repetition_penalty: 1.3,
            ..SamplingParams::default()
        };
        let history = [1usize, 4];
        let claimed =
            sampling_distribution(&logits, &params, PenaltyWindow::new(&[], &history), None);
        assert!((claimed.iter().sum::<f32>() - 1.0).abs() < 1e-5);

        let draws = 100_000;
        let mut counts = vec![0usize; logits.len()];
        let mut sampler = Sampler::new(0xC0FFEE);
        for _ in 0..draws {
            counts[sampler.sample(&logits, &params, PenaltyWindow::new(&[], &history))] += 1;
        }
        for (i, &c) in counts.iter().enumerate() {
            let empirical = c as f64 / draws as f64;
            assert!(
                (empirical - claimed[i] as f64).abs() < 0.01,
                "token {i}: sample() draws it {empirical:.4} of the time but \
                 sampling_distribution claims {:.4}",
                claimed[i]
            );
        }
    }

    #[test]
    fn degenerate_all_zero_probability_falls_back_to_greedy() {
        // top_k=1 combined with a top_p that would exclude even that
        // one surviving token is a contradictory/degenerate
        // configuration; must not panic or sample index 0 blindly.
        let logits = vec![0.1, 0.9, 0.3, -0.2];
        let params = SamplingParams {
            temperature: 1.0,
            top_k: 1,
            top_p: 1.0,
            ..SamplingParams::default()
        };
        let mut sampler = Sampler::new(1);
        assert_eq!(
            sampler.sample(&logits, &params, PenaltyWindow::new(&[], &[])),
            1
        );
    }
}