Gemma 4 compatible inference model. More...

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

Public Types
using	GemmaSlidingKvPolicy = Quant::KvCache::SlidingWindowKvCache
	KV policy for Gemma's LOCAL (sliding) layers.
using	ModelBase = LanguageModel<TDeviceType, TPrecision>
using	MR = typename DeviceTypeTraits<TDeviceType>::memory_resource
using	StagingMR = CpuMemoryResource
using	TensorType = Tensor<TPrecision, MR>
using	TokenIndexType = Tensor<dtype_t::INT32, MR>
Public Types inherited from Mila::Dnn::LanguageModel< TDeviceType, TPrecision >
using	Base = Model<TDeviceType, TPrecision>
Public Types inherited from Mila::Dnn::Model< TDeviceType, TPrecision >
using	NetworkType = Network<TDeviceType, TPrecision>

Public Member Functions
	GemmaModel (const GemmaModel &)=delete
	GemmaModel (GemmaModel &&)=default
dim_t	contextLength () const noexcept
	Deployment context length: the KV-cache depth the network was built with.
std::string	fingerprintPrefill (const std::vector< int32_t > &token_ids)
	Logits fingerprint for a fixed prompt, for comparing two loads of one model.
const GemmaModelConfig &	getModelConfig () const noexcept
	Deployment configuration (context length, weight-quant, kv-compression) this model was loaded with.
const GemmaConfig &	getNetworkConfig () const noexcept
	Architecture/network configuration (read from the checkpoint metadata).
GemmaModel &	operator= (const GemmaModel &)=delete
GemmaModel &	operator= (GemmaModel &&)=default
void	profilePrefill (const std::vector< int32_t > &token_ids)
std::string	toString () const override
	Human-readable summary of this model's configuration.
Public Member Functions inherited from Mila::Dnn::LanguageModel< TDeviceType, TPrecision >
	LanguageModel (const LanguageModel &)=delete
	LanguageModel (LanguageModel &&)=default
GenerateStatus	generate (std::span< const int32_t > prompt_tokens, const std::function< void(int32_t)> &on_token, const GenerateParams &params={}, std::stop_token stop={})
	Generate tokens from a prompt, streaming each through on_token.
LanguageModel &	operator= (const LanguageModel &)=delete
LanguageModel &	operator= (LanguageModel &&)=default
void	savePretrained (const std::filesystem::path &path) const
	Write this model's live weights as a safetensors artifact.
void	seedSampler (uint64_t seed)
	Seed the sampler's RNG for reproducible generation.
Public Member Functions inherited from Mila::Dnn::Model< TDeviceType, TPrecision >
	Model (const Model &)=delete
	Model (Model &&)=default
DeviceId	getDeviceId () const noexcept
	The device this model runs on.
MemoryStats	getMemoryStats () const
	Current memory allocation breakdown for this model.
RuntimeMode	getRuntimeMode () const noexcept
	The runtime mode this model was constructed for.
std::size_t	getScratchHighWaterBytes () const
	Context-owned scratch device memory, in bytes, at its high-water mark.
bool	isEval () const noexcept
	True if this model is currently in eval sub-state.
bool	isInferenceMode () const noexcept
	True if this model was constructed for inference.
bool	isTrainingMode () const noexcept
	True if this model was constructed for training.
Model &	operator= (const Model &)=delete
Model &	operator= (Model &&)=default
void	setEval (bool eval)
	Toggle eval sub-state for this model.
void	train ()
	Run the training loop for this model.

Static Public Member Functions
static std::unique_ptr< GemmaModel< TDeviceType, TPrecision > >	fromPretrained (const std::filesystem::path &path, const GemmaModelConfig &model_config, DeviceId device_id=DeviceId{ TDeviceType, 0 })
	Load from a Mila-converted Gemma 4 pretrained artifact.
static MemoryStats	getRequiredMemory (const std::filesystem::path &path, const GemmaModelConfig &model_config, DeviceId device_id=DeviceId{ TDeviceType, 0 })
	What loading this checkpoint at this context length would cost in VRAM.

Protected Member Functions
dim_t	maxSequenceLength () const noexcept override
GenerateStatus	onGenerating (std::span< const int32_t > prompt_tokens, const std::function< void(int32_t)> &on_token, const GenerateParams &params, std::stop_token stop) override
	Prefill + decode implementation hook.
void	onTraining () override
	Training loop hook – derived class owns the implementation.
dim_t	vocabSize () const noexcept override
Protected Member Functions inherited from Mila::Dnn::LanguageModel< TDeviceType, TPrecision >
	LanguageModel (std::unique_ptr< LanguageNetwork< TDeviceType, TPrecision > > network, RuntimeMode runtime_mode, Serialization::PretrainedMetadata source_metadata={}, WeightQuantization weight_quantization=WeightQuantization::None)
int32_t	awaitSampledToken ()
	Block until the last enqueueSampleNext()'s token id is host-visible.
void	enqueueSampleNext (const TensorType &logits, TokenTensor &token_out, const SamplingParams &params)
	Enqueue a sampling step without waiting for the host readback.
const LanguageNetwork< TDeviceType, TPrecision > &	getLanguageNetwork () const noexcept
LanguageNetwork< TDeviceType, TPrecision > &	getLanguageNetwork () noexcept
int32_t	sampleNext (const TensorType &logits, TokenTensor &token_out, const SamplingParams &params)
	Sample the next token from a logits row on the device.
Protected Member Functions inherited from Mila::Dnn::Model< TDeviceType, TPrecision >
	Model (std::unique_ptr< NetworkType > network, RuntimeMode runtime_mode)
	Construct with a fully built network and runtime mode.

Additional Inherited Members
Protected Types inherited from Mila::Dnn::LanguageModel< TDeviceType, TPrecision >
using	MR = typename DeviceTypeTraits<TDeviceType>::memory_resource
using	TensorType = Tensor<TPrecision, MR>
using	TokenTensor = Tensor<TensorDataType::INT32, MR>
Protected Attributes inherited from Mila::Dnn::LanguageModel< TDeviceType, TPrecision >
Serialization::PretrainedMetadata	source_metadata_
	The loaded artifact's metadata, carried so savePretrained can write it back verbatim.
WeightQuantization	weight_quantization_ { WeightQuantization::None }
	What the live weights are, not what the source file was.
Protected Attributes inherited from Mila::Dnn::Model< TDeviceType, TPrecision >
std::unique_ptr< NetworkType >	network_
	The owned Network instance.

Detailed Description

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

Gemma 4 compatible inference model.

Owns a loaded, built GemmaTransformer and drives the prefill + KV-cache decode two-phase generation loop. Construction is only possible via fromPretrained(); the network is always built, weights-loaded, and in inference mode when generation runs.

Thread safety: not thread-safe; external synchronization required if shared.

Member Typedef Documentation

◆ GemmaSlidingKvPolicy

template<DeviceType TDeviceType, TensorDataType TPrecision>

using Mila::Dnn::GemmaModel< TDeviceType, TPrecision >::GemmaSlidingKvPolicy = Quant::KvCache::SlidingWindowKvCache

KV policy for Gemma's LOCAL (sliding) layers.

Bounded sliding-window ring (SlidingWindowKvCache.md Phase 3): their cache is sized to the window working set instead of the full context. Strictly a memory optimization – tokens are identical to the full cache. GLOBAL (full-attention) layers are always NoKvCompression, hardwired in GemmaTransformer. Flip this alias to NoKvCompression to A/B the footprint against the full-context sliding cache.

Class scope rather than per-function so the load and footprint paths cannot be pointed at different policies – that would make a model report a figure for a cache it does not build.

Member Function Documentation

◆ fingerprintPrefill()

template<DeviceType TDeviceType, TensorDataType TPrecision>

std::string Mila::Dnn::GemmaModel< TDeviceType, TPrecision >::fingerprintPrefill ( const std::vector< int32_t > & token_ids )

inline

Logits fingerprint for a fixed prompt, for comparing two loads of one model.

Diagnostic. Parameters and config can be proven byte-identical between a quantize-on-load and a pre-quantized load while the models still behave differently, and nothing upstream of inference can see that. This runs one prefill and reports what the model actually computed.

Raw token ids rather than text so no tokenizer is involved and two runs are comparable by construction.

Returns: Digest of the last-position logits, the argmax token, and its value.

◆ fromPretrained()

template<DeviceType TDeviceType, TensorDataType TPrecision>

std::unique_ptr< GemmaModel< TDeviceType, TPrecision > > Mila::Dnn::GemmaModel< TDeviceType, TPrecision >::fromPretrained	(	const std::filesystem::path &	path,
		const GemmaModelConfig &	model_config,
		DeviceId	device_id = DeviceId{ TDeviceType, 0 } )

inlinestatic

Load from a Mila-converted Gemma 4 pretrained artifact.

The model_config carries the deployment decisions (context length, weight quantization, KV-cache compression); every architectural parameter is read from the checkpoint metadata.

Parameters

path	Path to the pretrained Gemma model artifact.
model_config	Deployment configuration for this load.
device_id	Target device; must match TDeviceType.

Returns: Inference-ready GemmaModel.

Exceptions

std::invalid_argument	on device type mismatch or zero context length.
std::runtime_error	on load failure or unsupported quantization.

◆ getRequiredMemory()

template<DeviceType TDeviceType, TensorDataType TPrecision>

MemoryStats Mila::Dnn::GemmaModel< TDeviceType, TPrecision >::getRequiredMemory	(	const std::filesystem::path &	path,
		const GemmaModelConfig &	model_config,
		DeviceId	device_id = DeviceId{ TDeviceType, 0 } )

inlinestatic

What loading this checkpoint at this context length would cost in VRAM.

Reads the artifact header for geometry and constructs the graph, then reports what build() would allocate – without building it, without reading a weight, and therefore without needing the device to have room. Answers before a multi-gigabyte download and for hardware the caller does not own.

Returns measurements only. Whether a given headroom is too tight is a deployment policy and belongs to the adaptor, not here; on Windows in particular WDDM oversubscribes rather than failing, so "fits" is not a property the runtime can decide. See Specifications/MemoryFootprint.md.

Exceptions

std::invalid_argument	on device type mismatch or zero context length.
std::runtime_error	on an unreadable artifact or unsupported quantization.

◆ maxSequenceLength()

template<DeviceType TDeviceType, TensorDataType TPrecision>

dim_t Mila::Dnn::GemmaModel< TDeviceType, TPrecision >::maxSequenceLength ( ) const

inlineoverrideprotectedvirtualnoexcept

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

◆ onGenerating()

template<DeviceType TDeviceType, TensorDataType TPrecision>

GenerateStatus Mila::Dnn::GemmaModel< TDeviceType, TPrecision >::onGenerating	(	std::span< const int32_t >	prompt_tokens,
		const std::function< void(int32_t)> &	on_token,
		const GenerateParams &	params,
		std::stop_token	stop )

inlineoverrideprotectedvirtual

Prefill + decode implementation hook.

Derived classes own the full autoregressive generation loop. on_token must be called for every generated token except EOS. stop.stop_requested() must be checked on each decode step and generation must abort early when signalled, returning the GenerateStatus that reflects why the loop stopped.

Parameters

prompt_tokens	Input token ids.
on_token	Per-token callback.
params	Per-call generation parameters (loop bound + sampling).
stop	Stop token for cooperative cancellation.

Returns: Why generation stopped.

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

◆ onTraining()

template<DeviceType TDeviceType, TensorDataType TPrecision>

void Mila::Dnn::GemmaModel< TDeviceType, TPrecision >::onTraining ( )

inlineoverrideprotectedvirtual

Training loop hook – derived class owns the implementation.

Called by train() after precondition enforcement. The derived class has total control over data loading, optimizer construction, loss computation, backward pass, checkpointing, and sampling.

Pure virtual – a model declaring RuntimeMode::Training must provide a training loop.

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

◆ toString()

template<DeviceType TDeviceType, TensorDataType TPrecision>

std::string Mila::Dnn::GemmaModel< TDeviceType, TPrecision >::toString ( ) const

inlineoverridevirtual

Human-readable summary of this model's configuration.

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

◆ vocabSize()

template<DeviceType TDeviceType, TensorDataType TPrecision>

dim_t Mila::Dnn::GemmaModel< TDeviceType, TPrecision >::vocabSize ( ) const

inlineoverrideprotectedvirtualnoexcept

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

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

Mila/Src/Dnn/Models/GemmaModel.ixx

Public Types

Public Member Functions

Static Public Member Functions

Protected Member Functions

Additional Inherited Members

Detailed Description

Member Typedef Documentation

◆ GemmaSlidingKvPolicy

Member Function Documentation

◆ fingerprintPrefill()

◆ fromPretrained()

◆ getRequiredMemory()

◆ maxSequenceLength()

◆ onGenerating()

◆ onTraining()

◆ toString()

◆ vocabSize()