precomputed_kv_layer

precomputed_kv_layer

Cross-attention layer with K/V dimension projection for bridging two embedding spaces.

PrecomputedKVCrossAttentionLayer

PrecomputedKVCrossAttentionLayer(embedding_dimension, conditioning_key_value_dimension, number_of_heads, number_of_key_value_heads, head_dimension, feedforward_dimension, normalization_type=value, conditioning_dimension=None, use_gating=False, dropout=0.1, activation=value)

Bases: Module

Projects precomputed conditioning K/V to local dimension, cross-attends, then feedforward.

Bridges two embedding spaces via learned K/V projections. The conditioning cache provides precomputed K/V from an external source which may have a different hidden dimension. Projections map them into the local attention space before cross-attention.

Initialize PrecomputedKVCrossAttentionLayer.

Parameters:

Name	Type	Description	Default
embedding_dimension	int	Hidden dimension of the main input stream.	required
conditioning_key_value_dimension	int	K/V dimension from the conditioning source.	required
number_of_heads	int	Number of attention heads.	required
number_of_key_value_heads	int	Number of K/V heads.	required
head_dimension	int	Dimension per attention head.	required
feedforward_dimension	int	FFN hidden dimension.	required
normalization_type	str	Normalization type for attention and FFN blocks.	value
conditioning_dimension	int | None	Dimension of conditioning vector for adaptive norm.	None
use_gating	bool	Whether to use gating in adaptive normalization.	False
dropout	float	Dropout rate for residual connections.	0.1
activation	str	Activation function for FFN.	value

Source code in src/versatil/models/layers/transformer/layer/precomputed_kv_layer.py

def __init__(
    self,
    embedding_dimension: int,
    conditioning_key_value_dimension: int,
    number_of_heads: int,
    number_of_key_value_heads: int,
    head_dimension: int,
    feedforward_dimension: int,
    normalization_type: str = NormalizationType.RMS_NORM.value,
    conditioning_dimension: int | None = None,
    use_gating: bool = False,
    dropout: float = 0.1,
    activation: str = ActivationFunction.SILU.value,
):
    """Initialize PrecomputedKVCrossAttentionLayer.

    Args:
        embedding_dimension: Hidden dimension of the main input stream.
        conditioning_key_value_dimension: K/V dimension from the conditioning source.
        number_of_heads: Number of attention heads.
        number_of_key_value_heads: Number of K/V heads.
        head_dimension: Dimension per attention head.
        feedforward_dimension: FFN hidden dimension.
        normalization_type: Normalization type for attention and FFN blocks.
        conditioning_dimension: Dimension of conditioning vector for adaptive norm.
        use_gating: Whether to use gating in adaptive normalization.
        dropout: Dropout rate for residual connections.
        activation: Activation function for FFN.
    """
    super().__init__()
    local_key_value_dimension = number_of_key_value_heads * head_dimension
    self.key_projection = nn.Linear(
        conditioning_key_value_dimension, local_key_value_dimension, bias=False
    )
    self.value_projection = nn.Linear(
        conditioning_key_value_dimension, local_key_value_dimension, bias=False
    )
    if number_of_key_value_heads == number_of_heads:
        attention_type = AttentionType.MULTI_HEAD.value
    else:
        attention_type = AttentionType.GROUPED_QUERY.value
    self.cross_attention_block = PrecomputedCrossAttentionBlock(
        attention=CachedAttention(
            embedding_dimension=embedding_dimension,
            number_of_heads=number_of_heads,
            number_of_key_value_heads=number_of_key_value_heads,
            head_dimension=head_dimension,
            dropout=dropout,
            bias=False,
            attention_type=attention_type,
        ),
        normalization=create_block_normalization(
            normalization_type=normalization_type,
            dimension=embedding_dimension,
            conditioning_dimension=conditioning_dimension,
            use_gating=use_gating,
        ),
        dropout=dropout,
    )
    self.feedforward_block = FeedforwardBlock(
        feedforward=build_feedforward(
            embedding_dimension=embedding_dimension,
            feedforward_dimension=feedforward_dimension,
            activation=activation,
            dropout=dropout,
            bias=False,
        ),
        normalization=create_block_normalization(
            normalization_type=normalization_type,
            dimension=embedding_dimension,
            conditioning_dimension=conditioning_dimension,
            use_gating=use_gating,
        ),
        dropout=dropout,
    )

forward

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

Project conditioning K/V, cross-attend with optional RoPE, then FFN.

Parameters:

Name	Type	Description	Default
hidden_states	Tensor	Local stream tokens (B, S, D).	required
conditioning_cache	ConditioningLayerCache	Precomputed K/V from conditioning source. Keys and values have shape (B, P, conditioning_kv_dim).	required
conditioning	Tensor | None	Conditioning vector for adaptive normalization (B, C). Ignored when normalization is unconditioned.	None
attention_mask	Tensor | None	Optional mask (B, 1, S, P).	None
precomputed_rope	tuple[Tensor, Tensor] | None	Optional precomputed (cos, sin) rotary positional encodings for query.	None

Returns:

Type	Description
Tensor	Updated hidden states (B, S, D).

Source code in src/versatil/models/layers/transformer/layer/precomputed_kv_layer.py

def forward(
    self,
    hidden_states: torch.Tensor,
    conditioning_cache: ConditioningLayerCache,
    conditioning: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    precomputed_rope: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> torch.Tensor:
    """Project conditioning K/V, cross-attend with optional RoPE, then FFN.

    Args:
        hidden_states: Local stream tokens (B, S, D).
        conditioning_cache: Precomputed K/V from conditioning source.
            Keys and values have shape (B, P, conditioning_kv_dim).
        conditioning: Conditioning vector for adaptive normalization (B, C).
            Ignored when normalization is unconditioned.
        attention_mask: Optional mask (B, 1, S, P).
        precomputed_rope: Optional precomputed (cos, sin) rotary positional encodings for query.

    Returns:
        Updated hidden states (B, S, D).
    """
    projected_keys = self.key_projection(conditioning_cache.keys)
    projected_values = self.value_projection(conditioning_cache.values)
    hidden_states = self.cross_attention_block(
        hidden_states=hidden_states,
        keys=projected_keys,
        values=projected_values,
        conditioning=conditioning,
        attention_mask=attention_mask,
        precomputed_query_rope=precomputed_rope,
    )
    hidden_states = self.feedforward_block(
        hidden_states=hidden_states,
        conditioning=conditioning,
    )
    return hidden_states