大语言模型（LLM）推理的核心瓶颈在于显存带宽和计算效率。本文深入剖析 KV Cache 的内存管理策略与 Flash Attention 的 IO-aware 算子融合技术，展示如何将推理吞吐量提升 3-5 倍。
一、自回归推理的内存瓶颈
LLM 的自回归推理每次生成一个 token，需要计算所有历史 token 的 Key 和 Value。如果不缓存，每次都要重新计算整个序列，计算量呈 O(n²) 增长。KV Cache 将已计算的 K、V 矩阵缓存起来，每次只需计算新 token 的注意力。
class KVCache:
def __init__(self, num_layers, num_heads, head_dim, max_seq_len):
self.cache = [{
'k': torch.zeros(max_seq_len, num_heads, head_dim),
'v': torch.zeros(max_seq_len, num_heads, head_dim),
} for _ in range(num_layers)]
self.seq_len = 0
def update(self, layer_idx, new_k, new_v):
self.cache[layer_idx]['k'][self.seq_len] = new_k
self.cache[layer_idx]['v'][self.seq_len] = new_v
def get(self, layer_idx):
return (self.cache[layer_idx]['k'][:self.seq_len],
self.cache[layer_idx]['v'][:self.seq_len])二、PagedAttention：分页式 KV Cache 管理
传统 KV Cache 为每个序列预分配连续的最大长度内存，造成严重浪费。PagedAttention（vLLM 采用）将 KV Cache 划分为固定大小的 Block，按需分配，类似操作系统的虚拟内存分页机制，显存利用率从 20-40% 提升到 96%。
class PagedKVCache:
def __init__(self, block_size=16):
self.block_size = block_size
self.blocks = {} # block_id -> (k_block, v_block)
self.block_table = {} # seq_id -> [block_ids]
def allocate_block(self):
block_id = len(self.blocks)
self.blocks[block_id] = (
torch.zeros(self.block_size, num_heads, head_dim),
torch.zeros(self.block_size, num_heads, head_dim)
)
return block_id
def append_token(self, seq_id, layer_idx, k, v):
if seq_id not in self.block_table:
self.block_table[seq_id] = []
table = self.block_table[seq_id]
if not table or self._block_full(table[-1]):
table.append(self.allocate_block())
block = self.blocks[table[-1]]
offset = self._next_offset(table[-1])
block[0][offset] = k
block[1][offset] = v三、Flash Attention：IO-Aware 精确注意力
标准 Attention 需要将完整的 QK^T 矩阵写入 HBM（显存），再读取进行 Softmax。Flash Attention 通过分块（Tiling）和重计算（Recomputation）技术，将中间结果保持在 SRAM 中，减少 HBM 读写次数，实现 2-4 倍加速。
def flash_attention(Q, K, V, block_size=128):
B, H, N, d = Q.shape
O = torch.zeros_like(Q)
L = torch.zeros(B, H, N, 1)
m = torch.full((B, H, N, 1), -float('inf'))
for i in range(0, N, block_size):
Qi = Q[:, :, i:i+block_size]
Oi = torch.zeros(B, H, block_size, d)
Li = torch.zeros(B, H, block_size, 1)
mi = torch.full((B, H, block_size, 1), -float('inf'))
for j in range(0, N, block_size):
Kj = K[:, :, j:j+block_size]
Vj = V[:, :, j:j+block_size]
Sij = torch.matmul(Qi, Kj.transpose(-2, -1)) / math.sqrt(d)
mij_new = torch.max(mi, Sij.max(dim=-1, keepdim=True).values)
Pij = torch.exp(Sij - mij_new)
Lij = torch.exp(mi - mij_new) * Li + Pij.sum(dim=-1, keepdim=True)
Oi = torch.exp(mi - mij_new) * Oi + torch.matmul(Pij, Vj)
mi, Li = mij_new, Lij
O[:, :, i:i+block_size] = Oi / Li
L[:, :, i:i+block_size] = Li + torch.log(Li)
return O四、INT4/INT8 量化技术
通过将 FP16 权重和激活值量化为 INT4/INT8，可将显存占用降低 4-8 倍。GPTQ 和 AWQ 是两种主流的权重量化方法，前者使用 OBQ 逐列量化，后者通过激活感知的缩放因子保护重要通道。
def quantize_int4(weight_fp16, group_size=128):
"""GPTQ-style 分组量化"""
num_groups = weight_fp16.shape[0] // group_size
scales = torch.zeros(num_groups, weight_fp16.shape[1])
zeros = torch.zeros(num_groups, weight_fp16.shape[1], dtype=torch.int32)
weight_int4 = torch.zeros_like(weight_fp16, dtype=torch.int8)
for g in range(num_groups):
group = weight_fp16[g*group_size:(g+1)*group_size]
w_max = group.max()
w_min = group.min()
scale = (w_max - w_min) / 15.0
scales[g] = scale
zeros[g] = torch.round(-w_min / scale).int()
weight_int4[g*group_size:(g+1)*group_size] = \
torch.clamp(torch.round(group / scale) + zeros[g], 0, 15)
return weight_int4, scales, zeros五、Continuous Batching
传统静态批处理要求所有请求等长，造成 GPU 空闲。Continuous Batching 允许不同长度的请求动态加入/离开批次，在每次迭代中重新组批，将 GPU 利用率从 30% 提升到 80% 以上。
LLM 推理优化是一个快速发展的领域，掌握 KV Cache、Flash Attention 和量化技术是构建高性能推理服务的基础。
← 返回首页