Mila
Deep Neural Network Library
Loading...
Searching...
No Matches
Mila::Dnn::ITransformerBlock< TDeviceType, TPrecision > Class Template Referenceabstractexport
module Dnn.Components.ITransformerBlock

Polymorphic inference interface for one transformer block. More...

Inheritance diagram for Mila::Dnn::ITransformerBlock< TDeviceType, TPrecision >:
Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, false, TWeightQuantization, TKvCachePolicy, kDelegatedFeedForward, kMixtureOfExperts > Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, true, TWeightQuantization, NoKvCompression, kDelegatedFeedForward, kMixtureOfExperts > Mila::Dnn::QwenAttentionBlock< TDeviceType, TPrecision, TWeightPlan, TKvCachePolicy > Mila::Dnn::QwenDeltaNetBlock< TDeviceType, TPrecision, TWeightPlan > Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy, kDelegatedFeedForward, kMixtureOfExperts > Mila::Dnn::QwenAttentionBlock< TDeviceType, TPrecision, TWeightPlan, TKvPolicy > Mila::Dnn::QwenDeltaNetBlock< TDeviceType, TPrecision, TWeightPlan >

Public Types

using MR = typename DeviceTypeTraits<TDeviceType>::memory_resource
using TensorType = Tensor<TPrecision, MR>

Public Member Functions

virtual TensorType & decode (const TensorType &input, dim_t position)=0
 Single-token decode at an absolute position (T == 1).
virtual TensorType & prefill (const TensorType &input, dim_t position_offset)=0
 Chunked prefill: process [B, T_chunk, model_dim] at an absolute offset.
virtual void resetKvCache ()=0
 Reset the KV cache (new generation session).
virtual bool rewindKvCache (dim_t position)=0
 Rewind the KV cache fill position for prompt-prefix reuse.
virtual void setState (const GqaState &state)=0
 Wire the shared GQA transient workspace (owned by the transformer).
virtual bool supportsKvCache () const noexcept=0
 True when the block's attention supports the KV-cache inference path.

Detailed Description

template<DeviceType TDeviceType, TensorDataType TPrecision>
class Mila::Dnn::ITransformerBlock< TDeviceType, TPrecision >

Polymorphic inference interface for one transformer block.

Template Parameters
TDeviceTypeCompile-time device.
TPrecisionActivation/compute precision (must match across the block list).

Member Function Documentation

◆ decode()

template<DeviceType TDeviceType, TensorDataType TPrecision>
virtual TensorType & Mila::Dnn::ITransformerBlock< TDeviceType, TPrecision >::decode ( const TensorType & input,
dim_t position )
pure virtual

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

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

Implemented in Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, false, TWeightQuantization, TKvCachePolicy, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, true, TWeightQuantization, NoKvCompression, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::QwenAttentionBlock< TDeviceType, TPrecision, TWeightPlan, TKvPolicy >, Mila::Dnn::QwenAttentionBlock< TDeviceType, TPrecision, TWeightPlan, TKvCachePolicy >, Mila::Dnn::QwenDeltaNetBlock< TDeviceType, TPrecision, TWeightPlan >, and Mila::Dnn::QwenDeltaNetBlock< TDeviceType, TPrecision, TWeightPlan >.

◆ prefill()

template<DeviceType TDeviceType, TensorDataType TPrecision>
virtual TensorType & Mila::Dnn::ITransformerBlock< TDeviceType, TPrecision >::prefill ( const TensorType & input,
dim_t position_offset )
pure virtual

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

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

Implemented in Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, false, TWeightQuantization, TKvCachePolicy, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, true, TWeightQuantization, NoKvCompression, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::QwenAttentionBlock< TDeviceType, TPrecision, TWeightPlan, TKvPolicy >, Mila::Dnn::QwenAttentionBlock< TDeviceType, TPrecision, TWeightPlan, TKvCachePolicy >, Mila::Dnn::QwenDeltaNetBlock< TDeviceType, TPrecision, TWeightPlan >, and Mila::Dnn::QwenDeltaNetBlock< TDeviceType, TPrecision, TWeightPlan >.

◆ resetKvCache()

template<DeviceType TDeviceType, TensorDataType TPrecision>
virtual void Mila::Dnn::ITransformerBlock< TDeviceType, TPrecision >::resetKvCache ( )
pure virtual

Reset the KV cache (new generation session).

Implemented in Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, false, TWeightQuantization, TKvCachePolicy, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, true, TWeightQuantization, NoKvCompression, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::QwenAttentionBlock< TDeviceType, TPrecision, TWeightPlan, TKvPolicy >, Mila::Dnn::QwenAttentionBlock< TDeviceType, TPrecision, TWeightPlan, TKvCachePolicy >, Mila::Dnn::QwenDeltaNetBlock< TDeviceType, TPrecision, TWeightPlan >, and Mila::Dnn::QwenDeltaNetBlock< TDeviceType, TPrecision, TWeightPlan >.

◆ rewindKvCache()

template<DeviceType TDeviceType, TensorDataType TPrecision>
virtual bool Mila::Dnn::ITransformerBlock< TDeviceType, TPrecision >::rewindKvCache ( dim_t position)
pure virtual

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).

Implemented in Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, false, TWeightQuantization, TKvCachePolicy, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, true, TWeightQuantization, NoKvCompression, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::QwenAttentionBlock< TDeviceType, TPrecision, TWeightPlan, TKvPolicy >, Mila::Dnn::QwenAttentionBlock< TDeviceType, TPrecision, TWeightPlan, TKvCachePolicy >, Mila::Dnn::QwenDeltaNetBlock< TDeviceType, TPrecision, TWeightPlan >, and Mila::Dnn::QwenDeltaNetBlock< TDeviceType, TPrecision, TWeightPlan >.

◆ setState()

template<DeviceType TDeviceType, TensorDataType TPrecision>
virtual void Mila::Dnn::ITransformerBlock< TDeviceType, TPrecision >::setState ( const GqaState & state)
pure virtual

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

Called once after build, before any prefill/decode.

Implemented in Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, false, TWeightQuantization, TKvCachePolicy, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, true, TWeightQuantization, NoKvCompression, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::QwenAttentionBlock< TDeviceType, TPrecision, TWeightPlan, TKvPolicy >, Mila::Dnn::QwenAttentionBlock< TDeviceType, TPrecision, TWeightPlan, TKvCachePolicy >, Mila::Dnn::QwenDeltaNetBlock< TDeviceType, TPrecision, TWeightPlan >, and Mila::Dnn::QwenDeltaNetBlock< TDeviceType, TPrecision, TWeightPlan >.

◆ supportsKvCache()

template<DeviceType TDeviceType, TensorDataType TPrecision>
virtual bool Mila::Dnn::ITransformerBlock< TDeviceType, TPrecision >::supportsKvCache ( ) const
pure virtualnoexcept

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

Implemented in Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, false, TWeightQuantization, TKvCachePolicy, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, true, TWeightQuantization, NoKvCompression, kDelegatedFeedForward, kMixtureOfExperts >, Mila::Dnn::QwenAttentionBlock< TDeviceType, TPrecision, TWeightPlan, TKvPolicy >, Mila::Dnn::QwenAttentionBlock< TDeviceType, TPrecision, TWeightPlan, TKvCachePolicy >, Mila::Dnn::QwenDeltaNetBlock< TDeviceType, TPrecision, TWeightPlan >, and Mila::Dnn::QwenDeltaNetBlock< TDeviceType, TPrecision, TWeightPlan >.

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