Attention 进阶优化：Flash Attention 与 Paged KV Cache 深度解析

从 O(N²) 到 O(N)：如何让大模型推理更快、更省内存？
本文深入讲解 Flash Attention 和 Paged KV Cache 两大核心优化技术，带你理解 vLLM、TensorRT-LLM 等高性能推理引擎的底层原理。不仅有算法讲解，更有完整的 Python 实现！

💡开源代码工程：https://github.com/rixin2025/attention-from-scratch/tree/main

*** 欢迎 star 和讨论 ***

📑 目录

• 📖 前言
• 🎯 为什么需要进阶优化
• 🚀 快速开始
• 📚 核心内容详解

  • 1. Flash Attention：内存高效的 Attention
  • 2. Paged KV Cache：分页管理的内存优化
• 📊 性能对比总结
• 🎓 学习路径
• 🔧 实际应用场景
• 💡 核心代码片段
• 🎯 项目特色
• 📈 学习成果
• 📚 参考资料
• 📝 总结

📖 前言

在基础篇中，我们实现了 Attention 的核心机制：

• ✅ Scaled Dot-Product Attention
• ✅ Multi-Head Attention (MHA)
• ✅ Grouped Query Attention (GQA)
• ✅ KV Cache

但在实际的大模型推理中，我们还面临两个关键挑战：

挑战 1：内存瓶颈

标准 Attention 需要存储完整的注意力矩阵 S = Q @ K^T：

• 序列长度 32K，32 个头 → 注意力矩阵约 256GB（FP16）
• 这在 GPU 上根本无法存储！

挑战 2：内存碎片

连续的 KV Cache 导致严重的内存碎片：

• 不同序列长度不同，需要不同大小的连续内存
• 删除序列后，内存无法有效重用
• 内存利用率低，浪费严重

本文将介绍两大解决方案：

1. Flash Attention：通过分块计算，将内存从 O(N²) 降到 O(N)
2. Paged KV Cache：通过分页管理，提高内存利用率，减少碎片

这两项技术是 vLLM、TensorRT-LLM 等高性能推理引擎的核心！

🎯 为什么需要进阶优化

标准 Attention 的问题

让我们用具体数字来看问题的严重性：

场景：LLaMA 2 70B 模型，处理 32K 上下文

• 模型配置：80 层，64 个 Q 头，8 个 KV 头，头维度 128
• 序列长度：32,768 tokens
• Batch size：32

内存占用计算：

1. 注意力矩阵（标准 Attention）：

   每层每个头：32768 × 32768 × 2 bytes (FP16) = 2GB
   64 个头：2GB × 64 = 128GB
   80 层：128GB × 80 = 10,240GB = 10TB！
   

2. KV Cache（连续存储）：

   每层：32 × 32768 × 8 × 128 × 2 × 2 (K+V) = 512MB
   80 层：512MB × 80 = 40GB
   

问题显而易见：

• ❌ 10TB 的注意力矩阵无法存储
• ❌ 40GB 的 KV Cache 在多序列场景下碎片严重
• ❌ 内存利用率低，推理速度慢

进阶优化的效果

使用 Flash Attention + Paged KV Cache：

1. Flash Attention：

   • 注意力矩阵内存：10TB → 0（不存储完整矩阵）
   • 只需要 O(N) 的临时缓冲区

2. Paged KV Cache：

   • 内存利用率：60% → 95%
   • 支持更大的 batch size
   • 动态内存管理，无碎片

这就是为什么 vLLM 能比 HuggingFace Transformers 快 24x！

🚀 快速开始

环境准备

# 克隆项目
git clone https://github.com/rixin2025/attention-from-scratch.git
cd attention-from-scratch

# 安装依赖
pip install -r requirements.txt

运行示例

import torch
from src.flash_attention import FlashAttention
from src.paged_kv_cache import PagedKVCache

# Flash Attention
flash_attn = FlashAttention(
    d_model=512, 
    num_heads=8, 
    block_size=64  # 分块大小
)
x = torch.randn(2, 100, 512)
output, lse = flash_attn(x)
print(f"输出形状: {output.shape}")  # [2, 100, 512]

# Paged KV Cache
paged_cache = PagedKVCache(
    num_heads=8,
    head_dim=64,
    page_size=16,  # 每页 16 个 tokens
    num_pages=1024
)

# 分配序列
seq_id = paged_cache.allocate_sequence(seq_len=100)
print(f"序列 {seq_id} 分配了 {len(paged_cache.page_table[seq_id])} 个页面")

# 更新缓存
k = torch.randn(1, 8, 100, 64)
v = torch.randn(1, 8, 100, 64)
paged_cache.update(seq_id, k, v, start_pos=0)

运行 Notebooks

# 启动 Jupyter
jupyter notebook notebooks/

# 按顺序学习进阶内容
# 05_flash_attention.ipynb - Flash Attention 详解
# 06_paged_kv_cache.ipynb - Paged KV Cache 详解

📚 核心内容详解

1. Flash Attention：内存高效的 Attention

1.1 核心问题：O(N²) 的内存瓶颈

标准 Attention 的计算流程：

# 步骤 1: 计算注意力分数矩阵
S = Q @ K.T / sqrt(d_k)  # [batch, heads, N, N] ← 需要存储！

# 步骤 2: Softmax 归一化
P = softmax(S, dim=-1)   # [batch, heads, N, N] ← 需要存储！

# 步骤 3: 加权求和
O = P @ V                # [batch, heads, N, d_v]

内存占用分析：

对于序列长度 N = 32,768：

• S 矩阵：32768 × 32768 × 2 bytes (FP16) = 2GB
• P 矩阵：32768 × 32768 × 2 bytes (FP16) = 2GB
• 总计： 4GB per head per layer

对于 LLaMA 2 70B（64 头，80 层）：

• 单个样本：4GB × 64 × 80 = 20TB
• 这在任何 GPU 上都无法存储！

1.2 Flash Attention 的解决方案

核心思想：分块计算（Tiling），不存储完整的 N×N 矩阵

三大技术：

1. Tiling（分块）：

   • 将 Q 分成块：Q₁, Q₂, ..., Qₘ
   • 将 K, V 分成块：K₁, K₂, ..., Kₙ 和 V₁, V₂, ..., Vₙ
   • 逐块计算，只存储小块矩阵

2. Online Softmax（在线 Softmax）：

   • 增量计算 softmax，避免两次遍历
   • 动态更新最大值和归一化因子
   • 支持跨块的 softmax 合并

3. Recomputation（重计算）：

   • 前向传播不存储中间结果
   • 反向传播时重新计算
   • 用计算换内存

算法流程：

输入: Q [N, d], K [N, d], V [N, d]
分块大小: B_q, B_kv

初始化:
  O = zeros(N, d)          # 输出
  l = zeros(N)             # softmax 归一化因子
  m = -inf * ones(N)       # softmax 最大值

外层循环 (遍历 Q 的块):
  for i = 1 to ceil(N / B_q):
    Q_i = Q[i*B_q : (i+1)*B_q]  # 加载 Q 的第 i 块
    O_i = zeros(B_q, d)
    l_i = zeros(B_q)
    m_i = -inf * ones(B_q)
    
    内层循环 (遍历 K, V 的块):
      for j = 1 to ceil(N / B_kv):
        K_j = K[j*B_kv : (j+1)*B_kv]  # 加载 K 的第 j 块
        V_j = V[j*B_kv : (j+1)*B_kv]  # 加载 V 的第 j 块
        
        # 计算注意力分数（小块）
        S_ij = Q_i @ K_j.T / sqrt(d_k)  # [B_q, B_kv]
        
        # 在线 Softmax 更新
        m_i_new = max(m_i, rowmax(S_ij))
        P_ij = exp(S_ij - m_i_new)
        l_i_new = exp(m_i - m_i_new) * l_i + rowsum(P_ij)
        
        # 更新输出
        O_i = exp(m_i - m_i_new) * O_i + P_ij @ V_j
        
        # 更新状态
        m_i = m_i_new
        l_i = l_i_new
    
    # 归一化
    O_i = O_i / l_i
    O[i*B_q : (i+1)*B_q] = O_i

返回: O

1.3 Online Softmax 算法详解

标准 Softmax 的问题：

# 需要两次遍历
def standard_softmax(x):
    # 第一次遍历：找最大值（数值稳定性）
    m = max(x)
    
    # 第二次遍历：计算 exp 和 sum
    exp_x = exp(x - m)
    s = sum(exp_x)
    
    return exp_x / s

Online Softmax 的解决方案：

一次遍历完成，并支持增量更新！

def online_softmax_update(old_max, old_sum, new_values):
    """
    增量更新 softmax
    
    Args:
        old_max: 旧的最大值
        old_sum: 旧的 exp 求和
        new_values: 新的值
    
    Returns:
        new_max: 更新后的最大值
        new_sum: 更新后的 exp 求和
    """
    # 计算新值的最大值
    new_max_local = max(new_values)
    
    # 全局最大值
    new_max = max(old_max, new_max_local)
    
    # 更新旧的 sum（重新归一化）
    old_sum_corrected = old_sum * exp(old_max - new_max)
    
    # 计算新值的 sum
    new_sum_local = sum(exp(new_values - new_max))
    
    # 合并
    new_sum = old_sum_corrected + new_sum_local
    
    return new_max, new_sum

数学原理：

假设我们已经计算了前 k 个值的 softmax：

m_k = max(x_1, ..., x_k)
l_k = sum(exp(x_i - m_k) for i in 1..k)

现在加入新值 x_{k+1}：

m_{k+1} = max(m_k, x_{k+1})
l_{k+1} = exp(m_k - m_{k+1}) * l_k + exp(x_{k+1} - m_{k+1})

这样就可以增量更新，无需重新计算所有值！

1.4 Python 实现

import torch
import torch.nn as nn
import math

def flash_attention_forward(
    Q: torch.Tensor,  # [batch, num_heads, seq_len, head_dim]
    K: torch.Tensor,
    V: torch.Tensor,
    block_size: int = 64
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Flash Attention 前向传播
    
    Returns:
        output: [batch, num_heads, seq_len, head_dim]
        lse: log-sum-exp，用于反向传播
    """
    batch, num_heads, seq_len, head_dim = Q.shape
    scale = 1.0 / math.sqrt(head_dim)
    
    # 初始化输出
    O = torch.zeros_like(Q)
    l = torch.zeros(batch, num_heads, seq_len, 1, device=Q.device)
    m = torch.full((batch, num_heads, seq_len, 1), float('-inf'), device=Q.device)
    
    # 外层循环：遍历 Q 的块
    num_q_blocks = (seq_len + block_size - 1) // block_size
    num_kv_blocks = (seq_len + block_size - 1) // block_size
    
    for i in range(num_q_blocks):
        q_start = i * block_size
        q_end = min((i + 1) * block_size, seq_len)
        Q_i = Q[:, :, q_start:q_end, :]  # [batch, heads, B_q, d]
        
        O_i = torch.zeros_like(Q_i)
        l_i = torch.zeros(batch, num_heads, q_end - q_start, 1, device=Q.device)
        m_i = torch.full((batch, num_heads, q_end - q_start, 1), float('-inf'), device=Q.device)
        
        # 内层循环：遍历 K, V 的块
        for j in range(num_kv_blocks):
            kv_start = j * block_size
            kv_end = min((j + 1) * block_size, seq_len)
            K_j = K[:, :, kv_start:kv_end, :]  # [batch, heads, B_kv, d]
            V_j = V[:, :, kv_start:kv_end, :]
            
            # 计算注意力分数（小块）
            S_ij = torch.matmul(Q_i, K_j.transpose(-2, -1)) * scale  # [batch, heads, B_q, B_kv]
            
            # Online Softmax 更新
            m_ij = S_ij.max(dim=-1, keepdim=True)[0]  # [batch, heads, B_q, 1]
            m_i_new = torch.maximum(m_i, m_ij)
            
            # 计算 exp 和更新
            P_ij = torch.exp(S_ij - m_i_new)  # [batch, heads, B_q, B_kv]
            l_i_new = torch.exp(m_i - m_i_new) * l_i + P_ij.sum(dim=-1, keepdim=True)
            
            # 更新输出
            O_i = torch.exp(m_i - m_i_new) * O_i + torch.matmul(P_ij, V_j)
            
            # 更新状态
            m_i = m_i_new
            l_i = l_i_new
        
        # 归一化
        O_i = O_i / l_i
        O[:, :, q_start:q_end, :] = O_i
        m[:, :, q_start:q_end, :] = m_i
        l[:, :, q_start:q_end, :] = l_i
    
    # 计算 log-sum-exp（用于反向传播）
    lse = m + torch.log(l)
    
    return O, lse


class FlashAttention(nn.Module):
    """Flash Attention 模块"""
    
    def __init__(self, d_model: int, num_heads: int, block_size: int = 64):
        super().__init__()
        assert d_model % num_heads == 0
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        self.block_size = block_size
        
        # Q, K, V 投影
        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: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            x: [batch, seq_len, d_model]
        
        Returns:
            output: [batch, seq_len, d_model]
            lse: log-sum-exp
        """
        batch, seq_len, _ = x.shape
        
        # 投影并分割成多头
        Q = self.W_q(x).view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        K = self.W_k(x).view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        V = self.W_v(x).view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        
        # Flash Attention
        attn_output, lse = flash_attention_forward(Q, K, V, self.block_size)
        
        # 合并多头
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch, seq_len, self.d_model)
        
        # 输出投影
        output = self.W_o(attn_output)
        
        return output, lse

1.5 性能分析

内存复杂度对比：

实际测试（batch=1, heads=32, d=128, FP16）：

速度对比（A100 GPU）：

2. Paged KV Cache：分页管理的内存优化

2.1 核心问题：连续 KV Cache 的内存碎片

在基础篇中，我们实现了 KV Cache 来加速推理。但连续的 KV Cache 存在严重问题：

连续 KV Cache 的内存布局：

序列1 (100 tokens): [████████████████████████████] 连续内存
序列2 (50 tokens):  [██████████████] 连续内存
序列3 (200 tokens): [████████████████████████████████████████████████] 连续内存

问题：

1. 预分配浪费：

   • 必须预分配最大长度的内存（如 2048 tokens）
   • 实际使用可能只有 100 tokens
   • 浪费率：(2048 - 100) / 2048 = 95%

2. 内存碎片：

   初始状态: [序列1][序列2][序列3]
   删除序列2: [序列1][空闲][序列3]  ← 产生碎片
   新序列4 (150 tokens): 无法使用序列2的空间（只有50 tokens）
   

3. 无法动态扩展：

   • 序列长度超过预分配大小时，需要重新分配
   • 重新分配需要复制数据，开销大

实际影响（batch=32, max_len=2048, 实际平均长度=200）：

理论需要: 32 × 200 × head_dim × 2 (K+V) = 12.8 MB
实际分配: 32 × 2048 × head_dim × 2 (K+V) = 131 MB
内存利用率: 12.8 / 131 = 9.8%  ← 浪费 90%！

2.2 Paged KV Cache 的解决方案

核心思想：借鉴操作系统的虚拟内存管理，将 KV Cache 分成固定大小的页面（pages）

三大组件：

1. 页面池（Page Pool）：

   • 全局的页面池，包含固定大小的页面
   • 页面大小通常为 16、32 或 64 tokens
   • 所有序列共享页面池

2. 页面表（Page Table）：

   • 记录每个序列使用的页面
   • 逻辑地址 → 物理页面的映射
   • 支持非连续的物理内存

3. 空闲列表（Free List）：

   • 管理未使用的页面
   • 支持动态分配和回收
   • 实现内存复用

内存布局：

页面池（全局）:
[Page 0: 16 tokens][Page 1: 16 tokens][Page 2: 16 tokens]...

页面表:
序列1 (50 tokens):  [0, 1, 2]      ← 使用 3 个页面
序列2 (30 tokens):  [3, 4]         ← 使用 2 个页面
序列3 (100 tokens): [5, 6, 7, 8, 9, 10] ← 使用 6 个页面

空闲列表: [11, 12, 13, ...]

优势：

1. 按需分配：

   • 只分配实际需要的页面
   • 50 tokens → 4 个页面（16×4=64）
   • 浪费：(64-50)/64 = 22%（vs 连续的 95%）

2. 无碎片：

   删除序列2: 页面 3, 4 回到空闲列表
   新序列4 (150 tokens): 可以使用任意 10 个空闲页面
   

3. 动态扩展：

   • 序列增长时，只需分配新页面
   • 无需复制已有数据

2.3 页面管理算法

页面分配：

def allocate_sequence(seq_len: int, page_size: int = 16) -> list[int]:
    """
    为序列分配页面
    
    Args:
        seq_len: 序列长度
        page_size: 页面大小
    
    Returns:
        page_ids: 分配的页面 ID 列表
    """
    # 计算需要的页面数
    num_pages = (seq_len + page_size - 1) // page_size
    
    # 从空闲列表分配页面
    page_ids = []
    for _ in range(num_pages):
        if not free_list:
            raise MemoryError("页面池已满")
        page_id = free_list.pop(0)
        page_ids.append(page_id)
    
    # 记录到页面表
    page_table[seq_id] = page_ids
    
    return page_ids

页面回收：

def free_sequence(seq_id: int):
    """
    回收序列的页面
    
    Args:
        seq_id: 序列 ID
    """
    # 获取页面列表
    page_ids = page_table[seq_id]
    
    # 回收到空闲列表
    free_list.extend(page_ids)
    
    # 从页面表删除
    del page_table[seq_id]

跨页面的 Attention 计算：

def paged_attention(Q, page_table, page_pool):
    """
    使用分页 KV Cache 计算 Attention
    
    Args:
        Q: Query [batch, heads, seq_q, d]
        page_table: 页面表 {seq_id: [page_ids]}
        page_pool: 页面池 [num_pages, heads, page_size, d]
    
    Returns:
        output: [batch, heads, seq_q, d]
    """
    outputs = []
    
    for seq_id in range(batch):
        # 获取该序列的页面
        page_ids = page_table[seq_id]
        
        # 拼接所有页面的 K, V
        K_pages = [page_pool[pid] for pid in page_ids]
        K = torch.cat(K_pages, dim=1)  # [heads, total_tokens, d]
        
        V_pages = [page_pool[pid] for pid in page_ids]
        V = torch.cat(V_pages, dim=1)
        
        # 计算 Attention
        output = scaled_dot_product_attention(Q[seq_id], K, V)
        outputs.append(output)
    
    return torch.stack(outputs)

2.4 Python 实现

import torch
from typing import Dict, List, Optional
from collections import deque

class PagedKVCache:
    """
    Paged KV Cache 实现
    
    内存布局:
    - 页面池: [num_pages, 2, num_heads, page_size, head_dim]
              (2 for K and V)
    - 页面表: {seq_id: [page_id1, page_id2, ...]}
    - 空闲列表: deque([page_id1, page_id2, ...])
    """
    
    def __init__(
        self,
        num_heads: int,
        head_dim: int,
        page_size: int = 16,
        num_pages: int = 1024,
        dtype: torch.dtype = torch.float16,
        device: str = 'cpu'
    ):
        """
        Args:
            num_heads: KV 头数量
            head_dim: 每个头的维度
            page_size: 每个页面的 token 数（通常 16, 32, 64）
            num_pages: 总页面数
            dtype: 数据类型
            device: 设备
        """
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.page_size = page_size
        self.num_pages = num_pages
        self.dtype = dtype
        self.device = device
        
        # 页面池: [num_pages, 2, num_heads, page_size, head_dim]
        self.page_pool = torch.zeros(
            num_pages, 2, num_heads, page_size, head_dim,
            dtype=dtype, device=device
        )
        
        # 页面表: {seq_id: [page_ids]}
        self.page_table: Dict[int, List[int]] = {}
        
        # 空闲列表
        self.free_list = deque(range(num_pages))
        
        # 序列计数器
        self.next_seq_id = 0
    
    def allocate_sequence(self, seq_len: int) -> int:
        """
        为新序列分配页面
        
        Args:
            seq_len: 序列长度
        
        Returns:
            seq_id: 分配的序列 ID
        """
        # 计算需要的页面数
        num_pages_needed = (seq_len + self.page_size - 1) // self.page_size
        
        if len(self.free_list) < num_pages_needed:
            raise MemoryError(
                f"页面池空间不足: 需要 {num_pages_needed} 个页面, "
                f"可用 {len(self.free_list)} 个"
            )
        
        # 分配页面
        page_ids = []
        for _ in range(num_pages_needed):
            page_id = self.free_list.popleft()
            page_ids.append(page_id)
        
        # 分配序列 ID
        seq_id = self.next_seq_id
        self.next_seq_id += 1
        
        # 记录到页面表
        self.page_table[seq_id] = page_ids
        
        return seq_id
    
    def free_sequence(self, seq_id: int):
        """
        释放序列的页面
        
        Args:
            seq_id: 序列 ID
        """
        if seq_id not in self.page_table:
            raise ValueError(f"序列 {seq_id} 不存在")
        
        # 获取页面列表
        page_ids = self.page_table[seq_id]
        
        # 回收到空闲列表
        self.free_list.extend(page_ids)
        
        # 从页面表删除
        del self.page_table[seq_id]
    
    def update(
        self,
        seq_id: int,
        key: torch.Tensor,
        value: torch.Tensor,
        start_pos: int = 0
    ):
        """
        更新序列的 KV Cache
        
        Args:
            seq_id: 序列 ID
            key: [1, num_heads, seq_len, head_dim]
            value: [1, num_heads, seq_len, head_dim]
            start_pos: 起始位置
        """
        if seq_id not in self.page_table:
            raise ValueError(f"序列 {seq_id} 不存在")
        
        page_ids = self.page_table[seq_id]
        seq_len = key.size(2)
        
        # 逐页更新
        for i, token_idx in enumerate(range(start_pos, start_pos + seq_len)):
            # 计算页面索引和页内偏移
            page_idx = token_idx // self.page_size
            offset = token_idx % self.page_size
            
            if page_idx >= len(page_ids):
                raise ValueError(f"Token 索引 {token_idx} 超出分配的页面范围")
            
            page_id = page_ids[page_idx]
            
            # 更新页面池
            self.page_pool[page_id, 0, :, offset, :] = key[0, :, i, :]  # K
            self.page_pool[page_id, 1, :, offset, :] = value[0, :, i, :]  # V
    
    def get_kv(self, seq_id: int) -> tuple[torch.Tensor, torch.Tensor]:
        """
        获取序列的完整 K, V
        
        Args:
            seq_id: 序列 ID
        
        Returns:
            key: [1, num_heads, seq_len, head_dim]
            value: [1, num_heads, seq_len, head_dim]
        """
        if seq_id not in self.page_table:
            raise ValueError(f"序列 {seq_id} 不存在")
        
        page_ids = self.page_table[seq_id]
        
        # 拼接所有页面
        k_pages = []
        v_pages = []
        for page_id in page_ids:
            k_pages.append(self.page_pool[page_id, 0])  # [num_heads, page_size, head_dim]
            v_pages.append(self.page_pool[page_id, 1])
        
        # 拼接: [num_heads, total_len, head_dim]
        key = torch.cat(k_pages, dim=1).unsqueeze(0)
        value = torch.cat(v_pages, dim=1).unsqueeze(0)
        
        return key, value
    
    def get_memory_stats(self) -> dict:
        """获取内存统计信息"""
        total_pages = self.num_pages
        used_pages = sum(len(pages) for pages in self.page_table.values())
        free_pages = len(self.free_list)
        
        total_memory = total_pages * self.page_size * self.num_heads * self.head_dim * 2
        used_memory = used_pages * self.page_size * self.num_heads * self.head_dim * 2
        
        return {
            'total_pages': total_pages,
            'used_pages': used_pages,
            'free_pages': free_pages,
            'utilization': used_pages / total_pages if total_pages > 0 else 0,
            'total_memory_mb': total_memory * 2 / 1024 / 1024,  # FP16
            'used_memory_mb': used_memory * 2 / 1024 / 1024,
            'num_sequences': len(self.page_table)
        }

2.5 性能分析

内存利用率对比：

场景：32 个序列，平均长度 200 tokens，最大长度 2048

内存节省：

支持的 Batch Size 对比（40GB GPU 内存）：

2.6 与 vLLM 的关系

vLLM 是首个大规模应用 Paged KV Cache 的推理框架：

vLLM 的 PagedAttention：

• 页面大小：16 tokens（默认）
• 支持 Copy-on-Write（写时复制）
• 支持 Prefix Caching（前缀缓存共享）

性能提升：

• 吞吐量提升： 24x vs HuggingFace Transformers
• 内存利用率： 95% vs 传统的 20-40%
• 支持更大的 batch size

📊 性能对比总结

Flash Attention vs 标准 Attention

Paged KV Cache vs 连续 KV Cache

组合效果（Flash Attention + Paged KV Cache）

这就是 vLLM、TensorRT-LLM 等高性能推理引擎的核心技术栈！

实际效果（LLaMA 2 70B，A100 80GB）：

提升：

• 吞吐量： 24x
• 内存占用： 49%
• Batch Size： 16x

🎓 学习路径

阶段 1：Flash Attention 基础（2-3 天）

📓 Notebook: 05_flash_attention.ipynb

学习内容：

1. 理解标准 Attention 的内存瓶颈
2. 掌握 Tiling（分块）技术
3. 理解 Online Softmax 算法
4. 实现 Python 版本的 Flash Attention
5. 分析内存和性能优势

关键问题：

• 为什么标准 Attention 需要 O(N²) 内存？
• Tiling 如何避免存储完整矩阵？
• Online Softmax 如何增量更新？
• 如何在分块间合并 softmax 结果？

实践任务：

# 1. 实现 Online Softmax
def online_softmax_merge(old_max, old_sum, new_max, new_sum):
    # TODO: 实现增量合并
    pass

# 2. 实现分块 Attention
def flash_attention_forward(Q, K, V, block_size):
    # TODO: 实现分块计算
    pass

# 3. 性能测试
# 对比标准 Attention 和 Flash Attention 的内存和速度

阶段 2：Paged KV Cache 基础（2-3 天）

📓 Notebook: 06_paged_kv_cache.ipynb

学习内容：

1. 理解连续 KV Cache 的内存碎片问题
2. 掌握分页管理的原理
3. 实现页面分配和回收算法
4. 实现跨页面的 Attention 计算
5. 分析内存利用率

关键问题：

• 连续 KV Cache 为什么会产生碎片？
• 页面大小如何选择？
• 如何实现页面表和空闲列表？
• 如何在分页存储上计算 Attention？

实践任务：

# 1. 实现页面分配
def allocate_sequence(seq_len, page_size):
    # TODO: 分配页面
    pass

# 2. 实现页面回收
def free_sequence(seq_id):
    # TODO: 回收页面
    pass

# 3. 实现分页 Attention
def paged_attention(Q, page_table, page_pool):
    # TODO: 跨页面计算 Attention
    pass

# 4. 内存利用率分析
# 对比连续 Cache 和 Paged Cache 的内存利用率

阶段 3：深入理解（3-5 天）

学习内容：

1. 阅读 Flash Attention 论文
2. 阅读 vLLM PagedAttention 论文
3. 对照 TensorRT-LLM XQA 源码
4. 理解 CUDA 优化技巧

推荐资源：

• Flash Attention 论文
• Flash Attention-2 论文
• vLLM 论文
• TensorRT-LLM GitHub

对照学习：

• Python 实现 → CUDA 实现
• 算法原理 → 工程优化
• 单机优化 → 分布式优化

🔧 实际应用场景

1. 大模型推理优化

场景：部署 LLaMA 2 70B 进行在线推理

优化方案：

# 使用 Flash Attention + Paged KV Cache
from src.flash_attention import FlashAttention
from src.paged_kv_cache import PagedKVCache

# 配置
config = {
    'd_model': 8192,
    'num_heads': 64,
    'num_kv_heads': 8,  # GQA
    'block_size': 64,   # Flash Attention
    'page_size': 16,    # Paged KV Cache
}

# 创建模型
attention = FlashAttention(
    d_model=config['d_model'],
    num_heads=config['num_heads'],
    block_size=config['block_size']
)

cache = PagedKVCache(
    num_heads=config['num_kv_heads'],
    head_dim=config['d_model'] // config['num_heads'],
    page_size=config['page_size'],
    num_pages=4096
)

# 推理
for batch in dataloader:
    # 分配序列
    seq_ids = [cache.allocate_sequence(len(seq)) for seq in batch]
    
    # 前向传播
    output, lse = attention(batch)
    
    # 更新缓存
    for seq_id, k, v in zip(seq_ids, keys, values):
        cache.update(seq_id, k, v)
    
    # 生成完成后释放
    for seq_id in seq_ids:
        cache.free_sequence(seq_id)

效果：

• 吞吐量提升： 20x
• 内存占用减少： 50%
• 支持更大的 batch size

2. 长文本处理

场景：处理 32K 上下文的文档问答

挑战：

• 标准 Attention：32K × 32K = 1B 元素，OOM
• Flash Attention：分块计算，内存 O(N)

实现：

# 支持长序列的 Attention
flash_attn = FlashAttention(
    d_model=4096,
    num_heads=32,
    block_size=128  # 更大的块以提高效率
)

# 处理 32K 上下文
long_text = tokenize(document)  # 32768 tokens
output, _ = flash_attn(long_text)

效果：

• 支持序列长度：2K → 32K
• 内存占用：128GB → 256MB

💡 核心代码片段

Flash Attention 核心算法

def flash_attention_forward(Q, K, V, block_size=64):
    """Flash Attention 核心算法"""
    batch, heads, seq_len, d = Q.shape
    scale = 1.0 / math.sqrt(d)
    
    # 初始化
    O = torch.zeros_like(Q)
    l = torch.zeros(batch, heads, seq_len, 1)
    m = torch.full((batch, heads, seq_len, 1), float('-inf'))
    
    # 外层循环：Q 的块
    for i in range(0, seq_len, block_size):
        Q_i = Q[:, :, i:i+block_size, :]
        O_i = torch.zeros_like(Q_i)
        l_i = torch.zeros(batch, heads, Q_i.size(2), 1)
        m_i = torch.full((batch, heads, Q_i.size(2), 1), float('-inf'))
        
        # 内层循环：K, V 的块
        for j in range(0, seq_len, block_size):
            K_j = K[:, :, j:j+block_size, :]
            V_j = V[:, :, j:j+block_size, :]
            
            # 计算注意力分数
            S_ij = torch.matmul(Q_i, K_j.transpose(-2, -1)) * scale
            
            # Online Softmax 更新
            m_ij = S_ij.max(dim=-1, keepdim=True)[0]
            m_i_new = torch.maximum(m_i, m_ij)
            P_ij = torch.exp(S_ij - m_i_new)
            l_i_new = torch.exp(m_i - m_i_new) * l_i + P_ij.sum(dim=-1, keepdim=True)
            
            # 更新输出
            O_i = torch.exp(m_i - m_i_new) * O_i + torch.matmul(P_ij, V_j)
            m_i = m_i_new
            l_i = l_i_new
        
        # 归一化并存储
        O[:, :, i:i+block_size, :] = O_i / l_i
    
    return O

Paged KV Cache 核心算法

class PagedKVCache:
    """Paged KV Cache 核心实现"""
    
    def __init__(self, num_heads, head_dim, page_size=16, num_pages=1024):
        # 页面池
        self.page_pool = torch.zeros(num_pages, 2, num_heads, page_size, head_dim)
        # 页面表
        self.page_table = {}
        # 空闲列表
        self.free_list = deque(range(num_pages))
    
    def allocate_sequence(self, seq_len):
        """分配页面"""
        num_pages = (seq_len + self.page_size - 1) // self.page_size
        page_ids = [self.free_list.popleft() for _ in range(num_pages)]
        seq_id = len(self.page_table)
        self.page_table[seq_id] = page_ids
        return seq_id
    
    def free_sequence(self, seq_id):
        """回收页面"""
        page_ids = self.page_table.pop(seq_id)
        self.free_list.extend(page_ids)
    
    def update(self, seq_id, key, value, start_pos=0):
        """更新缓存"""
        page_ids = self.page_table[seq_id]
        for i, token_idx in enumerate(range(start_pos, start_pos + key.size(2))):
            page_idx = token_idx // self.page_size
            offset = token_idx % self.page_size
            page_id = page_ids[page_idx]
            self.page_pool[page_id, 0, :, offset, :] = key[0, :, i, :]
            self.page_pool[page_id, 1, :, offset, :] = value[0, :, i, :]

📚 参考资料

论文

• Flash Attention: Fast and Memory-Efficient Exact Attention - Flash Attention 原始论文
• Flash Attention-2: Faster Attention with Better Parallelism - Flash Attention v2
• Efficient Memory Management for Large Language Model Serving with PagedAttention - vLLM 论文

开源项目

• Flash Attention - 官方 CUDA 实现
• vLLM - PagedAttention 实现
• TensorRT-LLM - NVIDIA 高性能推理引擎
• xFormers - Meta 的高效 Transformer 实现

博客和教程

• Flash Attention 详解
• vLLM 技术博客
• TensorRT-LLM 文档

📝 总结

Flash Attention 和 Paged KV Cache 是大模型推理优化的两大核心技术：

Flash Attention

• 问题：标准 Attention 需要 O(N²) 内存存储注意力矩阵
• 解决：通过 Tiling 和 Online Softmax，将内存降到 O(N)
• 效果：支持 32K+ 长序列，内存节省 99%+，速度提升 2-4x

Paged KV Cache

• 问题：连续 KV Cache 导致内存碎片，利用率低
• 解决：通过分页管理，动态分配和回收页面
• 效果：内存利用率从 20% 提升到 95%，支持 batch size 提升 4x

组合效果

• vLLM：吞吐量提升 24x
• TensorRT-LLM：端到端推理加速 8x
• 工业界标准：几乎所有高性能推理引擎都采用这两项技术

下一步

• 深入学习 CUDA 实现
• 研究 Flash Attention-2 的进一步优化
• 探索分布式推理优化
• 学习混合精度和量化技术

🚀 快速开始实践

系统命令

# 1. 克隆项目
git clone https://github.com/rixin2025/attention-from-scratch.git
cd attention-from-scratch

# 2. 创建虚拟环境（推荐）
python -m venv venv
.\venv\Scripts\Activate.ps1

# 3. 安装依赖
pip install -r requirements.txt

# 4. 运行 Jupyter Notebook
jupyter notebook notebooks\

# 5. 运行测试
pytest tests\ -v

# 6. 运行示例
python demo.py

学习建议

1. 先学基础篇：确保理解 MHA、GQA、KV Cache
2. 逐步深入：先理解原理，再看代码，最后动手实现
3. 对比学习：对比标准方法和优化方法的差异
4. 性能分析：实际测试内存和速度的提升
5. 阅读论文：深入理解算法的数学原理
6. 对照源码：学习 vLLM、TensorRT-LLM 的工程实现

🤝 贡献

欢迎提交 Issue 和 Pull Request！

如果你觉得这个项目对你有帮助，请给一个 ⭐ Star，这是对我最大的鼓励！

💡开源代码工程：https://github.com/rixin2025/attention-from-scratch/tree/main

让更多人理解大模型推理优化的核心技术，一起推动 AI 技术的发展！

*本文是《从零实现 Attention 机制》系列的进阶篇，基础篇请参考 https://blog.csdn.net/CSDN_3195/article/details/158179338?spm=1001.2014.3001.5502

*后续将继续深入 性能瓶颈分析/CUDA 优化 等更高级的主题，敬请期待！