metrale_sampling/
sample_impl.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The host sampling pipeline over little-endian f32 logits:
4//! penalties and logit bias, then either greedy argmax or top-n-sigma,
5//! temperature, top-k, min-p, top-p and a multinomial draw.
6//!
7//! Owner: sampling.
8//! Invariants:
9//! - `apply_penalties_and_bias` runs before the greedy bypass, so penalties
10//!   and bias also apply at temperature 0.
11
12use super::{SamplingParams, apply_penalties_and_bias, record_entropy};
13
14pub fn sample_with_params_history(
15    data: &[u8],
16    params: &SamplingParams,
17    token_history: &[u32],
18) -> u32 {
19    sample_with_params_seeded(data, params, token_history, params.seed)
20}
21
22/// 2026-09-25: Samples one token id from the little-endian f32 logits in
23/// `data`. `temperature <= 0.0` returns the greedy argmax. Otherwise
24/// `Some(seed)` draws from `StdRng::seed_from_u64(seed)` and `None` from
25/// `rand::random`.
26pub fn sample_with_params_seeded(
27    data: &[u8],
28    params: &SamplingParams,
29    token_history: &[u32],
30    seed: Option<u64>,
31) -> u32 {
32    let n = data.len() / 4;
33    let top_k = params.top_k as usize;
34    let top_p = params.top_p;
35    let top_n_sigma = params.top_n_sigma;
36    let min_p = params.min_p;
37
38    let mut raw_logits: Vec<f32> = data
39        .chunks_exact(4)
40        .take(n)
41        .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
42        .collect();
43
44    // 2026-09-25: Penalties and logit bias go first, before the greedy
45    // bypass, so they apply at temperature 0 too.
46    apply_penalties_and_bias(&mut raw_logits, params, token_history);
47
48    // 2026-09-25: Greedy returns the argmax of the penalised logits and skips
49    // every later stage. Those stages only drop tokens below the maximum or
50    // apply a monotonic scale, so none of them can move the argmax.
51    if params.temperature <= 0.0 {
52        return greedy_pick_last_wins(&raw_logits);
53    }
54    let temperature = params.temperature;
55
56    // 2026-09-25: Top-n-sigma, before temperature scaling: logits below
57    // mean - n*sigma become -inf.
58    if top_n_sigma > 0.0 {
59        let sum: f32 = raw_logits.iter().sum();
60        let mean = sum / n as f32;
61        let var: f32 = raw_logits.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / n as f32;
62        let sigma = var.sqrt();
63        if sigma > 0.0 {
64            let threshold = mean - top_n_sigma * sigma;
65            for logit in raw_logits.iter_mut() {
66                if *logit < threshold {
67                    *logit = f32::NEG_INFINITY;
68                }
69            }
70        }
71    }
72
73    let mut logits: Vec<(u32, f32)> = raw_logits
74        .iter()
75        .enumerate()
76        .filter(|(_, v)| v.is_finite())
77        .map(|(i, v)| (i as u32, v / temperature))
78        .collect();
79
80    if logits.is_empty() {
81        return raw_logits
82            .iter()
83            .enumerate()
84            .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
85            .map(|(i, _)| i as u32)
86            .unwrap_or(0);
87    }
88
89    // 2026-09-25: Top-k, min-p and top-p need descending order. With a top-k
90    // cap, quickselect isolates the k largest in O(n) and only those k are
91    // sorted. With none of the three, nothing is sorted: the distribution of
92    // the multinomial draw does not depend on order.
93    let cmp_desc =
94        |a: &(u32, f32), b: &(u32, f32)| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal);
95    let sorted = if top_k > 0 && top_k < logits.len() {
96        logits.select_nth_unstable_by(top_k, cmp_desc);
97        logits.truncate(top_k);
98        logits.sort_unstable_by(cmp_desc);
99        true
100    } else if min_p > 0.0 || top_p < 1.0 {
101        logits.sort_unstable_by(cmp_desc);
102        true
103    } else {
104        false
105    };
106
107    // 2026-09-25: `sorted` means logits[0] is the maximum; otherwise it is
108    // reduced for. The min-p and top-p blocks below run only when `sorted` is
109    // true, so they can rely on descending order.
110    let max_val = if sorted {
111        logits[0].1
112    } else {
113        logits
114            .iter()
115            .map(|&(_, v)| v)
116            .fold(f32::NEG_INFINITY, f32::max)
117    };
118    let mut probs: Vec<(u32, f32)> = logits
119        .iter()
120        .map(|&(idx, logit)| (idx, (logit - max_val).exp()))
121        .collect();
122
123    // 2026-09-25: Entropy H = -Σ p·ln(p), in nats, of the softmax over the
124    // candidates left after top-k and before min-p and top-p.
125    {
126        let sum: f32 = probs.iter().map(|p| p.1).sum();
127        if sum > 0.0 {
128            let inv = 1.0 / sum;
129            let h: f32 = probs
130                .iter()
131                .map(|&(_, w)| {
132                    let p = w * inv;
133                    if p > 1e-10 { -p * p.ln() } else { 0.0 }
134                })
135                .sum();
136            record_entropy(h);
137        }
138    }
139
140    if min_p > 0.0 {
141        let max_prob = probs[0].1;
142        let threshold = min_p * max_prob;
143        probs.retain(|p| p.1 >= threshold);
144    }
145
146    if top_p < 1.0 {
147        let sum: f32 = probs.iter().map(|p| p.1).sum();
148        let mut cumsum = 0.0f32;
149        let mut cutoff = probs.len();
150        for (i, &(_, prob)) in probs.iter().enumerate() {
151            cumsum += prob / sum;
152            if cumsum >= top_p {
153                cutoff = i + 1;
154                break;
155            }
156        }
157        probs.truncate(cutoff);
158    }
159
160    let sum: f32 = probs.iter().map(|p| p.1).sum();
161    let random_val: f32 = if let Some(s) = seed {
162        use rand::Rng;
163        use rand::SeedableRng;
164        let mut rng = rand::rngs::StdRng::seed_from_u64(s);
165        rng.random::<f32>()
166    } else {
167        rand::random::<f32>()
168    };
169    let threshold = random_val * sum;
170    let mut cumsum = 0.0f32;
171    for &(idx, prob) in &probs {
172        cumsum += prob;
173        if cumsum >= threshold {
174            return idx;
175        }
176    }
177    probs.last().map_or(0, |p| p.0)
178}
179
180/// 2026-09-25: Returns what `max_by(partial_cmp.unwrap_or(Equal))` returns:
181/// the last of several equal maxima. The verify path's `argmax_first_wins`
182/// takes the first, so the two are not interchangeable on exact ties.
183///
184/// A NaN compares `Equal`, and under `max_by` an equal later element
185/// displaces the current maximum, so a NaN after the maximum changes the
186/// result. A max-then-find pass cannot reproduce that. Pass 1 therefore takes
187/// the maximum over 8 lanes while detecting NaN. With any NaN present the
188/// `max_by` expression itself runs; otherwise the result is the last index
189/// equal to the maximum, which is what `max_by` returns on NaN-free input,
190/// -0.0/+0.0 ties included.
191fn greedy_pick_last_wins(v: &[f32]) -> u32 {
192    const LANES: usize = 8;
193    let mut acc = [f32::NEG_INFINITY; LANES];
194    let mut any_nan = false;
195    let mut chunks = v.chunks_exact(LANES);
196    for c in &mut chunks {
197        for (a, &x) in acc.iter_mut().zip(c) {
198            any_nan |= x.is_nan();
199            if x > *a {
200                *a = x;
201            }
202        }
203    }
204    let mut best = f32::NEG_INFINITY;
205    for &a in acc.iter() {
206        if a > best {
207            best = a;
208        }
209    }
210    for &x in chunks.remainder() {
211        any_nan |= x.is_nan();
212        if x > best {
213            best = x;
214        }
215    }
216    if any_nan {
217        return v
218            .iter()
219            .enumerate()
220            .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
221            .map(|(i, _)| i as u32)
222            .unwrap_or(0);
223    }
224    v.iter()
225        .rposition(|&x| x == best)
226        .unwrap_or(0)
227        .try_into()
228        .unwrap_or(0)
229}
230
231#[cfg(test)]
232mod greedy_tests {
233    use super::greedy_pick_last_wins;
234
235    fn reference(v: &[f32]) -> u32 {
236        v.iter()
237            .enumerate()
238            .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
239            .map(|(i, _)| i as u32)
240            .unwrap_or(0)
241    }
242
243    fn agree(v: &[f32]) {
244        assert_eq!(greedy_pick_last_wins(v), reference(v), "diverged on {v:?}");
245    }
246
247    #[test]
248    fn matches_max_by_reference() {
249        agree(&[]);
250        agree(&[1.0]);
251        agree(&[1.0, 3.0, 2.0]);
252        agree(&[1.0, 5.0, 5.0, 5.0, 2.0]);
253        agree(&[5.0, 1.0, 5.0]);
254        // 2026-09-25: a tie across the 8-lane boundary (index 7 in a lane,
255        // index 8 in the remainder).
256        agree(&[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 9.0, 9.0]);
257        agree(&[-5.0, -1.0, -3.0]);
258        agree(&[-0.0, 0.0, -0.0]);
259        agree(&[f32::NEG_INFINITY, f32::NEG_INFINITY]);
260        agree(&[f32::INFINITY, 1.0, f32::INFINITY]);
261        agree(&[f32::NAN, 1.0, 2.0]);
262        agree(&[1.0, 2.0, f32::NAN]);
263        agree(&[f32::NAN]);
264    }
265
266    #[test]
267    fn vocab_sized_last_wins() {
268        let mut v: Vec<f32> = (0..248_320)
269            .map(|i| (((i * 2654435761u64 as usize) % 100_003) as f32) / 1000.0 - 50.0)
270            .collect();
271        v[100_000] = 999.0;
272        v[200_000] = 999.0;
273        assert_eq!(greedy_pick_last_wins(&v), reference(&v));
274        assert_eq!(greedy_pick_last_wins(&v), 200_000);
275    }
276}
metrale_sampling/sample_impl.rs

metrale_sampling/
sample_impl.rs
