大模型推理加速技术全景：从FlashAttention到Speculative Decoding

一、大模型推理的瓶颈分析

1.1 推理的两阶段：Prefill vs Decode

LLM推理可以划分为两个阶段：Prefill（预填充）和Decode（自回归解码）。Prefill阶段一次性处理输入的prompt token，以计算复杂度O(L²d)为主；Decode阶段逐token生成，以内存带宽O(Ld)为主。两者对计算资源的需求截然不同，优化策略也因此分道扬镳。

1.2 瓶颈量化分析

以一个70B模型在A100-80G上为例，模型权重约140GB（FP16）。Decode阶段需要每次生成将模型权重从HBM加载到SRAM，加载耗时约140GB / 2TB/s = 70ms，而实际计算仅需数毫秒。这意味着Decode阶段的瓶颈本质上是"冯·诺依曼瓶颈"——数据搬移速度远快于实际运算速度。

大模型推理瓶颈示意图
═════════════════════════════════════════════════════════════════

                     ┌─────────────────────┐
                     │   GPU HBM (80GB)    │
                     │  带宽: 2TB/s        │ ◀── 主要瓶颈 !
                     │                     │
                     │ ┌───────────────┐   │
                     │ │ 模型权重 140GB │   │
                     │ │ (FP16 70B)    │   │
                     │ └───────────────┘   │
                     │ ┌───────────────┐   │
                     │ │ KV Cache 动态  │   │
                     │ │ (seq_len × d) │   │
                     │ └───────────────┘   │
                     └────────┬────────────┘
                              │ 加载权重（每次decode）
                              ▼
                     ┌─────────────────────┐
                     │  GPU SRAM (40MB)    │
                     │  带宽: 20TB/s       │  ← 60x HBM带宽
                     │                     │
                     │ 实际计算在这里完成   │
                     │ MatMul ≈ 1-2ms      │
                     └─────────────────────┘
        

二、FlashAttention：让注意力计算不再受限于内存带宽

2.1 标准Attention的内存瓶颈

标准Attention的计算流程是：QK^T → Softmax → SV。这需要将Q、K矩阵从HBM加载到SRAM，计算S矩阵并写回HBM，再从HBM读取S进行Softmax，写回softmax(S)，最后计算softmax(S)V。整个过程涉及多次HBM读写，其中间结果S矩阵的大小为N×N，这对长序列来说极为昂贵。

2.2 FlashAttention的核心思想

FlashAttention的核心突破在于"tiling"（分块）策略。它通过将Q、K、V切分为小块，在SRAM内完成整个注意力计算后再写回HBM，避免了中间结果的HBM读写。但这里有一个关键问题：Softmax需要全局规约（所有token的score），分块后如何正确计算？

FlashAttention的解决方案是"在线Softmax"（Online Softmax）：通过维护两个统计量（局部最大值m和局部指数和l），在合并分块时动态修正。这一技巧使得分块计算等价于全局Softmax。

FlashAttention 分块算法伪代码：
for each block Q_i, K_j, V_j:
    # 读取到SRAM
    load_block(Q_i, K_j, V_j)
    
    # 计算当前块的注意力分数
    S_ij = Q_i @ K_j.T   # 在SRAM内完成
    
    # 在线Softmax（关键创新）
    m_ij = max(m_i_prev, rowmax(S_ij))  
    P_ij = exp(S_ij - m_ij)
    l_ij = exp(m_i_prev - m_ij) * l_i_prev + rowsum(P_ij)
    
    # 加权累加
    O_i = diag(exp(m_i_prev - m_ij)) * O_i_prev + P_ij @ V_j
    
    # 最终归一化
    O_i = diag(l_i)^(-1) * O_i

# 结果：O = softmax(QK^T)V，无需HBM写中间矩阵

2.3 FlashAttention-2/3的演进

FlashAttention-2优化了线程束调度和并行策略，减少非计算开销，在A100上实现约2x的速度提升。FlashAttention-3（2024年）进一步利用Hopper架构的WGMMA（Warp Group Matrix Multiply-Accumulate）指令和异步拷贝特性，达到理论FLOPS的65-75%，几乎逼近硬件极限。

三、PagedAttention与vLLM：KV Cache的内存管理革命

3.1 KV Cache的内存碎片问题

在传统推理系统中，KV Cache为每个请求预分配最大长度（如2048或4096）的连续内存块。这导致严重的"内部碎片"——实际生成的序列往往远短于预分配长度，但剩余内存无法被其他请求复用。同时，这种固定分配方式在大规模连续批处理时，GPU显存的利用率通常只有30-50%。

3.2 操作系统的分页思想

PagedAttention的灵感直接来自虚拟内存管理：将KV Cache按固定大小（如16个token）分页（Page），每个请求持有一个页表，逻辑上连续的序列在物理内存中可以是分散的页面。这带来三个收益：消除内部碎片、按需分配、页面级共享（如同一prompt前缀的多个请求可共享KV页）。

PagedAttention 内存管理对比
═════════════════════════════════════════════════════════════════

传统方案（连续内存）:
┌────────────────────────────────────────┐
│ Request A：预分配4096连续空间           │
│ ┌─────────────┬──────────────────┐     │
│ │ 实际使用1024 │   内部碎片3072    │     │
│ └─────────────┴──────────────────┘     │
│ Request B：预分配4096连续空间           │
│ ┌──────┬─────────────────────────┐     │
│ │ 512  │   内部碎片3584           │     │
│ └──────┴─────────────────────────┘     │
│ ......                                 │
│ 内存利用率 ≈ 35%                        │
└────────────────────────────────────────┘

PagedAttention（分页管理）:
┌────────────────────────────────────────┐
│ 全局KV Block Pool (每个Block=16 token) │
│ ┌──┬──┬──┬──┬──┬──┬──┬──┬──┬──┬──┐  │
│ │A1│A2│B1│  │  │A3│  │A4│A5│B2│  │  │
│ └──┴──┴──┴──┴──┴──┴──┴──┴──┴──┴──┘  │
│                                        │
│ A的页表: [0, 1, 5, 7, 8]              │
│ B的页表: [2, 9]                        │
│ 内存利用率 ≈ 75%+                       │
└────────────────────────────────────────┘
        

四、Continuous Batching：连续批处理的吞吐量优化

4.1 静态批处理的问题

传统批处理方案以"请求到达→等批满→全部推理→返回"的方式运行。这导致"晚到"的请求必须等待下一次批处理，而批内所有请求的生成又受限于最慢的那个。整个系统的GPU利用率在长尾请求上表现极差。

4.2 Continuous Batching的工作机制

Continuous Batching（或称Iteration-level Batching）在每一步decode时动态调整batch。已完成生成的请求退出batch，新到达的请求立即加入下一轮iteration。这消除了等待时间，显著提高了GPU吞吐量。

五、Speculative Decoding：推测解码的策略与实践

5.1 核心思想

Speculative Decoding的核心洞察：小模型生成优质文本的速度远快于大模型（10-20x）。我们可以用小模型"打草稿"（初步生成多个候选token），再通过大模型并行验证。如果草稿被接受，一次大模型推理能"赚"多个token的生成量，从而突破自回归解码的"一步一推理"瓶颈。

5.2 两种主流策略

5.3 验证机制的数学保证

Speculative Decoding的一个重要特性是"无损"——它的输出分布与原始大模型的分布完全一致（不像量化或剪枝会改变输出）。这得益于拒绝采样（Rejection Sampling）的正确性保证：

Speculative Decoding 验证步骤：
1. 小模型生成K个候选token: x_1, x_2, ..., x_K
2. 大模型并行计算: q(x | prefix) 对所有候选位置
3. 对每个位置t = 1..K:
   - 以概率 min(1, p_target(x_t) / p_draft(x_t)) 接受
   - 如果拒绝，从分布 max(0, p_target - p_draft) 中采样一个token
4. 接受位置后的所有草稿被丢弃，继续下一轮

数学保证：最终分布 ≡ 纯大模型自回归分布

六、Quantization：量化推理的精度与性能权衡

6.1 量化策略对比

七、Prefix Caching与RadixAttention

7.1 系统提示词复用

在实际应用中，大量请求共享相同的System Prompt（如聊天机器人、代码助手等）。传统方案每次请求都从零计算KV Cache，浪费大量算力。SGLang提出的RadixAttention通过前缀树（Radix Tree）管理共享前缀的KV Cache，实现了细粒度的内存复用。

7.2 前缀树架构

Radix Tree 前缀缓存示例
═════════════════════════════════════════════════════════════════

请求A: "你是一个AI助手，帮我写一首诗..."
请求B: "你是一个AI助手，帮我算一下1+1..."
请求C: "你是一个AI助手，分析以下代码..."

           root
            │
            │ "你是一个AI助手"
            │
            ┌─────────┼─────────┐
            │         │         │
        "帮我写"  "帮我算"  "分析"
            │         │         │
         "一首诗"   "一下1+1"  "以下代码"
            │         │         │
         [缓存A]   [缓存B]   [缓存C]
        
共享前缀的KV Cache只存储一份
        

八、并行解码策略：从张量并行到序列并行

8.1 张量并行（Tensor Parallelism）

张量并行将单个Transformer层的权重沿hidden维度切分到多个GPU。每个GPU负责计算部分注意力头和FFN维度，通过all-reduce通信规约结果。对于超大模型（70B+），张量并行是必需的，但通信开销（约30-50%）在大规模集群上不可忽视。

8.2 流水线并行（Pipeline Parallelism）

流水线并行将Transformer层堆叠按深度切分，各GPU负责一组连续的层。优点：通信量小（仅传输激活值）；缺点：存在气泡（bubble）问题。通过微批次（Micro-batch）和1F1B调度可逼近100%利用率。

九、推理引擎对比：vLLM vs TensorRT-LLM vs SGLang

十、深挖点：Attention计算复杂度与FLOPs分析

10.1 Attention的算力需求公式

一个标准的Self-Attention层在N个token、隐含维度d下的计算复杂度为O(4Nd² + 2N²d)。前两项Nd²来自QKV投影和输出投影，后一项N²d来自attention矩阵的计算和加权。当N很小（如对话场景）时，Nd²占主导；当N很大（如文档理解）时，N²d成为瓶颈。

10.2 理论FLOPs与实际推理速度

以70B模型推理为例，理论FLOPs与实际吞吐量之间存在巨大差距——接近10倍。原因在于：(1) Decode阶段以内存带宽为主，GPU计算单元闲置；(2) Attention计算的矩阵乘法无法完全利用Tensor Core；(3) 通信开销（多卡部署）。