Mila::Dnn::Sampler< TDeviceType, TPrecision > Class Template Reference

Abstract base for token samplers. More...

Inheritance diagram for Mila::Dnn::Sampler< TDeviceType, TPrecision >:

Public Types
using	MR = typename DeviceTypeTraits<TDeviceType>::memory_resource
using	TokenTensor = Tensor<TensorDataType::INT32, MR>

Public Member Functions
virtual int32_t	sample (const ITensor &logits, TokenTensor &token_out, const SamplingParams &params)=0
	Select the next token from the final row of a logits tensor.

Detailed Description

template<DeviceType TDeviceType, TensorDataType TPrecision>
requires PrecisionSupportedOnDevice<TPrecision, TDeviceType>
class Mila::Dnn::Sampler< TDeviceType, TPrecision >

Abstract base for token samplers.

A Sampler maps a logits tensor to a single next-token id. The device-dispatched TokenSampler is the one concrete sampler; the base is retained as the seam for a future stateful strategy (e.g. Mirostat). The model owns the sampler and shares its ExecutionContext; the sampler is invoked by the generation loop, never by the operation graph.

Template Parameters

TDeviceType	Device the logits reside on (Cpu or Cuda).
TPrecision	Logits tensor precision (e.g. FP32, BF16).

Member Function Documentation

sample()

template<DeviceType TDeviceType, TensorDataType TPrecision>

virtual int32_t Mila::Dnn::Sampler< TDeviceType, TPrecision >::sample ( const ITensor & logits, TokenTensor & token_out, const SamplingParams & params )

pure virtual

Select the next token from the final row of a logits tensor.

Samples from the last position of logits (last dim == vocab_size), writing the chosen int32 token into token_out on the device stream and returning the same value on the host after the readback. token_out is the caller-owned [1, 1] INT32 device tensor (the model's decode-input buffer), so the sampled token is already in place for the next decode step.

Parameters

logits	Device logits tensor; the last vocab_size elements are the active distribution. Read in place, never copied to host.
token_out	Device INT32 [1, 1] tensor; written in place.
params	Per-call sampling parameters (temperature / top_k / top_p / seed).

Returns

The sampled token id (host value, after device readback).

Implemented in Mila::Dnn::TokenSampler< TDeviceType, TPrecision >.

---

The documentation for this class was generated from the following file:

• Mila/Src/Dnn/Core/SamplerBase.ixx