LLM 推理优化：KV Cache、FlashAttention 与投机解码
1. 引言
LLM 推理的核心瓶颈在于自回归解码------每个 token 都依赖前面所有 token 的 KV 值。本文将深入讲解三大推理优化技术：KV Cache 管理、FlashAttention 加速和投机解码（Speculative Decoding）。
推理性能瓶颈：
LLM 推理分为两个阶段：
1. Prefill 阶段：处理全部输入 token（计算密集）
2. Decode 阶段：逐个生成输出 token（内存密集）
主要瓶颈：
- KV Cache 内存随序列长度线性增长
- 注意力计算的内存访问瓶颈（Memory-bound）
- 自回归解码无法并行2. KV Cache 原理与优化
2.1 基础 KV Cache
python
class KVCache:
"""基础 KV Cache 实现"""
def __init__(self, num_layers, num_heads, head_dim, max_seq_len, batch_size=1):
self.cache = {}
for layer in range(num_layers):
self.cache[layer] = {
"k": torch.zeros(batch_size, num_heads, max_seq_len, head_dim),
"v": torch.zeros(batch_size, num_heads, max_seq_len, head_dim),
}
self.seq_len = 0
def update(self, layer, key, value):
"""更新 KV Cache"""
self.cache[layer]["k"][:, :, self.seq_len:self.seq_len+key.size(2)] = key
self.cache[layer]["v"][:, :, self.seq_len:self.seq_len+value.size(2)] = value
def get(self, layer):
"""获取当前层的 KV"""
return (
self.cache[layer]["k"][:, :, :self.seq_len],
self.cache[layer]["v"][:, :, :self.seq_len],
)2.2 GQA (Grouped Query Attention) KV Cache 压缩
python
# GQA：多个 Query Head 共享一组 KV Head
# 8 个 Query Head → 2 个 KV Head（KV Cache 减少 4x）
class GroupedQueryAttention(nn.Module):
def __init__(self, d_model, num_q_heads=32, num_kv_heads=8):
super().__init__()
self.num_q_heads = num_q_heads
self.num_kv_heads = num_kv_heads
self.num_groups = num_q_heads // num_kv_heads
self.head_dim = d_model // num_q_heads
self.q_proj = nn.Linear(d_model, num_q_heads * self.head_dim)
self.k_proj = nn.Linear(d_model, num_kv_heads * self.head_dim)
self.v_proj = nn.Linear(d_model, num_kv_heads * self.head_dim)
self.o_proj = nn.Linear(d_model, d_model)
def forward(self, x, kv_cache=None):
B, L, _ = x.shape
q = self.q_proj(x).view(B, L, self.num_q_heads, self.head_dim).transpose(1, 2)
k = self.k_proj(x).view(B, L, self.num_kv_heads, self.head_dim).transpose(1, 2)
v = self.v_proj(x).view(B, L, self.num_kv_heads, self.head_dim).transpose(1, 2)
# 更新 KV Cache
if kv_cache is not None:
k_cache, v_cache = kv_cache
k = torch.cat([k_cache, k], dim=2)
v = torch.cat([v_cache, v], dim=2)
# 扩展 KV 以匹配 Q 的头数
k = k.repeat_interleave(self.num_groups, dim=1)
v = v.repeat_interleave(self.num_groups, dim=1)
# 注意力计算
attn = (q @ k.transpose(-2, -1)) / (self.head_dim ** 0.5)
attn = attn.softmax(dim=-1)
out = (attn @ v).transpose(1, 2).reshape(B, L, -1)
return self.o_proj(out), (k[:, :, -L:], v[:, :, -L:])2.3 PagedAttention（vLLM）
python
# PagedAttention 将 KV Cache 分为固定大小的块（page）
# 类似操作系统的虚拟内存分页机制
class PagedKVCache:
"""分页 KV Cache"""
def __init__(self, num_blocks, block_size, num_heads, head_dim):
self.block_size = block_size
# 预分配物理块
self.k_blocks = torch.zeros(num_blocks, num_heads, block_size, head_dim)
self.v_blocks = torch.zeros(num_blocks, num_heads, block_size, head_dim)
self.free_blocks = list(range(num_blocks))
self.block_tables = {} # seq_id -> [block_indices]
def allocate(self, seq_id):
"""为序列分配新块"""
block_idx = self.free_blocks.pop()
if seq_id not in self.block_tables:
self.block_tables[seq_id] = []
self.block_tables[seq_id].append(block_idx)
return block_idx
def free(self, seq_id):
"""释放序列的所有块"""
if seq_id in self.block_tables:
self.free_blocks.extend(self.block_tables[seq_id])
del self.block_tables[seq_id]3. FlashAttention
3.1 原理
标准注意力：
S = Q @ K^T → 需要存储 N×N 矩阵（显存 O(N²)）
P = softmax(S)
O = P @ V
FlashAttention：
将 Q, K, V 分为小块（tile）
在 SRAM 中完成注意力计算
通过在线 softmax 算法避免存储完整注意力矩阵
显存 O(N)，速度提升 2-4x3.2 使用 FlashAttention
python
# PyTorch 2.0+ 原生支持
import torch.nn.functional as F
# 方式1：使用 torch.nn.functional.scaled_dot_product_attention
output = F.scaled_dot_product_attention(
query, key, value,
attn_mask=None,
is_causal=True, # 因果注意力
)
# 方式2：使用 flash-attn 库
from flash_attn import flash_attn_func
output = flash_attn_func(
query, key, value,
causal=True,
softmax_scale=1.0 / (head_dim ** 0.5),
)
# 方式3：在模型中启用
model = AutoModelForCausalLM.from_pretrained(
model_id,
attn_implementation="flash_attention_2", # 启用 FlashAttention
)3.3 FlashAttention 性能对比
| 序列长度 | 标准注意力 | FlashAttention v2 |
|---|---|---|
| 1K | 8ms | 4ms |
| 4K | 45ms | 12ms |
| 8K | 180ms | 25ms |
| 16K | OOM | 52ms |
| 32K | OOM | 108ms |
4. 投机解码（Speculative Decoding）
4.1 原理
核心思想：用小模型快速生成多个候选 token，大模型一次性验证
传统解码（7B模型，每步1个token）：
Step1 → t1 → Step2 → t2 → Step3 → t3 → Step4 → t4
耗时：4 × 20ms = 80ms
投机解码（160M小模型生成 + 7B大模型验证）：
小模型快速生成: t1, t2, t3, t4 (4ms)
大模型一次验证: [t1✓, t2✓, t3✗, t4-] (22ms)
修正: t1, t2, t3', t4' (4ms)
耗时：30ms（接受3个token）4.2 实现代码
python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
class SpeculativeDecoder:
"""投机解码器"""
def __init__(self, draft_model_id, target_model_id):
self.tokenizer = AutoTokenizer.from_pretrained(target_model_id)
# 小模型（Draft Model）
self.draft_model = AutoModelForCausalLM.from_pretrained(
draft_model_id, device_map="auto", torch_dtype=torch.float16
)
# 大模型（Target Model）
self.target_model = AutoModelForCausalLM.from_pretrained(
target_model_id, device_map="auto", torch_dtype=torch.float16
)
@torch.no_grad()
def generate(self, prompt, max_tokens=256, gamma=5):
"""投机解码生成"""
input_ids = self.tokenizer.encode(prompt, return_tensors="pt").to("cuda")
generated = input_ids.clone()
total_accepted = 0
total_speculated = 0
while generated.size(1) < max_tokens + input_ids.size(1):
# 1. Draft Model 生成 gamma 个候选 token
draft_tokens = []
draft_probs = []
current = generated.clone()
for _ in range(gamma):
draft_output = self.draft_model(current)
draft_logits = draft_output.logits[:, -1, :]
draft_prob = torch.softmax(draft_logits, dim=-1)
token = torch.multinomial(draft_prob, 1)
draft_tokens.append(token)
draft_probs.append(draft_prob)
current = torch.cat([current, token], dim=1)
draft_tokens = torch.cat(draft_tokens, dim=1) # (1, gamma)
# 2. Target Model 一次性验证所有候选 token
verify_input = torch.cat([generated, draft_tokens], dim=1)
target_output = self.target_model(verify_input)
target_logits = target_output.logits[:, -gamma-1:-1, :]
target_probs = torch.softmax(target_logits, dim=-1)
# 3. 逐个验证 token
accepted = 0
for i in range(gamma):
token = draft_tokens[0, i].item()
draft_p = draft_probs[i][0, token].item()
target_p = target_probs[0, i, token].item()
# 接受/拒绝采样
if torch.rand(1).item() < min(1, target_p / max(draft_p, 1e-10)):
accepted += 1
generated = torch.cat([generated, draft_tokens[:, i:i+1]], dim=1)
else:
# 从修正分布中采样
corrected = torch.clamp(target_probs[0, i] - draft_probs[i][0], min=0)
corrected = corrected / corrected.sum()
new_token = torch.multinomial(corrected, 1).unsqueeze(0)
generated = torch.cat([generated, new_token], dim=1)
break
total_accepted += accepted
total_speculated += gamma
# 如果所有 token 都被接受，还可以加上 target model 的额外 token
if accepted == gamma:
extra_logits = target_output.logits[:, -1, :]
extra_probs = torch.softmax(extra_logits, dim=-1)
extra_token = torch.multinomial(extra_probs, 1)
generated = torch.cat([generated, extra_token], dim=1)
if generated[0, -1].item() == self.tokenizer.eos_token_id:
break
acceptance_rate = total_accepted / total_speculated if total_speculated > 0 else 0
print(f"接受率: {acceptance_rate:.2%}")
return self.tokenizer.decode(generated[0][input_ids.size(1):], skip_special_tokens=True)4.3 投机解码性能
| 配置 | 速度提升 | 接受率 |
|---|---|---|
| 160M + 7B | 2.0-2.5x | 70-80% |
| 1B + 7B | 2.5-3.0x | 80-90% |
| 1B + 13B | 2.0-2.5x | 70-80% |
| 7B + 70B | 1.5-2.0x | 60-70% |
5. 量化推理优化
python
# INT8 量化推理
from transformers import BitsAndBytesConfig
bnb_config = BitsAndBytesConfig(
load_in_8bit=True,
llm_int8_threshold=6.0,
)
model = AutoModelForCausalLM.from_pretrained(
model_id,
quantization_config=bnb_config,
device_map="auto",
)6. vLLM 高性能推理
python
from vllm import LLM, SamplingParams
# 初始化 vLLM 引擎
llm = LLM(
model="meta-llama/Llama-2-7b-chat-hf",
tensor_parallel_size=1, # GPU 数量
gpu_memory_utilization=0.9, # GPU 显存利用率
max_model_len=4096, # 最大序列长度
dtype="half", # FP16
)
# 批量推理
prompts = ["Hello, how are you?", "What is AI?"]
sampling_params = SamplingParams(
temperature=0.7,
top_p=0.9,
max_tokens=256,
)
outputs = llm.generate(prompts, sampling_params)
for output in outputs:
print(output.outputs[0].text)7. 总结
LLM 推理优化的核心技术：
| 技术 | 优化目标 | 提升幅度 |
|---|---|---|
| KV Cache | 避免重复计算 | 基础必备 |
| GQA | KV Cache 内存 | 减少 4x |
| PagedAttention | KV Cache 利用率 | 提升 2-4x |
| FlashAttention | 注意力计算 | 加速 2-4x |
| 投机解码 | 解码并行度 | 加速 2-3x |
| 量化 | 模型大小和计算 | 加速 1.5-3x |
| vLLM | 综合优化 | 吞吐提升 5-24x |