Grouped-Query Attention module that accepts concatenated QKV input. More...

Inheritance diagram for Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >:

Public Types
using	ComponentBase = Component<TDeviceType, TComputePrecision>
using	KvCacheTensorType = Tensor<kCacheDtype, MR>
using	MR = typename DeviceTypeTraits<TDeviceType>::memory_resource
using	OpType = typename Compute::OperationTraits<Compute::OperationType::GroupedQueryAttentionOp, TDeviceType, TComputePrecision, TKvPolicy>::type
using	TensorType = Tensor<TComputePrecision, MR>

Public Member Functions
	GroupedQueryAttention (const std::string &name, const GqaConfig &config, std::optional< DeviceId > device_id=std::nullopt)
	Construct a GroupedQueryAttention component.
TensorType &	backward (const TensorType &, const TensorType &)
	Not implemented – GQA is inference only, and this is where that is decided.
TensorType &	decode (const TensorType &q, const TensorType &k, const TensorType &v, dim_t position_offset)
	Inference-only single-token decode pass.
TensorType &	forward (const TensorType &input)
	Standard forward pass.
const GqaConfig &	getConfig () const noexcept
DeviceId	getDeviceId () const override
	Get the compute device id associated with this component.
std::vector< ITensor * >	getGradients () const override
	Return non-owning pointers to parameter gradient tensors.
MemoryStats	getMemoryStats () const override
	Return the current memory allocation breakdown for this component.
dim_t	getModelDim () const noexcept
dim_t	getNumHeads () const noexcept
dim_t	getNumKvHeads () const noexcept
std::vector< ObservableStage >	getObservableStages () const override
	Stages this component publishes, and the passes each is published on.
std::vector< const ITensor * >	getOutputs () const override
	Return non-owning pointers to the tensors this component produces.
std::vector< ITensor * >	getParameters () const override
	Return non-owning pointers to parameter tensors.
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.
void	installSharedOutput (std::shared_ptr< TensorType > output)
	Install a shared prefill-output slot (activation pooling).
dim_t	parameterCount () const override
	Return number of trainable parameters.
TensorType &	prefill (const TensorType &q, const TensorType &k, const TensorType &v, dim_t position_offset)
	Chunked prefill pass with explicit position offset.
void	resetKvCache ()
	Reset the KV cache for a new generation session.
bool	rewindKvCache (dim_t position)
	Rewind the cache fill position for prompt-prefix reuse (PromptCaching.md).
void	save_ (ModelArchive &, SerializationMode) const override
void	setState (const GqaState &state)
	Forward the shared transient workspace to the underlying operation.
void	setUseFlashDecode (bool enabled)
	Route the BF16 decode through the fused decode-attention kernel.
void	setUseFlashPrefill (bool enabled)
	Route the BF16 prefill through the fused FlashAttention kernel.
bool	supportsKvCache () const noexcept
	Returns true when the underlying operation implements both IPositionalUnaryOp and IKvCacheLifecycle.
void	synchronize () override
	Wait for outstanding device work submitted by this component.
std::string	toString () const override
	Produce a short, human-readable description of the component.
Public Member Functions inherited from Mila::Dnn::Component< TDeviceType, TComputePrecision >
	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.
const std::string	getName () const
	Get the component's name identifier.
Compute::ComputePassMask	getObservedPasses () const noexcept
virtual std::vector< std::string >	getParameterNames () const
	List all available parameter names for this component.
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	load_ (ModelArchive &archive, SerializationMode mode)
	Restore this component's parameters from its archive scope.
virtual void	loadParameter (const std::string &, const Serialization::ITensorBlob &)
	Load a parameter from serialized tensor data.
virtual void	requireSerializableParameters () const
	Verify this component can serialize whatever parameters it owns.
virtual void	saveFlatTensors (Serialization::SafeTensorsWriter &writer, const std::string &prefix, Serialization::TensorSavePass pass) const
	Drive this component's tensors through one pass of a flat safetensors save.
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.

Static Public Attributes
static constexpr TensorDataType	kCacheDtype
static constexpr bool	kKvCompressed = TKvPolicy::kIsActive

Protected Member Functions
void	onBuilding (const BuildContext &context) override
	Hook invoked by build() to allocate component buffers.
void	onExecutionContextSet () override
	Lifecycle hook: Called immediately after ExecutionContext is set.
void	onTrainingModeChanging (TrainingMode training_mode) override
	Hook called before TrainingMode transitions.
Protected Member Functions inherited from Mila::Dnn::Component< TDeviceType, TComputePrecision >
bool	hasExecutionContext () const noexcept
	Check if execution context has been set.
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.
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.
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>".
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
Static Public Member Functions inherited from Mila::Dnn::Component< TDeviceType, TComputePrecision >
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, TComputePrecision >
using	HostStagingMemoryResource
	Host memory a device-resident parameter stages through.
Protected Attributes inherited from Mila::Dnn::Component< TDeviceType, TComputePrecision >
BuildContext	build_context_
	The BuildContext stored at build time.

Detailed Description

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>
requires PrecisionSupportedOnDevice<TComputePrecision, TDeviceType>
class Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >

Grouped-Query Attention module that accepts concatenated QKV input.

GQA generalises MHA by allowing num_kv_heads < num_heads. Each K/V head is shared by a group of (num_heads / num_kv_heads) Q heads, reducing KV cache memory and bandwidth during inference.

The module requires a single input tensor in model-layout containing concatenated Q, K and V along the feature axis:

input shape == [B, T, (num_heads + 2 * num_kv_heads) * head_dim] output shape == [B, T, model_dim] (model_dim = num_heads * head_dim)

The backend compute implementation (registered as "GroupedQueryAttentionOp") must accept this layout and produce the output above.

KV-cache inference is an optional backend capability. After build(), supportsKvCache() indicates whether the underlying operation implements both IPositionalUnaryOp (prefill/decode dispatch) and IKvCacheLifecycle (cache init/reset). Both pointers are resolved once at build time.

The cache self-initializes on the first prefill/forward; resetKvCache() is the public hook the owning decoder layer / transformer drives to start a new generation session.

REVIEW: resetKvCache() is public (initialization stays internal to prefill). When TransformerBase<> is introduced as the common base for GptTransformer, LlamaTransformer, MistralTransformer etc., revisit whether it should become private with 'friend class TransformerBase<TDeviceType, TPrecision>' to enforce that only the generate() orchestration path may manage the KV cache lifecycle.

Constructor & Destructor Documentation

◆ GroupedQueryAttention()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >::GroupedQueryAttention	(	const std::string &	name,
		const GqaConfig &	config,
		std::optional< DeviceId >	device_id = std::nullopt )

inlineexplicit

Construct a GroupedQueryAttention component.

Parameters

name	Component name identifier (mandatory).
config	GQA configuration (model_dim, num_heads, num_kv_heads).
device_id	Optional DeviceId to create an owned ExecutionContext (standalone / unit-test mode).

Member Function Documentation

◆ backward()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

TensorType & Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >::backward	(	const TensorType &	,
		const TensorType &	)

inline

Not implemented – GQA is inference only, and this is where that is decided.

No GQA backend implements a backward pass. The refusal is stated here, at the component's own boundary, rather than reached by dispatching into an operation whose body is an unconditional throw: the built/training-mode guards below it were checking preconditions for a call that could never succeed, and the tail after the dispatch was unreachable code the compiler reported.

The retired sketch ran on the expanded [B,NH,T,HS] layout that CudaGqaOp no longer allocates; whether a real backward keeps that layout for training or is derived on the compact NKV layout is an open design question recorded in GqaMemory.md.

MultiHeadAttention implements both directions and is the component to train through.

Exceptions

std::runtime_error Always.

◆ decode()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

TensorType & Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >::decode	(	const TensorType &	q,
		const TensorType &	k,
		const TensorType &	v,
		dim_t	position_offset )

inline

Inference-only single-token decode pass.

When the backend implements IPositionalUnaryOp and the cache has been populated by a prior forward() call, uses the fast O(n) KV cache path. When the backend does not support positional dispatch, falls back to forward(). The caller never needs to know which path was taken.

Precondition: forward() must have been called at least once to populate the KV cache before decode() is called.

Parameters

q	Query tensor [B, 1, Q * head_dim].
k	Key tensor [B, 1, KV * head_dim].
v	Value tensor [B, 1, KV * head_dim].
position_offset	Absolute position of the token (0-based).

Returns: Reference to component-owned single-token output tensor.

◆ forward()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

TensorType & Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >::forward ( const TensorType & input )

inline

Standard forward pass.

Always available regardless of backend. When the backend supports KV caching, the first forward() call initialises and populates the cache (prefill with position_offset=0). When called again after decode() steps, it automatically resets the cache and begins a new prefill session – no explicit session management required by callers.

Parameters

input Concatenated QKV input [B, T, (Q + 2*KV) * head_dim].

Returns: Reference to component-owned output tensor [B, T, model_dim].

◆ getDeviceId()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

DeviceId Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >::getDeviceId ( ) const

inlineoverridevirtual

Get the compute device id associated with this component.

Must return the device on which parameters and operations execute.

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

◆ getGradients()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

std::vector< ITensor * > Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >::getGradients ( ) const

inlineoverridevirtual

Return non-owning pointers to parameter gradient tensors.

Gradient buffers are allocated only when the component is built in training mode, so a component built for inference returns an empty vector. Stateless components return empty in either mode. This is the accessor counterpart to getParameters() and does not throw on mode.

Returns: Vector of gradient pointers; empty when built for inference.

Exceptions

std::runtime_error if called before the component has been built.

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

◆ getMemoryStats()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

MemoryStats Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, 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, TComputePrecision >.

◆ getObservableStages()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

std::vector< ObservableStage > Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >::getObservableStages ( ) const

inlineoverridevirtual

Stages this component publishes, and the passes each is published on.

Lets a caller tell an uninstrumented component from a clean one before running anything. The pass qualifier is load-bearing rather than decorative: the stage set genuinely differs by pass, since a prefill can produce intermediates its decode never materializes.

Returns: Vector of stages; empty for components that publish nothing.

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

◆ getOutputs()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

std::vector< const ITensor * > Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >::getOutputs ( ) const

inlineoverridevirtual

Return non-owning pointers to the tensors this component produces.

The counterpart to getParameters(): that reports what a component HOLDS, this reports what it PRODUCES. Tensors are the built allocations, so their extents are the ceiling a forward pass may narrow within – an observer wanting the live value of a particular call receives it through publication instead.

An empty result means the component does not describe itself. That absence cannot be mistaken for a finding – a view renders no row for it – so no distinction beyond empty is offered.

Returns: Vector of output pointers; empty before build, and for undescribed components.

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

◆ getParameters()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

std::vector< ITensor * > Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >::getParameters ( ) const

inlineoverridevirtual

Return non-owning pointers to parameter tensors.

The returned tensor pointers remain valid for the lifetime of the component. Order should be canonical (weights before biases).

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

◆ getRequiredMemory()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

MemoryStats Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >::getRequiredMemory ( const BuildContext & context ) const

inlineoverridevirtual

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

The KV cache – the dominant, context-scaling term – is reported by the operation, which is where it is allocated. See Specifications/MemoryFootprint.md.

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

◆ getType()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

const ComponentType Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, 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, TComputePrecision >.

◆ installSharedOutput()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

void Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >::installSharedOutput ( std::shared_ptr< TensorType > output )

inline

Install a shared prefill-output slot (activation pooling).

Must be called before build(): onBuilding then skips the prefill output self-allocation after validating the slot's storage covers the build shape; prefill() already always returns a shape-adjusted view. The tiny per-token decode output (decode_output_) stays component-owned. Mirrors Linear::installSharedWeight; self-allocation remains the default. The slot is owned and memory-accounted by the installer.

◆ onBuilding()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

void Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, 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, TComputePrecision >.

◆ onExecutionContextSet()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

void Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >::onExecutionContextSet ( )

inlineoverrideprotectedvirtual

Lifecycle hook: Called immediately after ExecutionContext is set.

Override this to perform initialization that requires a valid ExecutionContext. At the time this is called, getExecutionContext() is guaranteed to return a valid context.

Common uses:

Composite components: Create and configure child components.
Device resource allocation: Query device capabilities.

Default implementation does nothing.

Exceptions

Any	exception thrown will cause setExecutionContext() to fail and restore the component to a "context not set" state.

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

◆ onTrainingModeChanging()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

void Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >::onTrainingModeChanging ( TrainingMode training_mode )

inlineoverrideprotectedvirtual

Hook called before TrainingMode transitions.

Called by setTrainingMode() after validation and lock acquisition, before the internal state is updated. Derived classes override to respond to the transition – e.g. zeroing gradient buffers on transition to Eval, or re-enabling dropout on transition to Training.

The default implementation is a no-op.

Takes the incoming TrainingMode.

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

◆ parameterCount()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

dim_t Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >::parameterCount ( ) const

inlineoverridevirtual

Return number of trainable parameters.

For leaf components this is the element count of owned parameter tensors. CompositeComponent and Network implementations should return the recursive aggregate across all children.

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

◆ prefill()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

TensorType & Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >::prefill	(	const TensorType &	q,
		const TensorType &	k,
		const TensorType &	v,
		dim_t	position_offset )

inline

Chunked prefill pass with explicit position offset.

Called by the transformer block during chunked prefill. The KV cache must already be initialized (via onBuilding or forward()).

Parameters

q	Query tensor [B, T_chunk, Q * head_dim].
k	Key tensor [B, T_chunk, KV * head_dim].
v	Value tensor [B, T_chunk, KV * head_dim].
position_offset	Absolute position of the first token in this chunk.

Returns: Reference to component-owned output tensor.

◆ resetKvCache()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

void Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >::resetKvCache ( )

inline

Reset the KV cache for a new generation session.

Drops any active decode state so the next prefill starts a fresh sequence. A no-op on backends without KV-cache support, and harmless when no decode session is active.

◆ rewindKvCache()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

bool Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >::rewindKvCache ( dim_t position )

inline

Rewind the cache fill position for prompt-prefix reuse (PromptCaching.md).

Unlike resetKvCache() the cache session stays live: initialization state and device contents are untouched, and positions [0, position) remain valid for a subsequent prefillFrom.

Returns: true when the underlying operation accepted the rewind.

◆ save_()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

void Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >::save_	(	ModelArchive &	,
		SerializationMode	) const

inlineoverridevirtual

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

◆ setState()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

void Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >::setState ( const GqaState & state )

inline

Forward the shared transient workspace to the underlying operation.

Must be called after build() and before prefill() or decode(). A no-op on backends that do not implement setState (e.g. CPU stub).

Parameters

state Non-owning pointers to the shared GQA scratch tensors.

◆ setUseFlashDecode()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

void Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >::setUseFlashDecode ( bool enabled )

inline

Route the BF16 decode through the fused decode-attention kernel.

Only the CUDA BF16 ops honor this; other backends, FP32, and unsupported geometries keep the cuBLASLt decode pipeline. Unlike setUseFlashPrefill there is no workspace-width coupling – the fused path draws its split-K scratch from the shared execution-context buffer at decode time.

◆ setUseFlashPrefill()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

void Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >::setUseFlashPrefill ( bool enabled )

inline

Route the BF16 prefill through the fused FlashAttention kernel.

Only the CUDA BF16 ops honor this – the unbounded/global kernel and the bounded sliding-window ring variant; other backends and FP32 ignore it. The transformer couples this to the shared preatt/att workspace width (flash on -> the O(chunk x T_ctx) score buffer is reclaimed), so it must be set consistently with that sizing – a narrow workspace with flash off would overflow. See GqaFlashAttention.md 5.6.

◆ supportsKvCache()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

bool Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >::supportsKvCache ( ) const

inlinenoexcept

Returns true when the underlying operation implements both IPositionalUnaryOp and IKvCacheLifecycle.

Resolved once at build time. CPU backends return false; CUDA backends return true when CudaGroupedQueryAttentionOp is in use.

◆ synchronize()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

void Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >::synchronize ( )

inlineoverridevirtual

Wait for outstanding device work submitted by this component.

On CPU this may be a no-op. Use to ensure results are visible to the host or to measure synchronous timings.

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

◆ toString()

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

std::string Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >::toString ( ) const

inlineoverridevirtual

Produce a short, human-readable description of the component.

Implementations should keep output concise and avoid throwing.

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

Member Data Documentation

◆ kCacheDtype

template<DeviceType TDeviceType, TensorDataType TComputePrecision, KvCachePolicy TKvPolicy = NoKvCompression>

TensorDataType Mila::Dnn::GroupedQueryAttention< TDeviceType, TComputePrecision, TKvPolicy >::kCacheDtype

staticconstexpr

Initial value:

                                            = [] {
    if constexpr ( requires { TKvPolicy::kStorageDtype; } )
        return TKvPolicy::kStorageDtype;
    else
        return TComputePrecision;
}()

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

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

Public Types

Public Member Functions

Static Public Attributes

Protected Member Functions

Additional Inherited Members

Detailed Description

Constructor & Destructor Documentation

◆ GroupedQueryAttention()

Member Function Documentation

◆ backward()

◆ decode()

◆ forward()

◆ getDeviceId()

◆ getGradients()

◆ getMemoryStats()

◆ getObservableStages()

◆ getOutputs()

◆ getParameters()

◆ getRequiredMemory()

◆ getType()

◆ installSharedOutput()

◆ onBuilding()

◆ onExecutionContextSet()

◆ onTrainingModeChanging()

◆ parameterCount()

◆ prefill()

◆ resetKvCache()

◆ rewindKvCache()

◆ save_()

◆ setState()

◆ setUseFlashDecode()

◆ setUseFlashPrefill()

◆ supportsKvCache()

◆ synchronize()

◆ toString()

Member Data Documentation

◆ kCacheDtype