One Gemma 4 decoder block; kGlobal selects the global (full-attention) geometry. More...

Inheritance diagram for Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy >:

Classes
struct	BlockBuildContexts
	The per-child build contexts and split-scratch geometry this block implies. More...

Public Types
using	AttentionType = GroupedQueryAttention<TDeviceType, TPrecision, TKvPolicy>
using	CompositeComponentBase = CompositeComponent<TDeviceType, TPrecision>
using	GeGLUType = Swiglu<TDeviceType, TPrecision, ActivationType::Gelu>
using	LinearType = Linear<TDeviceType, TPrecision, TWeightQuant>
using	MR = typename DeviceTypeTraits<TDeviceType>::memory_resource
using	ResidualType = Residual<TDeviceType, TPrecision>
using	RmsNormType = RmsNorm<TDeviceType, TPrecision>
using	RopeType = Rope<TDeviceType, TPrecision>
using	TensorType = Tensor<TPrecision, MR>
Public Types inherited from Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >
using	ComponentBase = Component<TDeviceType, TPrecision>
using	ComponentPtr = std::shared_ptr<Component<TDeviceType, TPrecision>>
Public Types inherited from Mila::Dnn::IDecoderLayer< TDeviceType, TPrecision >
using	MR = typename DeviceTypeTraits<TDeviceType>::memory_resource
using	TensorType = Tensor<TPrecision, MR>

Public Member Functions
	GemmaBlock (const std::string &name, const GemmaConfig &config, std::optional< DeviceId > device_id=std::nullopt)
TensorType &	decode (const TensorType &input, dim_t position) override
	Single-token decode at an absolute position (T == 1).
MemoryStats	getMemoryStats () const override
	Return the current memory allocation breakdown for this component.
std::vector< std::string >	getParameterNames () const override
	List all available parameter names for this component.
MemoryStats	getRequiredMemory (const BuildContext &context) const override
	What onBuilding() would allocate for this context, without allocating.
const ComponentType	getType () const override
	Get the component type identifier.
dim_t	headDim () const noexcept
void	installSharedWorkspace (const GemmaBlockWorkspace< TDeviceType, TPrecision > &workspace)
	Install the transformer-owned shared activation workspace (pooling).
bool	keyEqualsValue () const noexcept
dim_t	kvProjWidth () const noexcept
void	load_ (ModelArchive &archive, SerializationMode mode) override
	Children through the base traversal, then this block's own layer_scalar.
void	loadParameter (const std::string &name, const ITensorBlob &blob) override
	Load the block's own parameters.
dim_t	numKVHeads () const noexcept
dim_t	packedQKVWidth () const noexcept
TensorType &	prefill (const TensorType &input, dim_t position_offset) override
	Chunked prefill: process [B, T_chunk, model_dim] at an absolute offset.
dim_t	qProjWidth () const noexcept
void	resetKVCache () override
	Reset the KV cache (new generation session).
bool	rewindKvCache (dim_t position) override
	Rewind the KV cache fill position for prompt-prefix reuse.
float	ropeTheta () const noexcept
dim_t	rotaryDim () const noexcept
void	save_ (ModelArchive &archive, SerializationMode mode) const override
void	saveFlatTensors (Serialization::SafeTensorsWriter &writer, const std::string &prefix, Serialization::TensorSavePass pass) const override
	Children through the base traversal, plus the block's own layer_scalar.
void	setState (const GqaState &state) override
	Wire the shared GQA transient workspace (owned by the transformer).
void	setUseFlashDecode (bool enabled)
void	setUseFlashPrefill (bool enabled)
bool	supportsKVCache () const noexcept override
	True when the block's attention supports the KV-cache inference path.
dim_t	window () const noexcept
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.
DeviceId	getDeviceId () const override
	Get the compute device for this composite.
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	synchronize () override
	Synchronize all child components.
std::string	toString () const override
	Generate a human-readable description.
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).
const std::string	getName () const
	Get the component's name identifier.
TrainingMode	getTrainingMode () const noexcept
	The current runtime behavioral mode of this Component.
virtual bool	isBuilt () const final
	Returns true if build() has completed successfully.
void	setTrainingMode (TrainingMode mode)
	Set the runtime behavioral mode for this Component.
virtual void	zeroGradients ()
	Clear all model-owned gradients for this component.

Static Public Member Functions
static constexpr bool	isGlobal () noexcept
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 Member Functions
void	onBuilding (const BuildContext &context) override
	Hook invoked by build() to allocate component buffers.
void	onTrainingModeChanging (TrainingMode training_mode) override
	Hook invoked when training mode is about to change.
BlockBuildContexts	resolveBlockBuildContexts (const BuildContext &context) const
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.
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.
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.

Additional Inherited Members
Protected Types inherited from Mila::Dnn::Component< TDeviceType, TPrecision >
using	HostStagingMemoryResource
	Host memory a device-resident parameter stages through.
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, bool kGlobal, WeightQuantPolicy TWeightQuant = NoWeightQuant, KvCachePolicy TKvPolicy = NoKvCompression>
requires PrecisionSupportedOnDevice<TPrecision, TDeviceType>
class Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy >

One Gemma 4 decoder block; kGlobal selects the global (full-attention) geometry.

Member Function Documentation

◆ decode()

template<DeviceType TDeviceType, TensorDataType TPrecision, bool kGlobal, WeightQuantPolicy TWeightQuant = NoWeightQuant, KvCachePolicy TKvPolicy = NoKvCompression>

TensorType & Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy >::decode	(	const TensorType &	input,
		dim_t	position )

inlineoverridevirtual

Single-token decode at an absolute position (T == 1).

Returns: Reference to the block-owned output [B, 1, model_dim].

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

◆ getMemoryStats()

template<DeviceType TDeviceType, TensorDataType TPrecision, bool kGlobal, WeightQuantPolicy TWeightQuant = NoWeightQuant, KvCachePolicy TKvPolicy = NoKvCompression>

MemoryStats Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy >::getMemoryStats ( ) const

inlineoverridevirtual

Return the current memory allocation breakdown for this component.

Reflects allocations at the moment of the call. The returned stats naturally track the component lifecycle:

After construction – nothing; construction allocates none After build( Inference ) – parameters + T=1 state buffers After build( Training ) – parameters + T=full state buffers After setTrainingMode( Train ) – parameters + state + gradients

For CompositeComponent and Network, the returned stats are the recursive aggregate of all child components.

May be called at any time – no lifecycle preconditions.

Returns: MemoryStats reflecting current allocations.

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

◆ getParameterNames()

template<DeviceType TDeviceType, TensorDataType TPrecision, bool kGlobal, WeightQuantPolicy TWeightQuant = NoWeightQuant, KvCachePolicy TKvPolicy = NoKvCompression>

std::vector< std::string > Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy >::getParameterNames ( ) const

inlineoverridevirtual

List all available parameter names for this component.

Returns an empty vector by default. Leaf components with parameters should override to return their canonical parameter name list in the same stable order used by save_() and loadParameter().

Reimplemented from Mila::Dnn::Component< TDeviceType, TPrecision >.

◆ getRequiredMemory()

template<DeviceType TDeviceType, TensorDataType TPrecision, bool kGlobal, WeightQuantPolicy TWeightQuant = NoWeightQuant, KvCachePolicy TKvPolicy = NoKvCompression>

MemoryStats Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy >::getRequiredMemory ( const BuildContext & context ) const

inlineoverridevirtual

What onBuilding() would allocate for this context, without allocating.

Children are named rather than walked, because each is built with its OWN context and a generic recursion over getComponents() would size every one of them against the block's shape. Children are fetched by name rather than read from the member pointers, which onBuilding() assigns and are therefore null before a build. See Specifications/MemoryFootprint.md section 4.4.

Reimplemented from Mila::Dnn::Component< TDeviceType, TPrecision >.

◆ getType()

template<DeviceType TDeviceType, TensorDataType TPrecision, bool kGlobal, WeightQuantPolicy TWeightQuant = NoWeightQuant, KvCachePolicy TKvPolicy = NoKvCompression>

const ComponentType Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy >::getType ( ) const

inlineoverridevirtual

Get the component type identifier.

Used for serialization and runtime type identification.

Returns: Component type enum value.

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

◆ installSharedWorkspace()

template<DeviceType TDeviceType, TensorDataType TPrecision, bool kGlobal, WeightQuantPolicy TWeightQuant = NoWeightQuant, KvCachePolicy TKvPolicy = NoKvCompression>

void Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy >::installSharedWorkspace ( const GemmaBlockWorkspace< TDeviceType, TPrecision > & workspace )

inline

Install the transformer-owned shared activation workspace (pooling).

One workspace serves every layer: the inference path is strictly sequential, so exactly one block is live at a time. Must be called before build(); onBuilding then routes each slot into the matching child component via installSharedOutput (and keeps the q/k/v split scratch for the block's own views), skipping all per-layer output self-allocation. Self-allocation remains the default for standalone blocks (tests).

◆ load_()

template<DeviceType TDeviceType, TensorDataType TPrecision, bool kGlobal, WeightQuantPolicy TWeightQuant = NoWeightQuant, KvCachePolicy TKvPolicy = NoKvCompression>

void Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy >::load_	(	ModelArchive &	archive,
		SerializationMode	mode )

inlineoverridevirtual

Children through the base traversal, then this block's own layer_scalar.

CompositeComponent::load_ only recurses – the own-parameter half is Component's default, which a composite does not inherit – so both are driven explicitly.

Reimplemented from Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >.

◆ loadParameter()

template<DeviceType TDeviceType, TensorDataType TPrecision, bool kGlobal, WeightQuantPolicy TWeightQuant = NoWeightQuant, KvCachePolicy TKvPolicy = NoKvCompression>

void Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy >::loadParameter	(	const std::string &	name,
		const ITensorBlob &	blob )

inlineoverridevirtual

Load the block's own parameters.

Sub-component weights are routed by the reader directly to the sub-components; the only parameter owned by the block itself is the Gemma 4 Unified per-layer output scale layer_scalar (a [1] FP32 value).

Reimplemented from Mila::Dnn::Component< TDeviceType, TPrecision >.

◆ onBuilding()

template<DeviceType TDeviceType, TensorDataType TPrecision, bool kGlobal, WeightQuantPolicy TWeightQuant = NoWeightQuant, KvCachePolicy TKvPolicy = NoKvCompression>

void Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy >::onBuilding ( const BuildContext & context )

inlineoverrideprotectedvirtual

Hook invoked by build() to allocate component buffers.

Receives the stored BuildContext. Implementations must use config.allocationSeqLen() when sizing output buffers – this is the single call that makes Inference and Training allocate the correct buffer sizes automatically without per-component logic.

// Example -- Linear component:
shape_t out_shape =
{
    config.batchSize(),
    config.allocationSeqLen(),   // 1 for Inference, T for Training
    config_.getOutputFeatures()
};
output_ = std::make_unique<TensorType>( device, out_shape,
    this->getName() + ".output" );

The default implementation forwards to the legacy onBuilding( const shape_t& ) overload for backwards compatibility. New components should override this overload directly.

Note: Do not call build() or onBuilding() from within this hook.; Implementations should either succeed fully or leave no partial state, as a failed build() may be retried.

Takes the build-time configuration; use its allocationSeqLen() to obtain the correct output buffer sequence dimension.

Reimplemented from Mila::Dnn::Component< TDeviceType, TPrecision >.

◆ onTrainingModeChanging()

template<DeviceType TDeviceType, TensorDataType TPrecision, bool kGlobal, WeightQuantPolicy TWeightQuant = NoWeightQuant, KvCachePolicy TKvPolicy = NoKvCompression>

void Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy >::onTrainingModeChanging ( TrainingMode training_mode )

inlineoverrideprotectedvirtual

Hook invoked when training mode is about to change.

Propagates the new mode to all child components. The hook runs with the Component's training mutex held; it MUST NOT call setTrainingMode().

Parameters

training_mode New training mode (Normal or Eval)

Reimplemented from Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >.

◆ prefill()

template<DeviceType TDeviceType, TensorDataType TPrecision, bool kGlobal, WeightQuantPolicy TWeightQuant = NoWeightQuant, KvCachePolicy TKvPolicy = NoKvCompression>

TensorType & Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy >::prefill	(	const TensorType &	input,
		dim_t	position_offset )

inlineoverridevirtual

Chunked prefill: process [B, T_chunk, model_dim] at an absolute offset.

Returns: Reference to the block-owned output [B, T_chunk, model_dim].

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

◆ resetKVCache()

template<DeviceType TDeviceType, TensorDataType TPrecision, bool kGlobal, WeightQuantPolicy TWeightQuant = NoWeightQuant, KvCachePolicy TKvPolicy = NoKvCompression>

void Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy >::resetKVCache ( )

inlineoverridevirtual

Reset the KV cache (new generation session).

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

◆ rewindKvCache()

template<DeviceType TDeviceType, TensorDataType TPrecision, bool kGlobal, WeightQuantPolicy TWeightQuant = NoWeightQuant, KvCachePolicy TKvPolicy = NoKvCompression>

bool Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy >::rewindKvCache ( dim_t position )

inlineoverridevirtual

Rewind the KV cache fill position for prompt-prefix reuse.

Keeps the cache session live; positions [0, position) stay valid.

Returns: true when the layer's attention accepted the rewind (a bounded sliding-window ring refuses when the stale tail has overwritten the window a continuation would attend to).

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

◆ save_()

template<DeviceType TDeviceType, TensorDataType TPrecision, bool kGlobal, WeightQuantPolicy TWeightQuant = NoWeightQuant, KvCachePolicy TKvPolicy = NoKvCompression>

void Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy >::save_	(	ModelArchive &	archive,
		SerializationMode	mode ) const

inlineoverridevirtual

Reimplemented from Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >.

◆ saveFlatTensors()

template<DeviceType TDeviceType, TensorDataType TPrecision, bool kGlobal, WeightQuantPolicy TWeightQuant = NoWeightQuant, KvCachePolicy TKvPolicy = NoKvCompression>

void Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy >::saveFlatTensors	(	Serialization::SafeTensorsWriter &	writer,
		const std::string &	prefix,
		Serialization::TensorSavePass	pass ) const

inlineoverridevirtual

Children through the base traversal, plus the block's own layer_scalar.

This block is the one composite in the tree that owns a parameter itself, so it needs both halves. The hand-rolled version this replaces did neither correctly: it re-implemented the base recursion without pushing a scope per child, so every child collided on one path, and it never wrote layer_scalar at all – a Gemma archive silently lost the per-layer output scales, which is a numerics change, not a missing extra.

Children through the base recursion, then this block's own layer_scalar.

CompositeComponent::saveFlatTensors only recurses – the own-parameter half is Component's default, which a composite overrides away – so a composite owning a parameter must drive both halves explicitly. Omitting this dropped all 48 layer_scalar tensors from a Gemma export while every structural check passed; they are in the converted .bin, so they are part of the flat vocabulary.

Order is identical across both passes, which the writer requires.

Reimplemented from Mila::Dnn::CompositeComponent< TDeviceType, TPrecision >.

◆ setState()

template<DeviceType TDeviceType, TensorDataType TPrecision, bool kGlobal, WeightQuantPolicy TWeightQuant = NoWeightQuant, KvCachePolicy TKvPolicy = NoKvCompression>

void Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy >::setState ( const GqaState & state )

inlineoverridevirtual

Wire the shared GQA transient workspace (owned by the transformer).

Called once after build, before any prefill/decode.

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

◆ supportsKVCache()

template<DeviceType TDeviceType, TensorDataType TPrecision, bool kGlobal, WeightQuantPolicy TWeightQuant = NoWeightQuant, KvCachePolicy TKvPolicy = NoKvCompression>

bool Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy >::supportsKVCache ( ) const

inlineoverridevirtualnoexcept

True when the block's attention supports the KV-cache inference path.

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

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

Mila/Src/Dnn/Components/Transformers/Gemma/Gemma.Block.ixx

Classes

Public Types

Public Member Functions

Static Public Member Functions

Protected Member Functions

Additional Inherited Members

Detailed Description

Member Function Documentation

◆ decode()

◆ getMemoryStats()

◆ getParameterNames()

◆ getRequiredMemory()

◆ getType()

◆ installSharedWorkspace()

◆ load_()

◆ loadParameter()

◆ onBuilding()

◆ onTrainingModeChanging()

◆ prefill()

◆ resetKVCache()

◆ rewindKvCache()

◆ save_()

◆ saveFlatTensors()

◆ setState()

◆ supportsKVCache()