LLM推理优化

📅 4周详细学习计划

1. KV Cache详解

1.1 问题背景

• 自回归生成：LLM逐token生成，每个token需要所有历史token的K,V
• 重复计算：如果不缓存，每个token都要重新计算所有K,V（O(n²)复杂度）
• KV Cache：缓存已计算的K,V，新token只需计算自己的Q,K,V

1.2 内存占用计算

// KV Cache内存计算公式
// 对于单个序列：
KV_Cache_Size = 2 × num_layers × num_heads × seq_len × head_dim × dtype_size

// 示例：Llama-2-70B
// num_layers=80, num_heads=64, head_dim=128, seq_len=4096, FP16
KV_Cache = 2 × 80 × 64 × 4096 × 128 × 2 bytes
         = 2 × 80 × 64 × 4096 × 128 × 2
         = 10.7 GB (单序列！)

// 这就是为什么大模型推理需要巨大显存

1.3 KV Cache优化策略

• GQA (Grouped Query Attention)：多个Q head共享K,V head
• MQA (Multi-Query Attention)：所有Q head共享一个K,V head
• 量化KV Cache：将K,V从FP16量化到INT8/INT4

2. FlashAttention详解

2.1 核心思想

• IO感知：优化GPU内存访问模式，减少HBM读写
• 分块计算(Tiling)：将Q,K,V分块加载到SRAM（Shared Memory）
• Online Softmax：分块计算Softmax，无需存储完整N×N矩阵
• 融合Kernel：将多个操作合并为一个CUDA Kernel

2.2 内存对比

// Standard Attention
// 1. 计算 S = Q × K^T (N×N矩阵，O(N²)内存)
// 2. 计算 P = softmax(S) (N×N矩阵)
// 3. 计算 O = P × V (输出)
// 总内存: O(N²) + O(N×d) = O(N²)

// FlashAttention
// 1. 分块加载Q,K,V到SRAM
// 2. 在SRAM中计算部分注意力
// 3. Online Softmax更新统计量
// 4. 逐步累加输出
// 总内存: O(N) - 只需要存储输出和统计量

2.3 使用示例

# FlashAttention Python接口
from flash_attn import flash_attn_func

# q, k, v: [batch, seqlen, nheads, headdim]
# 不同head_dim需要不同版本的FlashAttention
output = flash_attn_func(q, k, v, causal=True)

# FlashAttention-2 (更快)
from flash_attn import flash_attn_func
output = flash_attn_func(q, k, v, causal=True, window_size=(-1, -1))

3. PagedAttention (vLLM)

3.1 问题

• 连续分配：KV Cache需要连续内存，导致内存碎片
• 预分配：为最大序列长度预分配，浪费大量内存
• 内存利用率：传统方法只有20-40%的利用率

3.2 解决方案

• 分页管理：将KV Cache分为固定大小的Block（如16 tokens）
• 非连续存储：Block可以存储在任意位置
• 页表映射：逻辑Block → 物理Block的映射
• Copy-on-Write：多序列共享相同的KV Cache Block

# PagedAttention内存管理
# 传统方法：每个序列预分配max_seq_len内存
# PagedAttention：按需分配Block

# 示例：
# 序列1: 100 tokens → 7个Block (每个Block 16 tokens)
# 序列2: 50 tokens → 4个Block
# 序列3: 200 tokens → 13个Block
# 总共: 24个Block (而不是350个Block预分配)

3.3 vLLM使用

# vLLM安装和使用
pip install vllm

# 基本使用
from vllm import LLM, SamplingParams

# 初始化模型
llm = LLM(
    model="meta-llama/Llama-2-7b-hf",
    tensor_parallel_size=2,  # 2卡并行
    max_model_len=4096,
    gpu_memory_utilization=0.9
)

# 批量生成
prompts = ["Hello, how are you?", "What is AI?"]
sampling_params = SamplingParams(temperature=0.7, max_tokens=100)
outputs = llm.generate(prompts, sampling_params)

4. Continuous Batching

4.1 Static vs Continuous

• Static Batching：等所有序列完成才处理新请求，GPU空闲
• Continuous Batching：序列完成立即处理新请求，GPU持续忙碌
• 性能提升：2-3倍吞吐量提升

4.2 实现原理

• 迭代级调度：每个iteration重新调度batch
• 插入新请求：有空位时插入新序列
• 移除完成序列：序列完成时立即释放资源

5. 投机解码 (Speculative Decoding)

5.1 原理

• Draft模型：小模型快速生成多个token（如10个）
• Target模型：大模型验证这10个token
• 接受/拒绝：逐个验证，接受正确的，拒绝错误的
• 加速比：2-3倍速度提升，保持大模型质量

5.2 实现方案

• Medusa：多头预测，每个头预测不同位置
• EAGLE：特征级草稿，更准确
• Lookahead：推测多个未来token