Struct Sampler 

pub struct Sampler { /* private fields */ }

Sampler that picks tokens from logits.

Implementations

impl Sampler

pub fn new(vocab_size: usize) -> Self

pub fn sample( &mut self, logits_ptr: DevicePtr, params: &SamplingParams, gpu: &dyn GpuBackend, ) -> Result<u32>

Sample a token from logits on the GPU.

logits_ptr points to [vocab_size] BF16 values on device. Reads BF16, expands to FP32, then samples with full precision.

pub fn sample_batch( &mut self, logits_ptr: DevicePtr, batch_size: usize, params: &[&SamplingParams], gpu: &dyn GpuBackend, ) -> Result<Vec<u32>>

Sample a batch of tokens (one per sequence in the batch).

logits_ptr points to [batch_size, vocab_size] BF16 values.