spark_model/traits/
logprobs.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Pure log-softmax / top-k logprob extraction. SSOT for the math shared
4//! by the scheduler's decode-time logprobs (`spark-server`) and the
5//! model's prompt-logprob collection during prefill (legacy
6//! `/v1/completions` `echo` + `logprobs`).
7
8/// One scored prompt position: `log P(tokens[i+1] | tokens[..=i])` plus
9/// the top-k alternative tokens under the same distribution.
10#[derive(Clone, Debug)]
11pub struct PromptTokenLogprob {
12    /// The actual next prompt token (the one being scored).
13    pub token_id: u32,
14    /// Log-probability of `token_id` under this position's distribution.
15    pub logprob: f32,
16    /// Top-k `(token_id, logprob)` alternatives, sorted descending.
17    /// Empty when k=0 (logprob-only mode).
18    pub top: Vec<(u32, f32)>,
19}
20
21/// Log-softmax over `f32_logits`; returns the target token's logprob and
22/// the top-k alternatives sorted descending by logprob. `k=0` returns an
23/// empty top vector (chosen-token logprob only). A target outside the
24/// vocab yields `-inf` (fail-visible, never panics).
25pub fn logprob_of(f32_logits: &[f32], target: u32, k: usize) -> (f32, Vec<(u32, f32)>) {
26    if f32_logits.is_empty() {
27        return (f32::NEG_INFINITY, Vec::new());
28    }
29    // Log-softmax: logprob = logit - log(sum(exp(logits)))
30    let max_logit = f32_logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
31    let log_sum_exp = max_logit
32        + f32_logits
33            .iter()
34            .map(|&l| (l - max_logit).exp())
35            .sum::<f32>()
36            .ln();
37    let target_logprob = if (target as usize) < f32_logits.len() {
38        f32_logits[target as usize] - log_sum_exp
39    } else {
40        f32::NEG_INFINITY
41    };
42    if k == 0 {
43        return (target_logprob, Vec::new());
44    }
45    // Top-K by partial sort.
46    let mut indexed: Vec<(u32, f32)> = f32_logits
47        .iter()
48        .enumerate()
49        .map(|(j, &l)| (j as u32, l - log_sum_exp))
50        .collect();
51    let nth = k.min(indexed.len().saturating_sub(1));
52    indexed.select_nth_unstable_by(nth, |a, b| {
53        b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)
54    });
55    let mut top: Vec<(u32, f32)> = indexed[..k.min(indexed.len())].to_vec();
56    top.sort_unstable_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
57    (target_logprob, top)
58}
59
60/// BF16 → FP32 (upper 16 bits of the IEEE-754 f32 pattern).
61#[inline]
62pub fn bf16_to_f32(lo: u8, hi: u8) -> f32 {
63    f32::from_bits(((lo as u32) | ((hi as u32) << 8)) << 16)
64}
65
66/// Extract one position's `PromptTokenLogprob` from a BF16 `[vocab]`
67/// logits slice (the layout `lm_head_batched` writes per row).
68pub fn extract_bf16(bf16: &[u8], target: u32, k: usize, vocab: usize) -> PromptTokenLogprob {
69    let f32_logits: Vec<f32> = (0..vocab)
70        .map(|j| bf16_to_f32(bf16[j * 2], bf16[j * 2 + 1]))
71        .collect();
72    let (logprob, top) = logprob_of(&f32_logits, target, k);
73    PromptTokenLogprob {
74        token_id: target,
75        logprob,
76        top,
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83
84    #[test]
85    fn logprob_matches_manual_log_softmax() {
86        let logits = [1.0f32, 2.0, 3.0];
87        let sum: f32 = logits.iter().map(|l| l.exp()).sum();
88        let expect = 2.0 - sum.ln();
89        let (lp, top) = logprob_of(&logits, 1, 2);
90        assert!((lp - expect).abs() < 1e-6, "{lp} vs {expect}");
91        // Top-2 sorted descending: token 2 first, then token 1.
92        assert_eq!(top[0].0, 2);
93        assert_eq!(top[1].0, 1);
94        assert!(top[0].1 > top[1].1);
95    }
96
97    #[test]
98    fn k_zero_returns_empty_top() {
99        let (lp, top) = logprob_of(&[0.0, 1.0], 0, 0);
100        assert!(top.is_empty());
101        assert!(lp < 0.0);
102    }
103
104    #[test]
105    fn out_of_vocab_target_is_neg_inf_not_panic() {
106        let (lp, _) = logprob_of(&[0.0, 1.0], 99, 1);
107        assert_eq!(lp, f32::NEG_INFINITY);
108    }
109
110    #[test]
111    fn empty_vocab_is_neg_inf_with_no_alternatives() {
112        assert_eq!(logprob_of(&[], 0, 4), (f32::NEG_INFINITY, vec![]));
113    }
114
115    #[test]
116    fn bf16_slice_roundtrip_extract() {
117        // BF16 encodings of [0.0, 1.0, 2.0]: f32 bit patterns >> 16.
118        let vals = [0.0f32, 1.0, 2.0];
119        let mut bytes = Vec::new();
120        for v in vals {
121            let b = (v.to_bits() >> 16) as u16;
122            bytes.push((b & 0xFF) as u8);
123            bytes.push((b >> 8) as u8);
124        }
125        let r = extract_bf16(&bytes, 2, 1, 3);
126        let sum: f32 = vals.iter().map(|l| l.exp()).sum();
127        let expect = 2.0 - sum.ln();
128        assert_eq!(r.token_id, 2);
129        assert!((r.logprob - expect).abs() < 1e-3);
130        assert_eq!(r.top[0].0, 2);
131        assert!((r.top[0].1 - expect).abs() < 1e-3);
132    }
133}
spark_model/traits/logprobs.rs

spark_model/traits/
logprobs.rs
