flash attention的调用

### 调用Flash Attention实现

在现代深度学习框架中，调用Flash Attention已经变得更为简便。对于PyTorch 2及以上版本，可以直接利用内置函数`torch.nn.functional.scaled_dot_product_attention()`来应用缩放点积注意力机制[^2]。

下面是一个具体的Python代码片段展示如何在一个简单的多头自注意层中集成并使用Flash Attention：

import torch
from torch import nn

class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim=512, num_heads=8, dropout=0.1, bias=True):
        super(MultiHeadAttention, self).__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        
        self.qkv_proj = nn.Linear(embed_dim, 3 * embed_dim, bias=bias)
        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)

        self.dropout_p = dropout
    
    def _shape(self, tensor, seq_len, batch_size):
        return tensor.view(batch_size, seq_len, self.num_heads, -1).transpose(1, 2).contiguous()

    def forward(self, query, key, value, attn_mask=None, position_ids=None):    
        batch_size, tgt_len, src_len = query.size(0), query.size(1), key.size(1)

        qkv_same = torch.equal(query, key) and torch.equal(key, value)
        
        if qkv_same:
            # Self-attention
            qkv = self.qkv_proj(query).chunk(3, dim=-1)
            
            q, k, v = map(lambda t: self._shape(t, src_len, batch_size), qkv)
        else:
            raise NotImplementedError("Cross-attention not implemented")

        if position_ids is not None:
            # Pass additional positional information to Flash Attention implementation.
            pass
            
        output = torch.nn.functional.scaled_dot_product_attention(
            q, k, v,
            attn_mask=attn_mask,
            dropout_p=self.dropout_p if self.training else 0.,
            is_causal=False
        )
        
        output = output.transpose(1, 2).reshape(batch_size, tgt_len, self.embed_dim)
        output = self.out_proj(output)
        
        return output

此段代码定义了一个名为`MultiHeadAttention`的类，它继承自`nn.Module`。该模块实现了带有Flash Attention特性的多头自我关注机制。特别值得注意的是，在前向传播过程中，通过调用`scaled_dot_product_attention`方法执行实际的关注计算，并允许传递额外的位置ID给Flash Attention组件以增强其性能表现[^3]。