大语言模型推理引擎优化：从 KV Cache 到张量并行的完整架构

引言

大语言模型的推理是 AI 应用中计算密度最高的环节。以 Llama-2-70B 为例，生成一个 token 需要执行约 140GB 参数的矩阵乘法，并维护一个随序列长度线性增长的 KV Cache。本文系统讲解推理引擎的核心优化技术，帮助你在有限 GPU 资源下实现极致吞吐。

1. LLM 推理的计算特征

LLM 推理分为两个阶段：

• Prefill 阶段：并行处理全部输入 token，计算密集，受 GPU 算力（FLOPS）限制。
• Decode 阶段：逐个生成 token，内存带宽密集，受 GPU 内存带宽限制。

Decode 阶段每个 token 只需一次 Transformer 前向传播，但每层都要读取全部模型权重（70B 模型 ~140GB），计算强度（FLOPs/Byte）极低，通常是性能瓶颈所在。

2. KV Cache 与 PagedAttention

自回归生成中，每个 token 都要 attend 到之前所有 token 的 Key 和 Value。为避免重复计算，将历史的 K/V 缓存起来即为 KV Cache。

内存占用计算：

kv_cache_per_token = 2 × num_layers × num_kv_heads × head_dim × sizeof(dtype)
示例 (Llama-2-70B, FP16) = 2 × 80 × 8 × 128 × 2 = 320 KB/token

当序列长度为 8K token 时，单个请求的 KV Cache 就占用 2.5GB。PagedAttention（vLLM 的核心创新）借鉴操作系统虚拟内存思想：

• 将 KV Cache 划分为固定大小的「Block」（如 16 token），按需分配。
• 同一请求的不同 Block 不必连续存储，消除了内存碎片。
• 支持 Block 级别的 Copy-on-Write，实现 Beam Search 等场景的 KV Cache 共享。

相比传统连续分配，PagedAttention 将内存浪费从 ~20% 降至不到 4%。

3. 并行策略：TP / PP / DP / CP

张量并行（Tensor Parallelism, TP）：将单层矩阵乘法拆分到多 GPU，每 GPU 持有不同的列/行切片。适用于单机多卡（NVLink 高带宽通信）。通信量：每层两次 AllReduce。

流水线并行（Pipeline Parallelism, PP）：将不同层分配到不同 GPU，通过 micro-batch 流水线化。适用于跨节点部署。通信量：GPU 边界处点对点传输激活值。

数据并行（Data Parallelism, DP）：每个 GPU 持有完整模型副本，处理不同请求。最易于实现，但显存利用率低。

上下文并行（Context Parallelism, CP）：将超长序列拆分到多 GPU，各 GPU 独立处理片段并通过 AllGather 同步注意力。适用于