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

Public Types
using	MR = typename DeviceTypeTraits<TDeviceType>::memory_resource
using	NetworkBase = Network<TDeviceType, TPrecision>
using	StageProbe = std::function<void( std::string_view stage, const TensorType& value )>
	Observer called with each intermediate activation during prefill.
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
	LanguageNetwork (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.
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)
	Rewind the KV caches to position for prompt-prefix reuse (PromptCaching.md).
virtual void	setStageProbe (StageProbe probe)
	Install a stage probe, or clear it by passing an empty function.
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.
IExecutionContext *	getExecutionContext () const
	Public access to the network's shared execution context.
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.
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.
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)
void	saveFlatTensors (Serialization::SafeTensorsWriter &writer, const std::string &prefix, Serialization::TensorSavePass pass) const override
	Recurse into children, extending the flat dotted prefix.
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).
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< 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.
TrainingMode	getTrainingMode () const noexcept
	The current runtime behavioral mode of this Component.
virtual bool	isBuilt () const final
	Returns true if build() has completed successfully.
virtual void	loadParameter (const std::string &, const Serialization::ITensorBlob &)
	Load a parameter from serialized tensor data.
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
	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 PretrainedMetadata &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 >
IExecutionContext *	getExecutionContext () const
	Get the shared execution context.
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.
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.

Member Typedef Documentation

◆ StageProbe

template<DeviceType TDeviceType, TensorDataType TPrecision>

using Mila::Dnn::LanguageNetwork< TDeviceType, TPrecision >::StageProbe = std::function<void( std::string_view stage, const TensorType& value )>

Observer called with each intermediate activation during prefill.

Diagnostic. Two loads of one model can hold byte-identical parameters and still compute differently, and only the activations show where they diverge. A probe on the real prefill path rather than a parallel diagnostic one, because a second implementation is free to not reproduce the bug.

Member Function Documentation

◆ backward()

template<DeviceType TDeviceType, TensorDataType TPrecision>

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

pure virtual

Full backward pass (training).

Parameters

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

Returns: Gradient w.r.t. the input embeddings.

◆ decode()

template<DeviceType TDeviceType, TensorDataType TPrecision>

virtual TensorType & Mila::Dnn::LanguageNetwork< 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::LanguageNetwork< TDeviceType, TPrecision >::forward ( const TokenIndexType & input )

pure virtual

Full-sequence forward pass.

Parameters

input Token indices [B, T].

Returns: Logits [B, T, vocab_size].

◆ prefill()

template<DeviceType TDeviceType, TensorDataType TPrecision>

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

pure virtual

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

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::LanguageNetwork< 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.

Default: unsupported. Networks that implement the reuse path (Gemma) override both this and rewindKvCache; callers only reach prefillFrom after a successful rewind, so the default is never hit in practice.

◆ rewindKvCache()

template<DeviceType TDeviceType, TensorDataType TPrecision>

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

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.

Reimplemented in Mila::Dnn::GemmaTransformer< TDeviceType, TPrecision, TWeightQuantization, TKvCachePolicy >.

◆ setStageProbe()

template<DeviceType TDeviceType, TensorDataType TPrecision>

virtual void Mila::Dnn::LanguageNetwork< TDeviceType, TPrecision >::setStageProbe ( StageProbe probe )

inlinevirtual

Install a stage probe, or clear it by passing an empty function.

Default is a no-op so a network that has no instrumentation is unaffected.

Reimplemented in Mila::Dnn::GemmaTransformer< TDeviceType, TPrecision, TWeightQuantization, TKvCachePolicy >.

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

Mila/Src/Dnn/Core/LanguageNetwork.ixx

Public Types

Public Member Functions

Additional Inherited Members

Member Typedef Documentation

◆ StageProbe

Member Function Documentation

◆ backward()

◆ decode()

◆ forward()

◆ prefill()

◆ prefillFrom()

◆ rewindKvCache()

◆ setStageProbe()