metrale_model_engine/traits/
logprobs.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Log-softmax and top-k logprob extraction, used by the scheduler's decode-time
4//! logprobs (`logprob_of`) and by the prompt-logprob collection during prefill (`extract_bf16`).
5//!
6//! Owner: model-engine.
7//! Invariants: none beyond the types.
8
9/// 2026-09-25: One scored prompt position: `log P(tokens[i+1] | tokens[..=i])` plus the top-k
10/// alternatives under the same distribution.
11#[derive(Clone, Debug)]
12pub struct PromptTokenLogprob {
13    /// 2026-09-25: The next prompt token, the one scored.
14    pub token_id: u32,
15    pub logprob: f32,
16    /// 2026-09-25: Top-k `(token_id, logprob)` alternatives, highest first; empty when `k == 0`.
17    pub top: Vec<(u32, f32)>,
18}
19
20/// 2026-09-25: Log-softmax over `f32_logits`: the target token's logprob and the top-k
21/// alternatives, highest first. `k == 0` returns no alternatives. A target outside the vocab,
22/// or empty logits, gives `-inf` rather than a panic.
23pub fn logprob_of(f32_logits: &[f32], target: u32, k: usize) -> (f32, Vec<(u32, f32)>) {
24    if f32_logits.is_empty() {
25        return (f32::NEG_INFINITY, Vec::new());
26    }
27    let max_logit = f32_logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
28    let log_sum_exp = max_logit
29        + f32_logits
30            .iter()
31            .map(|&l| (l - max_logit).exp())
32            .sum::<f32>()
33            .ln();
34    let target_logprob = if (target as usize) < f32_logits.len() {
35        f32_logits[target as usize] - log_sum_exp
36    } else {
37        f32::NEG_INFINITY
38    };
39    if k == 0 {
40        return (target_logprob, Vec::new());
41    }
42    let mut indexed: Vec<(u32, f32)> = f32_logits
43        .iter()
44        .enumerate()
45        .map(|(j, &l)| (j as u32, l - log_sum_exp))
46        .collect();
47    let nth = k.min(indexed.len().saturating_sub(1));
48    indexed.select_nth_unstable_by(nth, |a, b| {
49        b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)
50    });
51    let mut top: Vec<(u32, f32)> = indexed[..k.min(indexed.len())].to_vec();
52    top.sort_unstable_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
53    (target_logprob, top)
54}
55
56/// 2026-09-25: BF16 to FP32: the two little-endian bytes become the upper 16 bits of an f32.
57#[inline]
58pub fn bf16_to_f32(lo: u8, hi: u8) -> f32 {
59    f32::from_bits(((lo as u32) | ((hi as u32) << 8)) << 16)
60}
61
62/// 2026-09-25: One position's `PromptTokenLogprob` from a row of `vocab` BF16 logits.
63pub fn extract_bf16(bf16: &[u8], target: u32, k: usize, vocab: usize) -> PromptTokenLogprob {
64    let f32_logits: Vec<f32> = (0..vocab)
65        .map(|j| bf16_to_f32(bf16[j * 2], bf16[j * 2 + 1]))
66        .collect();
67    let (logprob, top) = logprob_of(&f32_logits, target, k);
68    PromptTokenLogprob {
69        token_id: target,
70        logprob,
71        top,
72    }
73}
74
75#[cfg(test)]
76mod tests {
77    use super::*;
78
79    #[test]
80    fn logprob_matches_manual_log_softmax() {
81        let logits = [1.0f32, 2.0, 3.0];
82        let sum: f32 = logits.iter().map(|l| l.exp()).sum();
83        let expect = 2.0 - sum.ln();
84        let (lp, top) = logprob_of(&logits, 1, 2);
85        assert!((lp - expect).abs() < 1e-6, "{lp} vs {expect}");
86        assert_eq!(top[0].0, 2);
87        assert_eq!(top[1].0, 1);
88        assert!(top[0].1 > top[1].1);
89    }
90
91    #[test]
92    fn k_zero_returns_empty_top() {
93        let (lp, top) = logprob_of(&[0.0, 1.0], 0, 0);
94        assert!(top.is_empty());
95        assert!(lp < 0.0);
96    }
97
98    #[test]
99    fn out_of_vocab_target_is_neg_inf_not_panic() {
100        let (lp, _) = logprob_of(&[0.0, 1.0], 99, 1);
101        assert_eq!(lp, f32::NEG_INFINITY);
102    }
103
104    #[test]
105    fn empty_vocab_is_neg_inf_with_no_alternatives() {
106        assert_eq!(logprob_of(&[], 0, 4), (f32::NEG_INFINITY, vec![]));
107    }
108
109    #[test]
110    fn bf16_slice_roundtrip_extract() {
111        let vals = [0.0f32, 1.0, 2.0];
112        let mut bytes = Vec::new();
113        for v in vals {
114            let b = (v.to_bits() >> 16) as u16;
115            bytes.push((b & 0xFF) as u8);
116            bytes.push((b >> 8) as u8);
117        }
118        let r = extract_bf16(&bytes, 2, 1, 3);
119        let sum: f32 = vals.iter().map(|l| l.exp()).sum();
120        let expect = 2.0 - sum.ln();
121        assert_eq!(r.token_id, 2);
122        assert!((r.logprob - expect).abs() < 1e-3);
123        assert_eq!(r.top[0].0, 2);
124        assert!((r.top[0].1 - expect).abs() < 1e-3);
125    }
126}
metrale_model_engine/traits/logprobs.rs

metrale_model_engine/traits/
logprobs.rs
