autoregressive_decoder

autoregressive_decoder

GPT-style autoregressive decoder.

GPTDecoder

GPTDecoder(number_of_layers, embedding_dimension, number_of_heads, number_of_key_value_heads=None, feedforward_dimension=None, dropout=0.1, attention_dropout=0.0, activation=value, normalization_type=value, attention_type=value, use_cross_attention=False, positional_encoding_type=None, maximum_sequence_length=2048, bias=True, normalization_epsilon=1e-06, initializer_range=0.02)

Bases: TransformerMixin, Module

GPT-style autoregressive decoder, with KV caching, extended to support cross-attention.

Stacks multiple TransformerDecoderLayer modules and manages KV cache across layers. Applies causal masking for autoregressive generation.

Initialize GPT decoder.

Parameters:

Name	Type	Description	Default
number_of_layers	int	Number of decoder layers	required
embedding_dimension	int	Model embedding dimension	required
number_of_heads	int	Number of attention heads	required
number_of_key_value_heads	int | None	Number of K/V heads (for GQA)	None
feedforward_dimension	int | None	FFN hidden dimension	None
dropout	float	Dropout probability for residual connections	0.1
attention_dropout	float	Dropout probability for attention weights	0.0
activation	str	Activation function (use ActivationFunction enum values)	value
normalization_type	str	Type of normalization (use NormalizationType enum values)	value
attention_type	str	Type of attention (use AttentionType enum values)	value
use_cross_attention	bool	Whether to use cross-attention (False for decoder-only models)	False
positional_encoding_type	str | None	Type of positional encoding (use PositionalEncodingType enum values, or None)	None
maximum_sequence_length	int	Maximum sequence length for positional encoding	2048
bias	bool	Whether to use bias in linear layers	True
normalization_epsilon	float	Epsilon for normalization layers	1e-06
initializer_range	float	Standard deviation for weight initialization	0.02

Source code in src/versatil/models/layers/transformer/autoregressive_decoder.py

def __init__(
    self,
    number_of_layers: int,
    embedding_dimension: int,
    number_of_heads: int,
    number_of_key_value_heads: int | None = None,
    feedforward_dimension: int | None = None,
    dropout: float = 0.1,
    attention_dropout: float = 0.0,
    activation: str = ActivationFunction.SWIGLU.value,
    normalization_type: str = NormalizationType.RMS_NORM.value,
    attention_type: str = AttentionType.GROUPED_QUERY.value,
    use_cross_attention: bool = False,
    positional_encoding_type: str | None = None,
    maximum_sequence_length: int = 2048,
    bias: bool = True,
    normalization_epsilon: float = 1e-6,
    initializer_range: float = 0.02,
):
    """Initialize GPT decoder.

    Args:
        number_of_layers: Number of decoder layers
        embedding_dimension: Model embedding dimension
        number_of_heads: Number of attention heads
        number_of_key_value_heads: Number of K/V heads (for GQA)
        feedforward_dimension: FFN hidden dimension
        dropout: Dropout probability for residual connections
        attention_dropout: Dropout probability for attention weights
        activation: Activation function (use ActivationFunction enum values)
        normalization_type: Type of normalization (use NormalizationType enum values)
        attention_type: Type of attention (use AttentionType enum values)
        use_cross_attention: Whether to use cross-attention (False for decoder-only models)
        positional_encoding_type: Type of positional encoding (use PositionalEncodingType enum values, or None)
        maximum_sequence_length: Maximum sequence length for positional encoding
        bias: Whether to use bias in linear layers
        normalization_epsilon: Epsilon for normalization layers
        initializer_range: Standard deviation for weight initialization
    """
    super().__init__()

    self.number_of_layers = number_of_layers
    self.embedding_dimension = embedding_dimension
    self.number_of_heads = number_of_heads
    self.maximum_sequence_length = maximum_sequence_length
    self.use_cross_attention = use_cross_attention
    self.initializer_range = initializer_range
    self.number_of_residual_blocks = (
        3 if use_cross_attention else 2
    )  # Self-Attention + Feedforward
    self.number_of_key_value_heads, self.head_dimension = (
        self._resolve_attention_dimensions(
            embedding_dimension=embedding_dimension,
            number_of_heads=number_of_heads,
            number_of_key_value_heads=number_of_key_value_heads,
            attention_type=attention_type,
        )
    )
    self._setup_positional_encoding(
        positional_encoding_type=positional_encoding_type,
        embedding_dimension=embedding_dimension,
        maximum_sequence_length=maximum_sequence_length,
        number_of_heads=number_of_heads,
    )

    self.layers = nn.ModuleList(
        [
            TransformerDecoderLayer(
                embedding_dimension=embedding_dimension,
                number_of_heads=number_of_heads,
                number_of_key_value_heads=number_of_key_value_heads,
                feedforward_dimension=feedforward_dimension,
                dropout=dropout,
                attention_dropout=attention_dropout,
                activation=activation,
                normalization_type=normalization_type,
                attention_type=attention_type,
                use_cross_attention=use_cross_attention,
                bias=bias,
                normalization_epsilon=normalization_epsilon,
            )
            for _ in range(number_of_layers)
        ]
    )

    self.final_normalization = create_normalization_layer(
        normalization_type=normalization_type,
        dimension=embedding_dimension,
        epsilon=normalization_epsilon,
    )
    self.apply(self._init_weights)

create_empty_generation_cache

create_empty_generation_cache(batch_size, device, dtype=float32)

Create an initial empty GenerationCache for autoregressive generation.

Parameters:

Name	Type	Description	Default
batch_size	int	Batch size.	required
device	device | str	Device for cache tensors.	required
dtype	dtype	Data type for cache tensors.	float32

Returns:

Type	Description
GenerationCache	GenerationCache with empty layers ready for the first generation step.

Source code in src/versatil/models/layers/transformer/autoregressive_decoder.py

def create_empty_generation_cache(
    self,
    batch_size: int,
    device: torch.device | str,
    dtype: torch.dtype = torch.float32,
) -> GenerationCache:
    """Create an initial empty GenerationCache for autoregressive generation.

    Args:
        batch_size: Batch size.
        device: Device for cache tensors.
        dtype: Data type for cache tensors.

    Returns:
        GenerationCache with empty layers ready for the first generation step.
    """
    return GenerationCache(
        layers=initialize_generation_cache(
            batch_size=batch_size,
            num_layers=self.number_of_layers,
            number_of_heads=self.number_of_key_value_heads,
            head_dimension=self.head_dimension,
            device=device,
            dtype=dtype,
        ),
    )

precompute_conditioning_kv

precompute_conditioning_kv(encoded_features)

Precompute conditioning K/V for all layers for forward pass reuse.

Source code in src/versatil/models/layers/transformer/autoregressive_decoder.py

def precompute_conditioning_kv(
    self,
    encoded_features: torch.Tensor,
) -> ConditioningCache:
    """Precompute conditioning K/V for all layers for forward pass reuse."""
    return precompute_conditioning(
        layers=self.layers,
        encoded_features=encoded_features,
    )

forward

forward(hidden_states, encoded_features=None, self_attention_mask=None, cross_attention_mask=None, key_padding_mask=None, generation_cache=None, conditioning_cache=None)

Forward pass through decoder.

Parameters:

Name	Type	Description	Default
hidden_states	Tensor	Input token embeddings (B, query_length, D).	required
encoded_features	Tensor | None	Encoder features (B, num_features, D). Required when use_cross_attention=True and no conditioning_cache.	None
self_attention_mask	Tensor | None	Custom causal mask (B, 1, query_length, query_length), True = masked. If None, generates standard triangular causal mask.	None
cross_attention_mask	Tensor | None	Mask for cross-attention (B, 1, query_length, key_length), True = masked.	None
key_padding_mask	Tensor | None	Padding mask for observation tokens (B, query_length), True = masked.	None
generation_cache	GenerationCache | None	Cached K/V from previous generation steps. When provided, an updated cache is returned.	None
conditioning_cache	ConditioningCache | None	Precomputed K/V for static conditioning. When provided, encoded_features is not needed for cross-attention.	None

Returns:

Type	Description
tuple[Tensor, GenerationCache | None]	Tuple of (output (B, query_length, D), updated GenerationCache or None).

Source code in src/versatil/models/layers/transformer/autoregressive_decoder.py

def forward(
    self,
    hidden_states: torch.Tensor,
    encoded_features: torch.Tensor | None = None,
    self_attention_mask: torch.Tensor | None = None,
    cross_attention_mask: torch.Tensor | None = None,
    key_padding_mask: torch.Tensor | None = None,
    generation_cache: GenerationCache | None = None,
    conditioning_cache: ConditioningCache | None = None,
) -> tuple[torch.Tensor, GenerationCache | None]:
    """Forward pass through decoder.

    Args:
        hidden_states: Input token embeddings (B, query_length, D).
        encoded_features: Encoder features (B, num_features, D). Required when
            use_cross_attention=True and no conditioning_cache.
        self_attention_mask: Custom causal mask (B, 1, query_length, query_length),
            True = masked. If None, generates standard triangular causal mask.
        cross_attention_mask: Mask for cross-attention (B, 1, query_length, key_length),
            True = masked.
        key_padding_mask: Padding mask for observation tokens (B, query_length),
            True = masked.
        generation_cache: Cached K/V from previous generation steps. When provided,
            an updated cache is returned.
        conditioning_cache: Precomputed K/V for static conditioning. When provided,
            encoded_features is not needed for cross-attention.

    Returns:
        Tuple of (output (B, query_length, D), updated GenerationCache or None).
    """
    batch_size = hidden_states.shape[0]
    device = hidden_states.device
    query_length = hidden_states.shape[1]
    cache_length = generation_cache.get_length() if generation_cache else 0
    cached_key_padding_mask = (
        generation_cache.key_padding_mask if generation_cache else None
    )
    total_mask, full_key_padding_mask = create_full_padding_mask(
        key_padding_mask=key_padding_mask,
        cached_key_padding_mask=cached_key_padding_mask,
        self_attention_mask=self_attention_mask,
        batch_size=batch_size,
        query_length=query_length,
        cache_length=cache_length,
        device=device,
    )
    hidden_states, rope_pe = self._apply_positional_encoding(
        hidden_states=hidden_states, offset=cache_length
    )
    use_cache = generation_cache is not None
    new_layer_caches = []
    for layer_index, layer in enumerate(self.layers):
        layer_generation_cache = (
            generation_cache.layers[layer_index]
            if generation_cache is not None
            else None
        )
        hidden_states, new_layer_cache = layer(
            hidden_states=hidden_states,
            encoded_features=encoded_features,
            self_attention_mask=total_mask,
            cross_attention_mask=cross_attention_mask,
            generation_cache=layer_generation_cache,
            conditioning_cache=conditioning_cache[layer_index]
            if conditioning_cache
            else None,
            positional_encoding=rope_pe,
        )
        if use_cache:
            new_layer_caches.append(new_layer_cache)

    hidden_states = self.final_normalization(hidden_states)
    new_generation_cache = None
    if use_cache:
        new_generation_cache = GenerationCache(
            layers=new_layer_caches,
            key_padding_mask=full_key_padding_mask,
        )

    return hidden_states, new_generation_cache