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可以最大化推理性能。

常见问题解答

Flash Attention能快多少？

相比标准Attention快2-4倍，显存占用减少5-20倍。长序列(4K+)提升更明显。

哪些GPU支持Flash Attention？

需要Ampere架构及以上：A100、A10、RTX 3090/4090等。V100不支持。

怎么确认Flash Attention已启用？

用torch.backends.cuda.sdp_kernel()上下文管理器强制指定，或检查日志中的kernel选择信息。

Flash Attention v1和v2区别？

v2优化了并行度，减少了非矩阵运算，比v1快约2倍。PyTorch 2.1+默认使用v2。

## 相关文章推荐 - [AI在电商推荐中的应用](https://www.52runoob.com/2026/07/16/ai%e5%9c%a8%e7%94%b5%e5%95%86%e6%8e%a8%e8%8d%90%e4%b8%ad%e7%9a%84%e5%ba%94%e7%94%a8/) - [AI Agent智能体开发入门](https://www.52runoob.com/2026/05/25/ai-agent-1909/) - [RAG重排序Rerank怎么提升准确率](https://www.52runoob.com/2026/08/24/rag%e9%87%8d%e6%8e%92%e5%ba%8frerank%e6%80%8e%e4%b9%88%e6%8f%90%e5%8d%87%e5%87%86%e7%a1%ae%e7%8e%87/)

核心要点

本文详细介绍了AI模型推理优化Flash Attention原理的核心概念和实际应用。作为AI / 智能开发领域的重要主题，掌握AI模型推理优化Flash Attention原理对于提升开发效率和技术能力具有重要意义。

关键知识点

  • 基础概念：理解AI模型推理优化Flash Attention原理的基本原理和核心机制
  • 实践应用：AI模型推理优化Flash Attention原理在实际项目中的具体使用方法
  • 最佳实践：避免常见陷阱，掌握AI模型推理优化Flash Attention原理的推荐做法
  • 进阶技巧：深入了解AI模型推理优化Flash Attention原理的高级特性和优化策略

学习建议

建议在学习AI模型推理优化Flash Attention原理时，结合实际项目进行练习。可以从简单的示例开始，逐步深入理解其内部原理。同时，关注社区最新动态和最佳实践，确保掌握最新技术进展。

📚 相关文章推荐