metrale_sampling/
lib.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Host-side token sampling: `SamplingParams`, the `Sampler` that
4//! copies device logits to the host, the penalty stage, and the argmax helpers.
5//!
6//! Owner: metrale-sampling.
7//! Invariants:
8//! - In this crate, only the non-greedy path of `sample_with_params_seeded`
9//!   writes the entropy gauges (through `record_entropy`).
10
11use std::sync::atomic::Ordering;
12
13use anyhow::Result;
14use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
15
16// 2026-09-25: The entropy gauges are fields of
17// `metrale_telemetry::run_metrics::RunMetrics`; `reset_for_new_run` zeroes
18// them when a run starts.
19
20/// 2026-09-25: Entropy (nats) of the last non-greedy sample in this run.
21pub fn last_entropy() -> f32 {
22    f32::from_bits(
23        metrale_telemetry::run_metrics::metrics()
24            .last_entropy
25            .load(Ordering::Relaxed),
26    )
27}
28
29/// 2026-09-25: Non-greedy samples in this run whose entropy was below 0.3 nats.
30pub fn low_entropy_token_count() -> u64 {
31    metrale_telemetry::run_metrics::metrics()
32        .low_entropy_tokens
33        .load(Ordering::Relaxed)
34}
35
36/// 2026-09-25: Non-greedy samples in this run that recorded an entropy.
37pub fn total_sampled_token_count() -> u64 {
38    metrale_telemetry::run_metrics::metrics()
39        .total_sampled_tokens
40        .load(Ordering::Relaxed)
41}
42
43pub(crate) fn record_entropy(entropy: f32) {
44    let m = metrale_telemetry::run_metrics::metrics();
45    m.last_entropy.store(entropy.to_bits(), Ordering::Relaxed);
46    m.total_sampled_tokens.fetch_add(1, Ordering::Relaxed);
47    if entropy < 0.3 {
48        m.low_entropy_tokens.fetch_add(1, Ordering::Relaxed);
49    }
50}
51
52/// 2026-09-25: Sampling parameters for a request.
53#[derive(Debug, Clone)]
54pub struct SamplingParams {
55    /// 2026-09-25: Temperature. 0.0 is greedy (`is_greedy`).
56    pub temperature: f32,
57    /// 2026-09-25: Top-k: keep only the k highest-probability tokens.
58    /// 0 disables it.
59    pub top_k: u32,
60    /// 2026-09-25: Top-p (nucleus): keep the smallest set of tokens whose
61    /// cumulative probability reaches p. 1.0 or more disables it.
62    pub top_p: f32,
63    /// 2026-09-25: Top-n-sigma: before temperature scaling, drop tokens whose
64    /// logit is below mean - n*sigma. 0.0 disables it.
65    pub top_n_sigma: f32,
66    /// 2026-09-25: Min-p: keep tokens with prob >= min_p * max_prob
67    /// (post-softmax). 0.0 disables it.
68    pub min_p: f32,
69    /// 2026-09-25: Per-token logit bias as (token_id, bias) pairs, added after
70    /// the penalties and before any filter.
71    pub logit_bias: Vec<(u32, f32)>,
72    /// 2026-09-25: Repetition penalty: for every occurrence of a token in the
73    /// window, a positive logit is divided by it and a non-positive one
74    /// multiplied. 1.0 and values <= 0.0 disable it.
75    pub repetition_penalty: f32,
76    /// 2026-09-25: Window, in tokens, for the repetition, presence and
77    /// frequency penalties. 0 uses the whole history.
78    pub repetition_penalty_window: u32,
79    /// 2026-09-25: Presence penalty: subtracted once from the logit of each
80    /// token that appears in the window. 0.0 disables it.
81    pub presence_penalty: f32,
82    /// 2026-09-25: Frequency penalty: subtracted from a token's logit once per
83    /// occurrence in the window. 0.0 disables it.
84    pub frequency_penalty: f32,
85    /// 2026-09-25: Strength of `apply_lz_penalty`. Values <= 0.0 disable it.
86    pub lz_penalty: f32,
87    /// 2026-09-25: Multiplier of the DRY ("don't repeat yourself") penalty,
88    /// `apply_dry_penalty`. Values <= 0.0 disable it.
89    pub dry_multiplier: f32,
90    /// 2026-09-25: DRY base: penalty = multiplier * base^(match_len - allowed_len).
91    pub dry_base: f32,
92    /// 2026-09-25: DRY allowed length: only matches longer than this are penalised.
93    pub dry_allowed_length: u32,
94    /// 2026-09-25: DRY sequence breakers: token ids at which a match stops.
95    pub dry_sequence_breakers: Vec<u32>,
96    /// 2026-09-25: Maximum tokens to generate.
97    pub max_tokens: usize,
98    /// 2026-09-25: Stop token ids.
99    pub stop_token_ids: Vec<u32>,
100    /// 2026-09-25: RNG seed. `Some` seeds a `StdRng` for each draw; `None`
101    /// uses the thread RNG.
102    pub seed: Option<u64>,
103}
104
105impl SamplingParams {
106    /// 2026-09-25: Greedy sampling with every filter and penalty neutral.
107    pub fn greedy(max_tokens: usize) -> Self {
108        Self {
109            temperature: 0.0,
110            top_k: 0,
111            top_p: 1.0,
112            top_n_sigma: 0.0,
113            min_p: 0.0,
114            logit_bias: Vec::new(),
115            repetition_penalty: 1.0,
116            repetition_penalty_window: 0,
117            presence_penalty: 0.0,
118            frequency_penalty: 0.0,
119            lz_penalty: 0.0,
120            dry_multiplier: 0.0,
121            dry_base: 1.75,
122            dry_allowed_length: 2,
123            dry_sequence_breakers: Vec::new(),
124            max_tokens,
125            stop_token_ids: Vec::new(),
126            seed: None,
127        }
128    }
129
130    pub fn is_greedy(&self) -> bool {
131        self.temperature == 0.0
132    }
133}
134
135/// 2026-09-25: Copies device logits to the host and picks tokens there.
136pub struct Sampler {
137    /// 2026-09-25: Host buffer for the BF16 logits copy, reused across calls.
138    logits_host: Vec<u8>,
139    /// 2026-09-25: FP32 expansion of one row of `logits_host`.
140    logits_f32: Vec<f32>,
141    vocab_size: usize,
142}
143
144impl Sampler {
145    pub fn new(vocab_size: usize) -> Self {
146        let logits_host = vec![0u8; vocab_size * 2];
147        let logits_f32 = vec![0.0f32; vocab_size];
148        Self {
149            logits_host,
150            logits_f32,
151            vocab_size,
152        }
153    }
154
155    /// 2026-09-25: Copy one row of BF16 logits from the device, expand it to FP32 and return it.
156    fn fetch_logits_f32(&mut self, logits_ptr: DevicePtr, gpu: &dyn GpuBackend) -> Result<&[f32]> {
157        let byte_len = self.vocab_size * 2;
158        gpu.copy_d2h(logits_ptr, &mut self.logits_host[..byte_len])?;
159        for i in 0..self.vocab_size {
160            self.logits_f32[i] = bf16_to_f32(self.logits_host[i * 2], self.logits_host[i * 2 + 1]);
161        }
162        // 2026-09-25: With `METRALE_DUMP_LOGITS_PATH=<dir>`, append this row's
163        // FP32 logits to `<dir>/logits_fetch.bin`. The variable is read once per
164        // process. The server's `RunDumps` (scheduler/dumps.rs) reads the same
165        // variable and writes `logits_seq.bin` and `logits_stok.bin` in the same
166        // directory. Open and write errors are ignored.
167        static DUMP_DIR: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
168        if let Some(dir) = DUMP_DIR.get_or_init(|| std::env::var("METRALE_DUMP_LOGITS_PATH").ok()) {
169            use std::io::Write;
170            let path = std::path::Path::new(&dir).join("logits_fetch.bin");
171            if let Ok(mut f) = std::fs::OpenOptions::new()
172                .create(true)
173                .append(true)
174                .open(&path)
175            {
176                let bytes: &[u8] = unsafe {
177                    std::slice::from_raw_parts(
178                        self.logits_f32.as_ptr() as *const u8,
179                        self.vocab_size * 4,
180                    )
181                };
182                let _ = f.write_all(bytes);
183            }
184        }
185        Ok(&self.logits_f32[..self.vocab_size])
186    }
187
188    /// 2026-09-25: Pick one token from `[vocab_size]` BF16 logits on the device.
189    ///
190    /// The pick runs on the host. Greedy params take `argmax_bf16` of the raw
191    /// row, with no penalty or bias. Otherwise the row is expanded to FP32 and
192    /// passed to `sample_with_params` with an empty history.
193    pub fn sample(
194        &mut self,
195        logits_ptr: DevicePtr,
196        params: &SamplingParams,
197        gpu: &dyn GpuBackend,
198    ) -> Result<u32> {
199        if params.is_greedy() {
200            let byte_len = self.vocab_size * 2;
201            gpu.copy_d2h(logits_ptr, &mut self.logits_host[..byte_len])?;
202            return Ok(argmax_bf16(&self.logits_host[..byte_len]));
203        }
204        let f32_logits = self.fetch_logits_f32(logits_ptr, gpu)?;
205        let f32_bytes: &[u8] = unsafe {
206            std::slice::from_raw_parts(f32_logits.as_ptr() as *const u8, f32_logits.len() * 4)
207        };
208        Ok(sample_with_params(f32_bytes, params))
209    }
210
211    /// 2026-09-25: Pick one token per row of `[batch_size, vocab_size]` BF16
212    /// logits on the device, as `sample` does for one row.
213    ///
214    /// Row `i` uses `params[i]`, or `params[0]` when `params` is shorter; it
215    /// panics when `params` is empty.
216    pub fn sample_batch(
217        &mut self,
218        logits_ptr: DevicePtr,
219        batch_size: usize,
220        params: &[&SamplingParams],
221        gpu: &dyn GpuBackend,
222    ) -> Result<Vec<u32>> {
223        let total_bytes = batch_size * self.vocab_size * 2;
224        if self.logits_host.len() < total_bytes {
225            self.logits_host.resize(total_bytes, 0);
226        }
227        gpu.copy_d2h(logits_ptr, &mut self.logits_host[..total_bytes])?;
228
229        let stride_bf16 = self.vocab_size * 2;
230        let mut tokens = Vec::with_capacity(batch_size);
231        for i in 0..batch_size {
232            let start = i * stride_bf16;
233            let end = start + stride_bf16;
234            let p = params.get(i).copied().unwrap_or(params[0]);
235            tokens.push(if p.is_greedy() {
236                argmax_bf16(&self.logits_host[start..end])
237            } else {
238                if self.logits_f32.len() < self.vocab_size {
239                    self.logits_f32.resize(self.vocab_size, 0.0);
240                }
241                for j in 0..self.vocab_size {
242                    self.logits_f32[j] = bf16_to_f32(
243                        self.logits_host[start + j * 2],
244                        self.logits_host[start + j * 2 + 1],
245                    );
246                }
247                let f32_bytes: &[u8] = unsafe {
248                    std::slice::from_raw_parts(
249                        self.logits_f32.as_ptr() as *const u8,
250                        self.vocab_size * 4,
251                    )
252                };
253                sample_with_params(f32_bytes, p)
254            });
255        }
256        Ok(tokens)
257    }
258}
259
260pub mod feed_argmax;
261#[cfg(test)]
262#[path = "feed_argmax_tests.rs"]
263mod feed_argmax_tests;
264mod penalties;
265pub use penalties::{apply_dry_penalty, apply_lz_penalty, apply_penalties_and_bias};
266mod sample_impl;
267pub use sample_impl::{sample_with_params_history, sample_with_params_seeded};
268
269/// 2026-09-25: `sample_with_params_history` with an empty history, so of the
270/// penalty stage only the logit bias applies.
271pub fn sample_with_params(data: &[u8], params: &SamplingParams) -> u32 {
272    sample_with_params_history(data, params, &[])
273}
274
275/// 2026-09-25: Index of the first maximum of `v` under strict `>`.
276///
277/// Two passes: the maximum over 8 lanes plus the remainder, then the first
278/// index equal to it. NaN never wins, and empty or all-NaN input returns 0.
279/// -0.0 and +0.0 compare equal, so the first of them wins. The server's verify
280/// path calls this (`verify_pipeline_helper/argmax.rs`), and its tests compare
281/// it with a one-loop form.
282pub fn argmax_first_wins_f32(v: &[f32]) -> u32 {
283    const LANES: usize = 8;
284    let mut acc = [f32::NEG_INFINITY; LANES];
285    let mut chunks = v.chunks_exact(LANES);
286    for c in &mut chunks {
287        for (a, &x) in acc.iter_mut().zip(c) {
288            if x > *a {
289                *a = x;
290            }
291        }
292    }
293    let mut best = f32::NEG_INFINITY;
294    for &a in acc.iter() {
295        if a > best {
296            best = a;
297        }
298    }
299    for &x in chunks.remainder() {
300        if x > best {
301            best = x;
302        }
303    }
304    v.iter()
305        .position(|&x| x == best)
306        .unwrap_or(0)
307        .try_into()
308        .unwrap_or(0)
309}
310
311/// 2026-09-25: Argmax over little-endian FP32 bytes, with the same two passes
312/// and first-index tie-break as [`argmax_first_wins_f32`], reading the byte
313/// chunks directly instead of building a `Vec<f32>`.
314pub fn argmax_f32(data: &[u8]) -> u32 {
315    debug_assert!(data.len().is_multiple_of(4));
316    let vals = || {
317        data.chunks_exact(4)
318            .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
319    };
320    const LANES: usize = 8;
321    let mut acc = [f32::NEG_INFINITY; LANES];
322    let mut it = data.chunks_exact(4 * LANES);
323    for block in &mut it {
324        for (a, c) in acc.iter_mut().zip(block.chunks_exact(4)) {
325            let x = f32::from_le_bytes([c[0], c[1], c[2], c[3]]);
326            if x > *a {
327                *a = x;
328            }
329        }
330    }
331    let mut best = f32::NEG_INFINITY;
332    for &a in acc.iter() {
333        if a > best {
334            best = a;
335        }
336    }
337    for c in it.remainder().chunks_exact(4) {
338        let x = f32::from_le_bytes([c[0], c[1], c[2], c[3]]);
339        if x > best {
340            best = x;
341        }
342    }
343    vals()
344        .position(|x| x == best)
345        .unwrap_or(0)
346        .try_into()
347        .unwrap_or(0)
348}
349
350/// 2026-09-25: Argmax over little-endian BF16 bytes; the first maximum wins, and
351/// empty input returns 0. Used by `Sampler`.
352pub fn argmax_bf16(data: &[u8]) -> u32 {
353    debug_assert!(data.len().is_multiple_of(2));
354    let n = data.len() / 2;
355    if n == 0 {
356        return 0;
357    }
358    let mut best_idx: u32 = 0;
359    let mut best_val = bf16_to_f32(data[0], data[1]);
360    for i in 1..n {
361        let val = bf16_to_f32(data[i * 2], data[i * 2 + 1]);
362        if val > best_val {
363            best_val = val;
364            best_idx = i as u32;
365        }
366    }
367    best_idx
368}
369
370/// 2026-09-25: Convert one little-endian BF16 value to f32 (exact).
371#[inline]
372fn bf16_to_f32(lo: u8, hi: u8) -> f32 {
373    let bits = (lo as u32) | ((hi as u32) << 8);
374    f32::from_bits(bits << 16)
375}
376
377#[cfg(test)]
378mod tests;
metrale_sampling/lib.rs

metrale_sampling/
lib.rs
