xserv/crates/xserv-model/src/sampling.rs

use half::bf16;
use rand::Rng;
use xserv_tensor::{DType, Device, Tensor};

#[derive(Clone)]
pub struct SamplingParams {
    pub temperature: f32,
    pub top_k: usize,
    pub top_p: f32,
}

impl Default for SamplingParams {
    fn default() -> Self {
        Self {
            temperature: 0.0,
            top_k: 0,
            top_p: 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 {
    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
        && 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()),
    };

    // Greedy
    if params.temperature == 0.0 {
        return argmax(&last_row);
    }

    // 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()");
    }

    // 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
}