Project: s41467-025-60872-5 | Function: dot_product_attention (python)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Implements scaled dot-product attention mechanism for transformer architecture. Computes attention weights using query-key similarity, applies softmax normalization, and produces weighted value representations. Includes support for attention masks and key padding masks for variable-length sequences.

**KEY TERMS:** {
  "q": "query matrix Q",
  "k": "key matrix K",
  "v": "value matrix V",
  "attn": "attention weights",
  "c": "dimension d_k",
  "math.sqrt(c)": "scaling factor √d_k"
}

**CONTEXT CODE:** import torch
from torch import nn

import math

from rinalmo.model.rope import RotaryPositionEmbedding

from flash_attn import flash_attn_varlen_qkvpacked_func, flash_attn_qkvpacked_func
from flash_attn.layers.rotary import RotaryEmbedding

from flash_attn.bert_padding import unpad_input, pad_input

from einops import rearrange

def dot_product_attention(q, k, v, attn_mask=None, key_pad_mask=None, dropout=None):
    # MASKED: Implements scaled dot-product attention mechanism for transformer architecture with query-key similarity computation, softmax normalization, and weighted value representation including attention masks and key padding masks implementation
    raise NotImplementedError("This function is masked for testing purposes.")

class MultiHeadAttention(nn.Module):
    def __init__(self, c_in, num_heads, attention_dropout=0.0, use_rot_emb=True, bias=False):
        super().__init__()
        assert c_in % num_heads == 0, "Embedding dimensionality must be divisible with number of attention heads!"

        self.c_in = c_in
        self.num_heads = num_heads

        self.c_head = c_in // self.num_heads
        self.c_qkv = self.c_head * num_heads

        self.use_rot_emb = use_rot_emb
        if self.use_rot_emb:
            self.rotary_emb = RotaryPositionEmbedding(self.c_head)

        self.to_q = nn.Linear(self.c_in, self.c_qkv, bias=bias)
        self.to_k = nn.Linear(self.c_in, self.c_qkv, bias=bias)
        self.to_v = nn.Linear(self.c_in, self.c_qkv, bias=bias)

        self.attention_dropout = nn.Dropout(p=attention_dropout)

        self.out_proj = nn.Linear(c_in, c_in, bias=bias)

    def forward(self, q, k, v, attn_mask=None, key_pad_mask=None):
        bs = q.shape[0]

        q = self.to_q(q).view(bs, -1, self.num_heads, self.c_head).transpose(-2, -3)
        k = self.to_k(k).view(bs, -1, self.num_heads, self.c_head).transpose(-2, -3)
        v = self.to_v(v).view(bs, -1, self.num_heads, self.c_head).transpose(-2, -3)

        if self.use_rot_emb:
            q, k = self.rotary_emb(q, k)

        output, attn = dot_product_attention(q, k, v, attn_mask, key_pad_mask, self.attention_dropout)

        output = output.transpose(-2, -3).contiguous().view(bs, -1, self.num_heads * self.c_head)
        output = self.out_proj(output)

        return output, attn

class MultiHeadSelfAttention(nn.Module):
    def __init__(self, c_in, num_heads, attention_dropout=0.0, use_rot_emb=True, bias=False):
        super().__init__()

        self.mh_attn = MultiHeadAttention(c_in, num_heads, attention_dropout, use_rot_emb, bias)

    def forward(self, x, attn_mask=None, key_pad_mask=None):
        return self.mh_attn(x, x, x, attn_mask, key_pad_mask)

class FlashAttention(nn.Module):
    """
    Implement the scaled dot product attention with softmax.
    """
    def __init__(self, causal=False, softmax_scale=None, attention_dropout=0.0):
        """
        Args:
            causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
            softmax_scale: float. The scaling of QK^T before applying softmax.
                           Default to 1 / sqrt(headdim).
            attention_dropout: float. The dropout rate to apply to the attention
                               (default: 0.0)
        """
        super().__init__()
        self.softmax_scale = softmax_scale
        self.attention_dropout = attention_dropout
        self.causal = causal

    def forward(self, qkv, cu_seqlens=None, max_seqlen=None, return_attn_probs=False):
        """
        Arguments
        ---------
            qkv: The tensor containing the query, key, and value.
                 If cu_seqlens is None and max_seqlen is None, then qkv has shape (B, S, 3, H, D).
                 If cu_seqlens is not None and max_seqlen is not None, then qkv has shape
                 (total, 3, H, D), where total is the sum of the sequence lengths in the batch.
            cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
                        of the sequences in the batch, used to index into qkv.
            max_seqlen: int. Maximum sequence length in the batch.
            return_attn_probs: bool. Whether to return the attention probabilities. This option is for
                               testing only. The returned probabilities are not guaranteed to be correct
                               (they might not have the right scaling).
        Returns:
        --------
            out: (total, H, D) if cu_seqlens is not None and max_seqlen is not None,
                 else (B, S, H, D).
        """
        assert qkv.dtype in [torch.float16, torch.bfloat16]
        assert qkv.is_cuda

        unpadded = cu_seqlens is not None

        if unpadded:
            assert cu_seqlens.dtype == torch.int32
            assert max_seqlen is not None
            assert isinstance(max_seqlen, int)
            return flash_attn_varlen_qkvpacked_func(


Complete the function implementation replacing "raise NotImplementedError".