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

Polymorphic inference interface for one decoder layer. More...

Inheritance diagram for Mila::Dnn::IDecoderLayer< TDeviceType, TPrecision >:
Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, false, TWeightQuantization, TKvCachePolicy > Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, true, TWeightQuantization, NoKvCompression > Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, kGlobal, TWeightQuant, TKvPolicy >

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::IDecoderLayer< TDeviceType, TPrecision >

Polymorphic inference interface for one decoder layer.

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

Member Function Documentation

◆ decode()

template<DeviceType TDeviceType, TensorDataType TPrecision>
virtual TensorType & Mila::Dnn::IDecoderLayer< 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 >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, false, TWeightQuantization, TKvCachePolicy >, and Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, true, TWeightQuantization, NoKvCompression >.

◆ prefill()

template<DeviceType TDeviceType, TensorDataType TPrecision>
virtual TensorType & Mila::Dnn::IDecoderLayer< 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 >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, false, TWeightQuantization, TKvCachePolicy >, and Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, true, TWeightQuantization, NoKvCompression >.

◆ resetKVCache()

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

Reset the KV cache (new generation session).

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

◆ rewindKvCache()

template<DeviceType TDeviceType, TensorDataType TPrecision>
virtual bool Mila::Dnn::IDecoderLayer< 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 >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, false, TWeightQuantization, TKvCachePolicy >, and Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, true, TWeightQuantization, NoKvCompression >.

◆ setState()

template<DeviceType TDeviceType, TensorDataType TPrecision>
virtual void Mila::Dnn::IDecoderLayer< 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 >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, false, TWeightQuantization, TKvCachePolicy >, and Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, true, TWeightQuantization, NoKvCompression >.

◆ supportsKVCache()

template<DeviceType TDeviceType, TensorDataType TPrecision>
virtual bool Mila::Dnn::IDecoderLayer< 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 >, Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, false, TWeightQuantization, TKvCachePolicy >, and Mila::Dnn::GemmaBlock< TDeviceType, TPrecision, true, TWeightQuantization, NoKvCompression >.

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