cross_attention

cross_attention

Cross-attention block: norm -> cross-attention -> gated residual.

CrossAttentionBlock

CrossAttentionBlock(attention, normalization, dropout=0.1)

Bases: TransformerBlock

Norm -> cross-attention to encoder hidden states -> gated residual.

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

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

precompute_kv

precompute_kv(encoded_features)

Precompute K/V projections for static conditioning.

Parameters:

Name	Type	Description	Default
encoded_features	Tensor	Encoder features (B, memory_length, D).	required

Returns:

Type	Description
ConditioningLayerCache	ConditioningLayerCache with projected keys and values.

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

def precompute_kv(self, encoded_features: torch.Tensor) -> ConditioningLayerCache:
    """Precompute K/V projections for static conditioning.

    Args:
        encoded_features: Encoder features (B, memory_length, D).

    Returns:
        ConditioningLayerCache with projected keys and values.
    """
    return ConditioningLayerCache(
        keys=self.attention.compute_key(encoded_features),
        values=self.attention.compute_value(encoded_features),
    )

forward

forward(hidden_states, encoder_hidden_states=None, conditioning=None, attention_mask=None, conditioning_cache=None)

Norm -> cross-attention -> gated residual.

Parameters:

Name	Type	Description	Default
hidden_states	Tensor	Query input (B, T, D).	required
encoder_hidden_states	Tensor | None	Encoder output for K/V projection (B, S, D). Can be None when using conditioning_cache.	None
conditioning	Tensor | None	Conditioning vector for AdaNorm (B, C). Ignored by UnconditionedNorm.	None
attention_mask	Tensor | None	Bool mask (B, 1, T, S), True = masked.	None
conditioning_cache	ConditioningLayerCache | None	Precomputed K/V for static conditioning.	None

Returns:

Type	Description
Tensor	Output hidden states (B, T, D).

Raises:

Type	Description
ValueError	If both encoder_hidden_states and conditioning_cache are None.

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

def forward(
    self,
    hidden_states: torch.Tensor,
    encoder_hidden_states: torch.Tensor | None = None,
    conditioning: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    conditioning_cache: ConditioningLayerCache | None = None,
) -> torch.Tensor:
    """Norm -> cross-attention -> gated residual.

    Args:
        hidden_states: Query input (B, T, D).
        encoder_hidden_states: Encoder output for K/V projection (B, S, D).
            Can be None when using conditioning_cache.
        conditioning: Conditioning vector for AdaNorm (B, C). Ignored by UnconditionedNorm.
        attention_mask: Bool mask (B, 1, T, S), True = masked.
        conditioning_cache: Precomputed K/V for static conditioning.

    Returns:
        Output hidden states (B, T, D).

    Raises:
        ValueError: If both encoder_hidden_states and conditioning_cache are None.
    """
    if encoder_hidden_states is None and conditioning_cache is None:
        raise ValueError(
            "Either encoder_hidden_states or conditioning_cache must be provided"
        )
    residual = hidden_states
    hidden_states, gate = self.normalization(
        x=hidden_states, condition=conditioning
    )
    attention_output, _ = self.attention(
        query_input=hidden_states,
        key_input=encoder_hidden_states if conditioning_cache is None else None,
        value_input=encoder_hidden_states if conditioning_cache is None else None,
        attention_mask=attention_mask,
        conditioning_cache=conditioning_cache,
    )
    hidden_states = self.apply_residual(residual, attention_output, gate)
    return hidden_states