The type erasure boundary between a language model and its transformer. More...

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

Public Types
using	MR = typename DeviceTypeTraits<TDeviceType>::memory_resource
using	NetworkBase = Network<TDeviceType, TPrecision>
using	TensorType = Tensor<TPrecision, MR>
using	TokenIndexType = Tensor<TensorDataType::INT32, MR>
Public Types inherited from Mila::Dnn::Network< TDeviceType, TPrecision >
using	ComponentPtr = typename CompositeBase::ComponentPtr
using	CompositeBase = CompositeComponent<TDeviceType, TPrecision>
using	MR = typename DeviceTypeTraits<TDeviceType>::memory_resource
Public Types inherited from Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >
using	ComponentBase = Component<TDeviceType, TPrecision>
using	ComponentPtr = std::shared_ptr<Component<TDeviceType, TPrecision>>

Public Member Functions
	LanguageModelNetwork (const std::string &name)
virtual TokenIndexType &	backward (const TokenIndexType &input, const TensorType &output_grad)=0
	Full backward pass (training).
virtual TensorType &	decode (const TokenIndexType &input, dim_t position)=0
	Inference decode – single-token autoregressive step.
virtual TensorType &	forward (const TokenIndexType &input)=0
	Full-sequence forward pass – the training path, not the generation path.
virtual TensorType &	prefill (const TokenIndexType &input)=0
	Inference prefill – process full prompt and populate the KV cache.
virtual TensorType &	prefillFrom (const TokenIndexType &input, dim_t start_offset)
	Chunked prefill starting at an absolute position (prompt-prefix reuse).
virtual bool	rewindKvCache (dim_t position)
virtual SequenceLogLikelihood	scoreTokens (const TokenIndexType &input)
	Rewind the KV caches to position for prompt-prefix reuse (PromptCaching.md).
Public Member Functions inherited from Mila::Dnn::Network< TDeviceType, TPrecision >
	Network (const std::string &name)
	Construct network (context managed by derived class).
template<typename TOptimizer, typename TConfig>
std::shared_ptr< TOptimizer >	createOptimizer (const TConfig &config)
	Create and configure an optimizer for this network's parameters.
DeviceId	getDeviceId () const noexcept override
	Get the compute device for this composite.
const ComponentType	getType () const override
	Get the component type identifier.
void	load (ModelArchive &archive, SerializationMode mode)
	Restore this network's parameters from an archive.
void	save (ModelArchive &archive, SerializationMode mode) const
	Save network to archive.
void	synchronize () override
	Synchronize all child components.
std::string	toString () const override
	Generate a human-readable description.
Public Member Functions inherited from Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >
	CompositeComponent (CompositeComponent &&) noexcept=default
	CompositeComponent (const CompositeComponent &)=delete
	CompositeComponent (const std::string &name)
	Construct composite component with name.
CompositeComponent &	addComponent (ComponentPtr component)
	Add a pre-constructed child component (chainable).
size_t	childCount () const noexcept
	Get the number of direct children.
void	clearComponents ()
	Clear all child components.
std::vector< std::string >	componentPaths () const
	Every component path in this subtree, in traversal order.
ComponentPtr	findComponent (const std::string &path) const
	Resolve a dot-separated component path within this composite.
ComponentPtr	getComponent (const std::string &name) const
	Retrieve a direct child component by name.
const std::vector< ComponentPtr > &	getComponents () const
	Get all child components in insertion order.
std::vector< ITensor * >	getGradients () const override
	Get all parameter gradients from all children.
std::vector< ITensor * >	getParameters () const override
	Get all parameters from all children.
bool	hasChildren () const noexcept
	Check if this composite has any children.
bool	hasComponent (const std::string &name) const
	Check if a named child component exists.
size_t	observe (std::string_view pattern, Compute::ComputePassMask passes, Compute::ActivationObserver sink)
	Attach sink to every component in this subtree whose path matches.
CompositeComponent &	operator= (CompositeComponent &&) noexcept=default
CompositeComponent &	operator= (const CompositeComponent &)=delete
dim_t	parameterCount () const override
	Count parameters across all children.
bool	removeComponent (const std::string &name)
	Remove a child component by name.
void	saveFlatTensors (Serialization::SafeTensorsWriter &writer, const std::string &prefix, Serialization::TensorSavePass pass) const override
	Recurse into children, extending the flat dotted prefix.
void	stopObserving ()
	Detach every observer in this subtree and clear the shared sink.
ComponentPtr	tryFindComponent (const std::string &path) const
	Try to resolve a dot-separated component path within this composite.
Public Member Functions inherited from Mila::Dnn::Component< TDeviceType, TPrecision >
	Component (const std::string &name)
	Construct component with required name identifier.
virtual void	build (const BuildContext &context) final
	Build the component with the provided BuildContext (canonical overload).
IExecutionContext *	getExecutionContext () const
	Get the shared execution context.
virtual MemoryStats	getMemoryStats () const =0
	Return the current memory allocation breakdown for this component.
const std::string	getName () const
	Get the component's name identifier.
virtual std::vector< ObservableStage >	getObservableStages () const
	Stages this component publishes, and the passes each is published on.
Compute::ComputePassMask	getObservedPasses () const noexcept
virtual std::vector< const ITensor * >	getOutputs () const
	Return non-owning pointers to the tensors this component produces.
virtual std::vector< std::string >	getParameterNames () const
	List all available parameter names for this component.
virtual MemoryStats	getRequiredMemory (const BuildContext &context) const
	Report what build( context ) would allocate, without allocating it.
RuntimeMode	getRuntimeMode () const noexcept
TrainingMode	getTrainingMode () const noexcept
	The current runtime behavioral mode of this Component.
virtual bool	isBuilt () const final
	Returns true if build() has completed successfully.
bool	isInferenceMode () const noexcept
bool	isTrainingMode () const noexcept
virtual void	loadParameter (const std::string &, const Serialization::ITensorBlob &)
	Load a parameter from serialized tensor data.
void	setObservedPasses (Compute::ComputePassMask passes) noexcept
	Select which passes this component publishes on.
void	setTrainingMode (TrainingMode mode)
	Set the runtime behavioral mode for this Component.
virtual void	zeroGradients ()
	Clear all model-owned gradients for this component.

Additional Inherited Members
Static Public Member Functions inherited from Mila::Dnn::Component< TDeviceType, TPrecision >
static constexpr DeviceType	getDeviceType ()
	Compile-time device type for this component instance.
static constexpr TensorDataType	getPrecision () noexcept
	Compile-time tensor precision for this component instance.
Protected Types inherited from Mila::Dnn::Component< TDeviceType, TPrecision >
using	HostStagingMemoryResource = typename DeviceTypeTraits<TDeviceType>::host_staging_memory_resource
	Host memory a device-resident parameter stages through.
Protected Member Functions inherited from Mila::Dnn::Network< TDeviceType, TPrecision >
virtual void	load_ (ModelArchive &archive, SerializationMode mode) override
	Hook for concrete classes to validate type-specific state on load.
virtual void	save_ (ModelArchive &archive, SerializationMode mode) const override=0
	Hook for concrete classes to save type-specific state.
void	verifyArchitectureCompatibility (const WeightsMetadata &metadata)
	Verify that imported model is compatible with network architecture.
Protected Member Functions inherited from Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >
template<typename TComponent>
std::shared_ptr< TComponent >	getComponentAs (const std::string &name) const
	Retrieve a typed child component by name.
void	onExecutionContextSet () override
	Hook invoked after ExecutionContext is set.
void	onTrainingModeChanging (TrainingMode training_mode) override
	Hook invoked when training mode is about to change.
virtual void	optimize ()
	Virtual hook for graph optimization after construction.
void	requireSerializableParameters () const override
	No-op override: a composite names no parameters of its own.
Protected Member Functions inherited from Mila::Dnn::Component< TDeviceType, TPrecision >
bool	hasExecutionContext () const noexcept
	Check if execution context has been set.
template<TensorDataType TParameterPrecision, typename TMemoryResource>
void	loadParameterFromBlob (const std::string &param_name, const Serialization::ITensorBlob &blob, Tensor< TParameterPrecision, TMemoryResource > &target, const shape_t &expected_shape)
	Load a tensor blob into a parameter tensor with validation.
virtual void	onBuilding (const BuildContext &)
	Hook invoked by build() to allocate component buffers.
void	publish (Compute::ComputePass pass, std::string_view stage, const ITensor &value) const
	Publish one activation to the observer, if this pass is being observed.
template<TensorDataType TParameterPrecision, typename TMemoryResource>
void	saveParameterToArchive (ModelArchive &archive, const std::string &parameter_name, const Tensor< TParameterPrecision, TMemoryResource > &parameter) const
	Write one parameter tensor into the archive under "tensors/<name>".
template<TensorDataType TParameterPrecision, typename TMemoryResource>
void	saveParameterToWriter (Serialization::SafeTensorsWriter &writer, const std::string &flat_name, const Tensor< TParameterPrecision, TMemoryResource > &parameter, Serialization::TensorSavePass pass) const
	Drive one parameter through one pass of a flat safetensors save.
void	setExecutionContext (IExecutionContext *context)
	Set the execution context for this component.
Protected Attributes inherited from Mila::Dnn::Component< TDeviceType, TPrecision >
BuildContext	build_context_ { shape_t{ 1 }, RuntimeMode::Training }
	The BuildContext stored at build time.

Detailed Description

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

The type erasure boundary between a language model and its transformer.

A concrete transformer is templated on its weight-quantization and KV-cache policies, and those parameters have no business reaching the model layer – GemmaModel should not be a different type because its weights are FP4. LanguageModelNetwork is where they stop. LanguageModel holds one of these and drives generation through it, so the interface below is the whole vocabulary a model needs from a transformer: prefill/decode, the prefix-reuse pair (prefillFrom + rewindKvCache), and nothing else.

The virtual boundary is deliberately coarse. One dispatch per layer per token step is negligible against the per-layer GEMMs, so the interface is drawn at whole passes rather than at anything finer.

Implemented by all four transformer families:

Network<TDeviceType, TPrecision> +- LanguageModelNetwork<TDeviceType, TPrecision> +- GptTransformer<...> training + inference +- LlamaTransformer<...> training + inference +- GemmaTransformer<...> inference only +- QwenTransformer<...> inference only

Not every member is implemented by every family; each states its own support.

Member Function Documentation

◆ backward()

template<DeviceType TDeviceType, TensorDataType TPrecision>

virtual TokenIndexType & Mila::Dnn::LanguageModelNetwork< TDeviceType, TPrecision >::backward	(	const TokenIndexType &	input,
		const TensorType &	output_grad )

pure virtual

Full backward pass (training).

Requires training mode and a prior forward.

Implemented by GptTransformer and LlamaTransformer; the inference-only families throw.

The return type is the INT32 token-index tensor and carries no usable gradient: the input is discrete, so the token embedding has no input gradient to produce. Callers should discard it.

Parameters

input	Token indices [B, T].
output_grad	Gradient of the loss w.r.t. logits.

◆ decode()

template<DeviceType TDeviceType, TensorDataType TPrecision>

virtual TensorType & Mila::Dnn::LanguageModelNetwork< TDeviceType, TPrecision >::decode	(	const TokenIndexType &	input,
		dim_t	position )

pure virtual

Inference decode – single-token autoregressive step.

Parameters

input	Single token index [B, 1].
position	Current sequence position (0-based).

Returns: Logits [B, 1, vocab_size].

◆ forward()

template<DeviceType TDeviceType, TensorDataType TPrecision>

virtual TensorType & Mila::Dnn::LanguageModelNetwork< TDeviceType, TPrecision >::forward ( const TokenIndexType & input )

pure virtual

Full-sequence forward pass – the training path, not the generation path.

Retains per-layer activations for backward. Implemented by GptTransformer and LlamaTransformer; the inference-only families (GemmaTransformer, QwenTransformer) throw. Generation does not use this – see prefill/decode below.

Parameters

input Token indices [B, T].

Returns: Logits [B, T, vocab_size].

◆ prefill()

template<DeviceType TDeviceType, TensorDataType TPrecision>

virtual TensorType & Mila::Dnn::LanguageModelNetwork< TDeviceType, TPrecision >::prefill ( const TokenIndexType & input )

pure virtual

Inference prefill – process full prompt and populate the KV cache.

Equivalent to prefillFrom( input, 0 ). Starting from position 0 also discards any state carried from a previous sequence, so a fresh prompt needs no explicit reset: the KV cache, the DeltaNet recurrent state and the convolution windows all treat offset 0 as a new sequence.

Parameters

input Full prompt token indices [B, T].

Returns: Logits for the last token position.

◆ prefillFrom()

template<DeviceType TDeviceType, TensorDataType TPrecision>

virtual TensorType & Mila::Dnn::LanguageModelNetwork< TDeviceType, TPrecision >::prefillFrom	(	const TokenIndexType &	input,
		dim_t	start_offset )

inlinevirtual

Chunked prefill starting at an absolute position (prompt-prefix reuse).

Parameters

input	The FULL prompt token indices [B, T] – not a pre-sliced tail; token index and absolute position coincide.
start_offset	First position to prefill; [0, start_offset) must already be resident in the KV caches (see rewindKvCache).

Returns: Logits for the last token position.

Implemented by GemmaTransformer and QwenTransformer, which override both this and rewindKvCache; on any other network it throws. Because rewindKvCache defaults to false and a caller reaches this only after a successful rewind, the throw is not reachable through the intended sequence.

◆ rewindKvCache()

template<DeviceType TDeviceType, TensorDataType TPrecision>

virtual bool Mila::Dnn::LanguageModelNetwork< TDeviceType, TPrecision >::rewindKvCache ( dim_t position )

inlinevirtual

Reimplemented in Mila::Dnn::GemmaTransformer< TDeviceType, TPrecision, TWeightQuantization, TKvCachePolicy, kDelegatedFeedForward, kMixtureOfExperts >, and Mila::Dnn::QwenTransformer< TDeviceType, TPrecision, TWeightPlan, TKvCachePolicy >.

◆ scoreTokens()

template<DeviceType TDeviceType, TensorDataType TPrecision>

virtual SequenceLogLikelihood Mila::Dnn::LanguageModelNetwork< TDeviceType, TPrecision >::scoreTokens ( const TokenIndexType & input )

inlinevirtual

Rewind the KV caches to position for prompt-prefix reuse (PromptCaching.md).

Positions [0, position) stay valid; device contents are untouched.

Returns: true when every layer accepted the rewind. Default: false (no reuse capability); a full prefill positionally overwrites regardless, so a refused or partial rewind never needs cleanup.

Teacher-forced scoring: how well the model predicted a sequence it is given.

Runs the sequence through the prefill path and, at every position, reads the probability the model assigned to the token that actually came next. Nothing is sampled and nothing is generated, so the result is a property of the model and the text alone – which is what makes it usable as a quality measure between two quantizations of the same weights.

Distinct from prefill() rather than an option on it: prefill returns logits for one position and its callers sample that row, so widening what prefill returns would silently move which row a sampler reads.

Requires a head built wide enough to evaluate more than the final position – see the family's head-width configuration. Implemented by QwenTransformer; the other families throw until their heads take a width.

Parameters

input Token indices [1, T], T >= 2. Batching is not supported: the targets are the sequence's own next tokens, so two sequences in one call would need per-row lengths the shape cannot carry.

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

Mila/Src/Dnn/Core/LanguageModelNetwork.ixx

Public Types

Public Member Functions

Additional Inherited Members

Detailed Description

Member Function Documentation

◆ backward()

◆ decode()

◆ forward()

◆ prefill()

◆ prefillFrom()

◆ rewindKvCache()

◆ scoreTokens()