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

The standard token sampler: temperature / top-k / top-p multinomial. More...

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

Public Types
using	Base = Sampler<TDeviceType, TPrecision>
using	SamplingOpType = typename OperationTraits<OperationType::SamplingOp, TDeviceType, TPrecision>::type
using	TokenTensor = typename Base::TokenTensor
Public Types inherited from Mila::Dnn::Sampler< TDeviceType, TPrecision >
using	MR = typename DeviceTypeTraits<TDeviceType>::memory_resource
using	TokenTensor = Tensor<TensorDataType::INT32, MR>

Public Member Functions
	TokenSampler (IExecutionContext *context, const SamplingConfig &config)
int32_t	awaitToken ()
	Block until the last enqueueSample()'s token id is host-visible and return it.
void	enqueueSample (const ITensor &logits, TokenTensor &token_out, const SamplingParams &params)
	Enqueue one sampling step without waiting for the token readback.
void	reseed (uint64_t seed)
	Reseed the host RNG for reproducible sampling.
int32_t	sample (const ITensor &logits, TokenTensor &token_out, const SamplingParams &params) override
	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::TokenSampler< TDeviceType, TPrecision >

The standard token sampler: temperature / top-k / top-p multinomial.

Dispatches to the device SamplingOp resolved by OperationTraits (CudaSamplingOp / CpuSamplingOp). Shares the model's ExecutionContext so the op runs on the decode stream and writes the token in place. The host-drawn uniform is generated here so the op stays pure and deterministic (Phase A: greedy only).

Template Parameters

TDeviceType	Device the logits reside on.
TPrecision	Logits precision (FP32 or BF16).

Member Function Documentation

enqueueSample()

template<DeviceType TDeviceType, TensorDataType TPrecision>

void Mila::Dnn::TokenSampler< TDeviceType, TPrecision >::enqueueSample ( const ITensor & logits, TokenTensor & token_out, const SamplingParams & params )

inline

Enqueue one sampling step without waiting for the token readback.

Decode-ahead half of the pipelined generation loop: the op samples on the model's stream (ordered after the forward pass that produced logits) and writes the token into token_out in place, so the next decode step can be enqueued before the host knows the token id. The host uniform is drawn here at enqueue time – one draw per sampled token, same RNG sequence as sample(). awaitToken() must be called before the next enqueueSample().

reseed()

template<DeviceType TDeviceType, TensorDataType TPrecision>

void Mila::Dnn::TokenSampler< TDeviceType, TPrecision >::reseed ( uint64_t seed )

inline

Reseed the host RNG for reproducible sampling.

Called once per generation run (not per token) when the caller supplies a seed.

sample()

template<DeviceType TDeviceType, TensorDataType TPrecision>

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

inlineoverridevirtual

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).

Implements Mila::Dnn::Sampler< TDeviceType, TPrecision >.

---

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

• Mila/Src/Dnn/Samplers/TokenSampler.ixx