252 lines
8.3 KiB
Rust
252 lines
8.3 KiB
Rust
use half::bf16;
|
|
use rand::Rng;
|
|
use std::collections::HashSet;
|
|
use xserv_tensor::{DType, Device, Tensor};
|
|
|
|
#[derive(Clone)]
|
|
pub struct SamplingParams {
|
|
pub temperature: f32,
|
|
pub top_k: usize,
|
|
pub top_p: f32,
|
|
pub presence_penalty: f32,
|
|
pub repetition_penalty: f32,
|
|
}
|
|
|
|
impl Default for SamplingParams {
|
|
fn default() -> Self {
|
|
Self {
|
|
temperature: 0.0,
|
|
top_k: 0,
|
|
top_p: 1.0,
|
|
presence_penalty: 0.0,
|
|
repetition_penalty: 1.0,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Sample a token from logits with shape [seq_len, vocab_size].
|
|
/// Uses the last position's logits. Handles both F32 and BF16 dtypes.
|
|
pub fn sample(logits: &Tensor, params: &SamplingParams) -> u32 {
|
|
sample_with_history(logits, params, &[])
|
|
}
|
|
|
|
/// Sample while applying penalties to tokens already present in the request.
|
|
/// Presence penalty follows the OpenAI convention (subtract once per distinct
|
|
/// token); repetition penalty follows the HF convention.
|
|
pub fn sample_with_history(logits: &Tensor, params: &SamplingParams, history: &[u32]) -> u32 {
|
|
assert_eq!(logits.ndim(), 2);
|
|
// Greedy fast path: GPU argmax + 4-byte D2H instead of copying the whole
|
|
// [seq, vocab] logits to the host and scanning it (~201k bf16/token).
|
|
// NaN logits lose every `>` comparison in the kernel, matching the
|
|
// NaN-safe host argmax below.
|
|
if params.temperature == 0.0
|
|
&& params.presence_penalty == 0.0
|
|
&& params.repetition_penalty == 1.0
|
|
&& logits.dtype() == DType::BF16
|
|
&& matches!(logits.device(), Device::Cuda(_))
|
|
&& logits.is_contiguous()
|
|
{
|
|
let ids = xserv_kernels::argmax_bf16_to_host(logits);
|
|
return *ids.last().unwrap();
|
|
}
|
|
let vocab_size = logits.shape()[1];
|
|
let seq_len = logits.shape()[0];
|
|
let logits_cpu = logits.to_device(Device::Cpu);
|
|
|
|
// Extract last row as f32
|
|
let mut last_row: Vec<f32> = match logits.dtype() {
|
|
DType::F32 => {
|
|
let data = logits_cpu.as_slice::<f32>();
|
|
data[(seq_len - 1) * vocab_size..seq_len * vocab_size].to_vec()
|
|
}
|
|
DType::BF16 => {
|
|
let data = logits_cpu.as_slice::<bf16>();
|
|
data[(seq_len - 1) * vocab_size..seq_len * vocab_size]
|
|
.iter()
|
|
.map(|v| v.to_f32())
|
|
.collect()
|
|
}
|
|
_ => panic!("unsupported dtype for sampling: {:?}", logits.dtype()),
|
|
};
|
|
|
|
// NaN-safe: sampling path uses partial_cmp().unwrap() in top-k/top-p
|
|
// sorts and softmax; a single NaN logit would panic the engine thread.
|
|
// Replace NaN with -inf (equivalent to masking) instead.
|
|
let mut nan_seen = false;
|
|
for v in last_row.iter_mut() {
|
|
if v.is_nan() {
|
|
nan_seen = true;
|
|
*v = f32::NEG_INFINITY;
|
|
}
|
|
}
|
|
if nan_seen {
|
|
eprintln!("[sampling] WARNING: NaN logits encountered in sample()");
|
|
}
|
|
|
|
if !history.is_empty() && (params.presence_penalty != 0.0 || params.repetition_penalty != 1.0) {
|
|
let seen: HashSet<u32> = history.iter().copied().collect();
|
|
for id in seen {
|
|
let i = id as usize;
|
|
if i >= last_row.len() {
|
|
continue;
|
|
}
|
|
if params.repetition_penalty != 1.0 {
|
|
let value = last_row[i];
|
|
last_row[i] = if value > 0.0 {
|
|
value / params.repetition_penalty
|
|
} else {
|
|
value * params.repetition_penalty
|
|
};
|
|
}
|
|
last_row[i] -= params.presence_penalty;
|
|
}
|
|
}
|
|
|
|
if params.temperature == 0.0 {
|
|
return argmax(&last_row);
|
|
}
|
|
|
|
// Apply temperature
|
|
let mut logits_f32: Vec<f32> = last_row.iter().map(|v| v / params.temperature).collect();
|
|
|
|
// Top-k filtering
|
|
if params.top_k > 0 && params.top_k < vocab_size {
|
|
let mut indices: Vec<usize> = (0..vocab_size).collect();
|
|
indices.select_nth_unstable_by(params.top_k, |&a, &b| {
|
|
logits_f32[b].partial_cmp(&logits_f32[a]).unwrap()
|
|
});
|
|
// Everything after top_k should be masked
|
|
for &i in &indices[params.top_k..] {
|
|
logits_f32[i] = f32::NEG_INFINITY;
|
|
}
|
|
}
|
|
|
|
// Top-p (nucleus) filtering
|
|
if params.top_p < 1.0 {
|
|
// Sort indices by descending logit value
|
|
let mut indices: Vec<usize> = (0..vocab_size).collect();
|
|
indices.sort_unstable_by(|&a, &b| logits_f32[b].partial_cmp(&logits_f32[a]).unwrap());
|
|
|
|
// Compute softmax probabilities for the sorted order
|
|
let max_val = logits_f32[indices[0]];
|
|
let sorted_probs: Vec<f32> = indices
|
|
.iter()
|
|
.map(|&i| (logits_f32[i] - max_val).exp())
|
|
.collect();
|
|
let sum: f32 = sorted_probs.iter().sum();
|
|
let sorted_probs: Vec<f32> = sorted_probs.iter().map(|v| v / sum).collect();
|
|
|
|
// Cumulative sum, find cutoff
|
|
let mut cumsum = 0.0f32;
|
|
let mut cutoff = indices.len();
|
|
for (rank, &prob) in sorted_probs.iter().enumerate() {
|
|
cumsum += prob;
|
|
if cumsum > params.top_p {
|
|
cutoff = rank + 1; // keep at least this many
|
|
break;
|
|
}
|
|
}
|
|
|
|
// Mask everything beyond cutoff
|
|
for &i in &indices[cutoff..] {
|
|
logits_f32[i] = f32::NEG_INFINITY;
|
|
}
|
|
}
|
|
|
|
// Softmax
|
|
let max_val = logits_f32.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
|
|
let exps: Vec<f32> = logits_f32.iter().map(|v| (v - max_val).exp()).collect();
|
|
let sum: f32 = exps.iter().sum();
|
|
let probs: Vec<f32> = exps.iter().map(|v| v / sum).collect();
|
|
|
|
// Weighted random sampling
|
|
let mut rng = rand::thread_rng();
|
|
let r: f32 = rng.r#gen();
|
|
let mut cumsum = 0.0f32;
|
|
for (i, &p) in probs.iter().enumerate() {
|
|
cumsum += p;
|
|
if cumsum > r {
|
|
return i as u32;
|
|
}
|
|
}
|
|
|
|
// Fallback (rounding edge case)
|
|
(vocab_size - 1) as u32
|
|
}
|
|
|
|
/// Greedy argmax with a repetition penalty applied to `recent` token ids
|
|
/// (HF-style: divide positive logits by `penalty`, multiply negative by it).
|
|
/// `penalty <= 1.0` is a no-op. Mitigates greedy repetition loops on reasoning
|
|
/// models without changing the forward pass. NaN-safe.
|
|
pub fn sample_greedy_penalized(logits: &Tensor, recent: &[u32], penalty: f32) -> u32 {
|
|
assert_eq!(logits.ndim(), 2);
|
|
let vocab_size = logits.shape()[1];
|
|
let seq_len = logits.shape()[0];
|
|
let logits_cpu = logits.to_device(Device::Cpu);
|
|
let mut last_row: Vec<f32> = match logits.dtype() {
|
|
DType::F32 => {
|
|
logits_cpu.as_slice::<f32>()[(seq_len - 1) * vocab_size..seq_len * vocab_size].to_vec()
|
|
}
|
|
DType::BF16 => logits_cpu.as_slice::<bf16>()
|
|
[(seq_len - 1) * vocab_size..seq_len * vocab_size]
|
|
.iter()
|
|
.map(|v| v.to_f32())
|
|
.collect(),
|
|
_ => panic!("unsupported dtype for sampling: {:?}", logits.dtype()),
|
|
};
|
|
if penalty > 1.0 {
|
|
for &id in recent {
|
|
let i = id as usize;
|
|
if i < last_row.len() {
|
|
let v = last_row[i];
|
|
last_row[i] = if v > 0.0 { v / penalty } else { v * penalty };
|
|
}
|
|
}
|
|
}
|
|
argmax(&last_row)
|
|
}
|
|
|
|
fn argmax(data: &[f32]) -> u32 {
|
|
// NaN-safe: a single NaN logit must not crash the engine thread (a
|
|
// partial_cmp().unwrap() panics on NaN). Skip NaNs; warn once if seen.
|
|
let mut best_i = 0usize;
|
|
let mut best = f32::NEG_INFINITY;
|
|
let mut nan_seen = false;
|
|
for (i, &v) in data.iter().enumerate() {
|
|
if v.is_nan() {
|
|
nan_seen = true;
|
|
continue;
|
|
}
|
|
if v > best {
|
|
best = v;
|
|
best_i = i;
|
|
}
|
|
}
|
|
if nan_seen {
|
|
eprintln!("[sampling] WARNING: NaN logits encountered in argmax");
|
|
}
|
|
best_i as u32
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn greedy_sampling_applies_history_penalties() {
|
|
let logits = Tensor::from_slice(&[10.0f32, 9.0, 0.0], &[1, 3]);
|
|
|
|
let presence = SamplingParams {
|
|
presence_penalty: 2.0,
|
|
..SamplingParams::default()
|
|
};
|
|
assert_eq!(sample_with_history(&logits, &presence, &[0]), 1);
|
|
|
|
let repetition = SamplingParams {
|
|
repetition_penalty: 2.0,
|
|
..SamplingParams::default()
|
|
};
|
|
assert_eq!(sample_with_history(&logits, &repetition, &[0]), 1);
|
|
}
|
|
}
|