self_attention

self_attention

Self-attention block: norm -> self-attention -> gated residual.

SelfAttentionBlock

SelfAttentionBlock(attention, normalization, dropout=0.1)

Bases: TransformerBlock

Norm -> self-attention -> gated residual.

Source code in src/versatil/models/layers/transformer/block/self_attention.py

def __init__(
    self,
    attention: CachedAttention,
    normalization: BlockNormalization,
    dropout: float = 0.1,
):
    super().__init__(normalization=normalization, dropout=dropout)
    self.attention = attention

forward

forward(hidden_states, conditioning=None, attention_mask=None, positional_encoding=None, generation_cache=None)

Norm -> self-attention -> gated residual.

Parameters:

Name	Type	Description	Default
hidden_states	Tensor	Input embeddings (B, S, D).	required
conditioning	Tensor | None	Conditioning vector for AdaNorm (B, C). Ignored by UnconditionedNorm.	None
attention_mask	Tensor | None	Bool mask (B, 1, S, S), True = masked.	None
positional_encoding	RotaryPositionalEncoding | None	Optional RoPE module.	None
generation_cache	GenerationLayerCache | None	Cached K/V from previous autoregressive steps. When provided, an updated cache is returned.	None

Returns:

Type	Description
tuple[Tensor, GenerationLayerCache | None]	Tuple of (output hidden states (B, S, D), updated GenerationLayerCache or None).

Source code in src/versatil/models/layers/transformer/block/self_attention.py

def forward(
    self,
    hidden_states: torch.Tensor,
    conditioning: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    positional_encoding: RotaryPositionalEncoding | None = None,
    generation_cache: GenerationLayerCache | None = None,
) -> tuple[torch.Tensor, GenerationLayerCache | None]:
    """Norm -> self-attention -> gated residual.

    Args:
        hidden_states: Input embeddings (B, S, D).
        conditioning: Conditioning vector for AdaNorm (B, C). Ignored by UnconditionedNorm.
        attention_mask: Bool mask (B, 1, S, S), True = masked.
        positional_encoding: Optional RoPE module.
        generation_cache: Cached K/V from previous autoregressive steps. When
            provided, an updated cache is returned.

    Returns:
        Tuple of (output hidden states (B, S, D), updated GenerationLayerCache or None).
    """
    residual = hidden_states
    hidden_states, gate = self.normalization(
        x=hidden_states, condition=conditioning
    )
    attention_output, new_cache = self.attention(
        query_input=hidden_states,
        key_input=hidden_states,
        value_input=hidden_states,
        attention_mask=attention_mask,
        positional_encoding=positional_encoding,
        generation_cache=generation_cache,
    )
    hidden_states = self.apply_residual(residual, attention_output, gate)
    return hidden_states, new_cache