Mila::Dnn::GqaConfig Class Reference

Configuration class for the Grouped-Query Attention module. More...

Inheritance diagram for Mila::Dnn::GqaConfig:

Public Member Functions
	GqaConfig (dim_t model_dim, dim_t num_heads, dim_t num_kv_heads)
	Constructor with all required parameters.
void	fromMetadata (const SerializationMetadata &meta) override
	Populate configuration from serialization metadata.
float	getAttentionScale () const noexcept
	Attention logit scale: the explicit value, or 1/sqrt(head_dim) when unset.
dim_t	getGroupSize () const noexcept
	Number of Q heads sharing each K/V head (num_heads / num_kv_heads).
dim_t	getHeadDim () const noexcept
	Per-head feature dimension (model_dim / num_heads).
dim_t	getModelDim () const noexcept
	Total Q-projection output width (= num_heads * head_dim).
dim_t	getNumHeads () const noexcept
	Number of Q attention heads.
dim_t	getNumKvHeads () const noexcept
	Number of K/V attention heads.
dim_t	getWindow () const noexcept
	Sliding-attention window size (0 = unbounded/global causal).
SerializationMetadata	toMetadata () const override
	Convert configuration to serialization metadata.
std::string	toString () const override
	Human-readable description of the configuration.
void	validate () const override
	Validate all configuration parameters.
template<typename Self>
decltype(auto)	withAttentionScale (this Self &&self, float scale)
	Fluent setter for the attention logit scale (QK softmax multiplier).
template<typename Self>
decltype(auto)	withModelDim (this Self &&self, dim_t model_dim)
	Fluent setter for model dimension.
template<typename Self>
decltype(auto)	withNumHeads (this Self &&self, dim_t num_heads)
	Fluent setter for number of Q heads.
template<typename Self>
decltype(auto)	withNumKvHeads (this Self &&self, dim_t num_kv_heads)
	Fluent setter for number of K/V heads.
template<typename Self>
decltype(auto)	withWindow (this Self &&self, dim_t window)
	Fluent setter for the sliding-attention window size.
Public Member Functions inherited from Mila::Dnn::ComponentConfig
virtual	~ComponentConfig ()=default
	Virtual destructor for polymorphic base.

Detailed Description

Configuration class for the Grouped-Query Attention module.

Carries the three parameters that uniquely define a GQA layer:

• model_dim : total Q-projection output width (= num_heads * head_dim)
• num_heads : number of Q attention heads
• num_kv_heads: number of K/V attention heads (must divide num_heads)

The derived head_dim = model_dim / num_heads and group_size = num_heads / num_kv_heads are computed on demand rather than stored, so they always stay consistent with the three primary fields.

Fluent setters follow the C++23 explicit-object-parameter pattern used throughout the codebase, enabling value-category-preserving method chaining on both lvalue and rvalue configs.

Constructor & Destructor Documentation

GqaConfig()

Mila::Dnn::GqaConfig::GqaConfig ( dim_t model_dim, dim_t num_heads, dim_t num_kv_heads )

inline

Constructor with all required parameters.

Parameters

model_dim	Total Q-projection output width.
num_heads	Number of Q attention heads.
num_kv_heads	Number of K/V attention heads.

Member Function Documentation

fromMetadata()

void Mila::Dnn::GqaConfig::fromMetadata ( const SerializationMetadata & meta )

inlineoverridevirtual

Populate configuration from serialization metadata.

Reads available fields from the provided metadata and updates the configuration object in place. Missing keys are silently ignored so that older checkpoints without num_kv_heads fall back to the constructor-supplied default.

Implements Mila::Dnn::ComponentConfig.

getGroupSize()

dim_t Mila::Dnn::GqaConfig::getGroupSize ( ) const

inlinenoexcept

Number of Q heads sharing each K/V head (num_heads / num_kv_heads).

Equivalent to the GQA "group size". A value of 1 recovers standard MHA; a value equal to num_heads recovers Multi-Query Attention.

getHeadDim()

dim_t Mila::Dnn::GqaConfig::getHeadDim ( ) const

inlinenoexcept

Per-head feature dimension (model_dim / num_heads).

Derived quantity – always consistent with model_dim and num_heads.

getNumKvHeads()

dim_t Mila::Dnn::GqaConfig::getNumKvHeads ( ) const

inlinenoexcept

Number of K/V attention heads.

The KV cache is sized proportionally to this value, giving GQA its memory bandwidth advantage over full MHA.

toMetadata()

SerializationMetadata Mila::Dnn::GqaConfig::toMetadata ( ) const

inlineoverridevirtual

Convert configuration to serialization metadata.

Produces a SerializationMetadata object containing all configuration fields suitable for writing into an archive by the caller.

Implements Mila::Dnn::ComponentConfig.

toString()

std::string Mila::Dnn::GqaConfig::toString ( ) const

inlineoverridevirtual

Human-readable description of the configuration.

Implements Mila::Dnn::ComponentConfig.

validate()

void Mila::Dnn::GqaConfig::validate ( ) const

inlineoverridevirtual

Validate all configuration parameters.

Checks:

1. model_dim > 0
2. num_heads >= 2
3. model_dim % num_heads == 0 (integer head_dim)
4. num_kv_heads >= 1
5. num_kv_heads <= num_heads (KV heads cannot exceed Q heads)
6. num_heads % num_kv_heads == 0 (integer group size)

Exceptions

std::invalid_argument

If any constraint is violated.

Implements Mila::Dnn::ComponentConfig.

withAttentionScale()

template<typename Self>

decltype(auto) Mila::Dnn::GqaConfig::withAttentionScale ( this Self && self, float scale )

inline

Fluent setter for the attention logit scale (QK softmax multiplier).

0 (the default) means "derive 1/sqrt(head_dim)" – the standard scaling for Llama/Qwen/MHA. A positive value overrides it; Gemma uses 1.0 because its QK-norm already controls the dot-product magnitude (HF self.scaling = 1.0).

withWindow()

template<typename Self>

decltype(auto) Mila::Dnn::GqaConfig::withWindow ( this Self && self, dim_t window )

inline

Fluent setter for the sliding-attention window size.

The number of most-recent keys a query may attend to. 0 (the default) means unbounded/global causal attention - every query attends over [0, abs_t], which is the behavior for Llama / Qwen / Gemma global layers. A positive value bounds the lower edge to [abs_t - window + 1, abs_t] (Gemma sliding layers: 1024).

---

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

• Mila/Src/Dnn/Components/Attention/GQA/GroupedQueryAttention.Config.ixx