src/models/attention.py

"""Attention building blocks used by RiboSphere."""

from __future__ import annotations

from typing import Any

import torch
from torch import Tensor, nn
import torch.nn.functional as F
from torch.nn.attention.flex_attention import create_block_mask, flex_attention

from .layers import FeedForward
from .rotary import RotaryEmbedding

AttentionArguments = dict[str, Any]


def root_mean_square_norm(tensor: Tensor) -> Tensor:
    """Apply parameter-free RMS normalization over the final dimension."""

    return F.rms_norm(tensor, (tensor.shape[-1],))


class TransformerStack(nn.Module):
    """Stack of local self-attention blocks."""

    def __init__(
        self,
        *,
        num_channels: int,
        num_heads: int,
        mlp_factor: int,
        window_size: int,
        num_layers: int,
        attention_backend: str = "flex",
        dropout: float = 0.1,
        pairwise_channels: int = 0,
        is_causal: bool = False,
    ) -> None:
        super().__init__()
        if num_channels <= 0 or num_heads <= 0 or num_layers <= 0:
            raise ValueError(
                "num_channels, num_heads, and num_layers must be positive."
            )
        if num_channels % num_heads != 0:
            raise ValueError("num_channels must be divisible by num_heads.")
        if window_size <= 0:
            raise ValueError("window_size must be positive.")
        if pairwise_channels < 0:
            raise ValueError("pairwise_channels cannot be negative.")

        attention_backend = attention_backend.lower()
        if attention_backend not in {"sdpa", "flex"}:
            raise ValueError("attention_backend must be 'sdpa' or 'flex'.")

        use_pair_bias = pairwise_channels > 0
        self.blocks = nn.ModuleList(
            [
                TransformerBlock(
                    num_channels=num_channels,
                    num_heads=num_heads,
                    mlp_factor=mlp_factor,
                    attention_backend=attention_backend,
                    dropout=dropout,
                    use_pairwise_bias=use_pair_bias,
                    pairwise_channels=pairwise_channels,
                )
                for _ in range(num_layers)
            ]
        )
        self.window_size = window_size
        self.is_causal = is_causal
        self.attention_backend = attention_backend

    def _window_mask(
        self,
        batch_index: Tensor,
        head_index: Tensor,
        query_index: Tensor,
        key_value_index: Tensor,
    ) -> Tensor:
        del batch_index, head_index
        within_window = (query_index - key_value_index).abs() <= self.window_size
        if self.is_causal:
            within_window = within_window & (query_index >= key_value_index)
        return within_window

    def forward(
        self,
        hidden_states: Tensor,
        pairwise_features: Tensor | None = None,
    ) -> Tensor:
        """Transform ``[B, L, D]`` token features."""

        if hidden_states.ndim != 3:
            raise ValueError("hidden_states must have shape [B, L, D].")

        sequence_length = hidden_states.shape[1]
        if pairwise_features is not None and pairwise_features.shape[:3] != (
            hidden_states.shape[0],
            sequence_length,
            sequence_length,
        ):
            raise ValueError(
                "pairwise_features must have shape [B, L, L, P]."
            )

        if self.attention_backend == "flex":
            attention_arguments: AttentionArguments = {
                "block_mask": create_block_mask(
                    self._window_mask,
                    B=None,
                    H=None,
                    Q_LEN=sequence_length,
                    KV_LEN=sequence_length,
                    device=hidden_states.device,
                ),
                "score_mod": None,
            }
        else:
            positions = torch.arange(
                sequence_length,
                device=hidden_states.device,
            )
            attention_mask = (
                positions[:, None] - positions[None, :]
            ).abs() <= self.window_size
            if self.is_causal:
                attention_mask = attention_mask & (
                    positions[:, None] >= positions[None, :]
                )
            attention_arguments = {"attn_mask": attention_mask.unsqueeze(0)}

        for block in self.blocks:
            hidden_states = block(
                hidden_states,
                pairwise_features=pairwise_features,
                **attention_arguments,
            )
        return hidden_states


class TransformerBlock(nn.Module):
    """Pre-normalized self-attention and feed-forward block."""

    def __init__(
        self,
        *,
        num_channels: int,
        num_heads: int,
        mlp_factor: int,
        attention_backend: str = "flex",
        dropout: float = 0.1,
        use_pairwise_bias: bool = False,
        pairwise_channels: int = 0,
    ) -> None:
        super().__init__()
        self.attention_backend = attention_backend
        self.attention = SelfAttention(
            model_dimension=num_channels,
            num_heads=num_heads,
            dropout=dropout,
            attention_backend=attention_backend,
        )
        self.feed_forward = FeedForward(
            num_channels,
            num_channels * mlp_factor,
            num_channels,
            activation=nn.GELU,
            dropout=dropout,
        )

        self.use_pairwise_bias = use_pairwise_bias
        if use_pairwise_bias:
            if pairwise_channels <= 0:
                raise ValueError(
                    "pairwise_channels must be positive when pair bias is enabled."
                )
            self.pair_bias_projection = nn.Linear(
                pairwise_channels, 1, bias=True
            )
            self.pair_bias_norm = nn.LayerNorm(pairwise_channels)
        else:
            self.pair_bias_projection = None
            self.pair_bias_norm = None

    def _add_pair_bias(
        self,
        pairwise_features: Tensor,
        attention_arguments: AttentionArguments,
    ) -> AttentionArguments:
        if self.pair_bias_projection is None or self.pair_bias_norm is None:
            return attention_arguments

        pair_bias = self.pair_bias_projection(
            self.pair_bias_norm(pairwise_features)
        ).squeeze(-1)
        attention_arguments = dict(attention_arguments)

        if self.attention_backend == "flex":

            def pair_biased_score(
                score: Tensor,
                batch_index: Tensor,
                head_index: Tensor,
                query_index: Tensor,
                key_value_index: Tensor,
            ) -> Tensor:
                del head_index
                return score + pair_bias[
                    batch_index,
                    query_index,
                    key_value_index,
                ]

            attention_arguments["score_mod"] = pair_biased_score
        else:
            attention_mask = attention_arguments["attn_mask"]
            additive_pair_bias = torch.where(
                attention_mask,
                pair_bias,
                torch.full_like(pair_bias, -torch.inf),
            )
            attention_arguments["attn_mask"] = additive_pair_bias.unsqueeze(1)
        return attention_arguments

    def forward(
        self,
        hidden_states: Tensor,
        *,
        pairwise_features: Tensor | None = None,
        **attention_arguments: Any,
    ) -> Tensor:
        if self.use_pairwise_bias:
            if pairwise_features is None:
                raise ValueError(
                    "pairwise_features are required when pair bias is enabled."
                )
            attention_arguments = self._add_pair_bias(
                pairwise_features,
                attention_arguments,
            )

        hidden_states = hidden_states + self.attention(
            root_mean_square_norm(hidden_states),
            **attention_arguments,
        )
        hidden_states = hidden_states + self.feed_forward(
            root_mean_square_norm(hidden_states)
        )
        return hidden_states


class SelfAttention(nn.Module):
    """Multi-head self-attention with rotary position embeddings."""

    def __init__(
        self,
        model_dimension: int,
        num_heads: int,
        *,
        normalize_queries_and_keys: bool = False,
        attention_backend: str = "flex",
        dropout: float = 0.1,
    ) -> None:
        super().__init__()
        if model_dimension <= 0 or num_heads <= 0:
            raise ValueError("model_dimension and num_heads must be positive.")
        if model_dimension % num_heads != 0:
            raise ValueError("model_dimension must be divisible by num_heads.")
        if not 0.0 <= dropout < 1.0:
            raise ValueError("dropout must be in [0, 1).")

        attention_backend = attention_backend.lower()
        if attention_backend not in {"flex", "sdpa"}:
            raise ValueError("backend must be 'flex' or 'sdpa'.")

        self.model_dimension = model_dimension
        self.num_heads = num_heads
        self.head_dimension = self.model_dimension // self.num_heads
        self.attention_dropout = nn.Dropout(dropout)
        self.dropout = dropout
        self.normalize_queries_and_keys = normalize_queries_and_keys

        self.rotary_embedding = RotaryEmbedding(self.head_dimension)
        self.qkv_projection = nn.Linear(
            model_dimension, 3 * model_dimension, bias=True
        )
        self.output_projection = nn.Linear(model_dimension, model_dimension)
        self.residual_dropout = nn.Dropout(dropout)
        self.attention_backend = attention_backend

    def forward(
        self,
        hidden_states: Tensor,
        **attention_arguments: Any,
    ) -> Tensor:
        """Apply self-attention to ``[B, L, D]`` hidden states."""

        if hidden_states.ndim != 3:
            raise ValueError("hidden_states must have shape [B, L, D].")
        batch_size, sequence_length, hidden_dimension = hidden_states.shape
        if hidden_dimension != self.model_dimension:
            raise ValueError(
                f"Expected hidden dimension {self.model_dimension}, "
                f"received {hidden_dimension}."
            )

        query, key, value = self.qkv_projection(hidden_states).split(
            self.model_dimension, dim=-1
        )

        def split_heads(tensor: Tensor) -> Tensor:
            return tensor.reshape(
                batch_size,
                sequence_length,
                self.num_heads,
                self.head_dimension,
            ).transpose(1, 2)

        query, key, value = map(split_heads, (query, key, value))
        if self.normalize_queries_and_keys:
            query = root_mean_square_norm(query)
            key = root_mean_square_norm(key)
        query, key = self.rotary_embedding(query, key)

        if self.attention_backend == "flex":
            attention_output = flex_attention(
                query,
                key,
                value,
                block_mask=attention_arguments.get("block_mask"),
                score_mod=attention_arguments.get("score_mod"),
            )
        else:
            attention_output = F.scaled_dot_product_attention(
                query,
                key,
                value,
                **attention_arguments,
            )

        attention_output = self.attention_dropout(attention_output)
        attention_output = attention_output.transpose(1, 2).contiguous().view(
            batch_size,
            sequence_length,
            self.model_dimension,
        )
        attention_output = self.residual_dropout(
            self.output_projection(attention_output)
        )
        return attention_output