AI模型推理优化Flash Attention原理

Flash Attention是LLM推理加速的核心技术之一，理解了它就知道为什么大模型推理越来越快。把学习笔记整理一下。

先搞清楚原理

Flash Attention的核心思想是通过分块计算减少GPU显存访问。传统Attention需要把完整的注意力矩阵存到HBM，而Flash Attention在SRAM中分块计算，避免频繁读写HBM。看下面这段代码：

import torch
import torch.nn.functional as F
import math

# 传统Attention
def standard_attention(Q, K, V):
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(K.size(-1))
    attn = F.softmax(scores, dim=-1)
    return torch.matmul(attn, V)

# Flash Attention (PyTorch 2.0+)
batch, heads, seq_len, dim = 2, 8, 4096, 64
Q = torch.randn(batch, heads, seq_len, dim, device='cuda')
K = torch.randn(batch, heads, seq_len, dim, device='cuda')
V = torch.randn(batch, heads, seq_len, dim, device='cuda')

output = F.scaled_dot_product_attention(Q, K, V)
print(f'输出形状: {output.shape}')

实际怎么用

实际项目中直接用PyTorch 2.0+的SDPA接口，底层自动启用Flash Attention：

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, n_heads=8):
        super().__init__()
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x):
        batch, seq_len, _ = x.shape
        Q = self.W_q(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        attn = F.scaled_dot_product_attention(Q, K, V, is_causal=True)
        attn = attn.transpose(1, 2).contiguous().view(batch, seq_len, -1)
        return self.W_o(attn)

# KV Cache加速
class KVCache:
    def __init__(self, batch, heads, max_len, d_k):
        self.k = torch.zeros(batch, heads, max_len, d_k, device='cuda')
        self.v = torch.zeros(batch, heads, max_len, d_k, device='cuda')
        self.pos = 0
    def update(self, k, v):
        n = k.size(2)
        self.k[:,:,self.pos:self.pos+n] = k
        self.v[:,:,self.pos:self.pos+n] = v
        self.pos += n
        return self.k[:,:,:self.pos], self.v[:,:,:self.pos]

进阶一点的用法

用torch.compile进一步优化：

import torch
import torch.nn.functional as F

@torch.compile(mode='max-autotune')
def compiled_attention(Q, K, V):
    return F.scaled_dot_product_attention(Q, K, V, is_causal=True)

# 性能对比
import time
Q = torch.randn(2, 32, 8192, 128, device='cuda')
K = torch.randn(2, 32, 8192, 128, device='cuda')
V = torch.randn(2, 32, 8192, 128, device='cuda')

# warmup
for _ in range(3): compiled_attention(Q, K, V)
torch.cuda.synchronize()

start = time.time()
for _ in range(10): compiled_attention(Q, K, V)
torch.cuda.synchronize()
print(f'Flash Attention: {(time.time()-start)/10*1000:.1f}ms')

容易踩的坑

第一个是PyTorch版本太低不支持SDPA，需要2.0+。第二个是GPU架构不支持Flash Attention，需要Ampere及以上。第三个是mask和is_causal不能同时指定。

写在最后

Flash Attention的核心是分块计算减少显存访问。实际开发中直接用PyTorch SDPA接口，配合torch.compile和KV Cache可以最大化推理性能。