precomputed_cross_attention

precomputed_cross_attention

Cross-attention block with precomputed K/V and optional query RoPE.

PrecomputedCrossAttentionBlock

PrecomputedCrossAttentionBlock(attention, normalization, dropout=0.1)

Bases: TransformerBlock

Norm -> query projection -> optional RoPE -> cross-attention -> gated residual.

Accepts precomputed K/V tensors (already in head-split format) and only projects queries.

Source code in src/versatil/models/layers/transformer/block/precomputed_cross_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, keys, values, conditioning=None, attention_mask=None, precomputed_query_rope=None)

Norm -> cross-attention with precomputed K/V -> gated residual.

Parameters:

Name	Type	Description	Default
hidden_states	Tensor	Query input (B, T, D).	required
keys	Tensor	Precomputed keys (B, S, kv_dim).	required
values	Tensor	Precomputed values (B, S, kv_dim).	required
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
precomputed_query_rope	tuple[Tensor, Tensor] | None	Precomputed (cos, sin) for query positions. Applied via half-rotation after query projection.	None

Returns:

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

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

def forward(
    self,
    hidden_states: torch.Tensor,
    keys: torch.Tensor,
    values: torch.Tensor,
    conditioning: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    precomputed_query_rope: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> torch.Tensor:
    """Norm -> cross-attention with precomputed K/V -> gated residual.

    Args:
        hidden_states: Query input (B, T, D).
        keys: Precomputed keys (B, S, kv_dim).
        values: Precomputed values (B, S, kv_dim).
        conditioning: Conditioning vector for AdaNorm (B, C). Ignored by UnconditionedNorm.
        attention_mask: Bool mask (B, 1, T, S), True = masked.
        precomputed_query_rope: Precomputed (cos, sin) for query positions.
            Applied via half-rotation after query projection.

    Returns:
        Output hidden states (B, T, D).
    """
    residual = hidden_states
    hidden_states, gate = self.normalization(
        x=hidden_states, condition=conditioning
    )
    keys = self._reshape_to_heads(keys)
    values = self._reshape_to_heads(values)
    queries = self.attention.compute_query(hidden_states)  # (B, H, T, D_head)
    if precomputed_query_rope is not None:
        cos, sin = precomputed_query_rope
        queries = RotaryPositionalEncoding.apply_rotation_half(queries, sin, cos)
    attention_output = self.attention.compute_attention(
        queries=queries,
        keys=keys,
        values=values,
        attention_mask=attention_mask,
    )
    hidden_states = self.apply_residual(residual, attention_output, gate)
    return hidden_states