AI大模型推理优化：从KV Cache到生产部署

大模型推理分为两个阶段：Prefill 并行处理输入prompt，计算所有token的K、V；Decode 自回归逐token生成。Decode阶段每步只产生一个token，却要与历史全部token做注意力。若每步重算历史K、V，单步复杂度为O(n²)，总复杂度O(n³)。KV Cache的核心思想是：缓存历史K、V，每步只计算新token的Q、K、V，使单步注意力降为O(n)，总复杂度O(n²)。这是所有推理优化的起点。

KV Cache显存占用可估算为：

2 × layers × heads × head_dim × seq_len × batch × dtype_bytes

以7B模型、32层、32头、head_dim=128、FP16、seq_len=4096、batch=1为例，约2GB。序列翻倍，显存翻倍；batch增大，线性增长。因此长上下文和高并发场景下，KV Cache管理直接决定吞吐上限。

在HuggingFace Transformers中，use_cache=True即可启用：

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-7B-Instruct", torch_dtype=torch.float16, device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct")

inputs = tokenizer("解释KV Cache的原理", return_tensors="pt").to(model.device)
with torch.no_grad():
    out = model(**inputs, use_cache=True)
    past = out.past_key_values
    next_token = out.logits[:, -1].argmax(-1, keepdim=True)
    out2 = model(next_token, past_key_values=past, use_cache=True)

past_key_values是一个按层组织的元组，每层包含K和V张量，形状为[batch, heads, seq_len, head_dim]。手工管理时需注意：每生成一个token，seq_len维度追加一列；batch内序列长度不同时，需要padding和attention mask，否则会引入无效计算。

传统KV Cache按请求预分配最大序列长度的连续显存，造成严重内部碎片。vLLM提出的PagedAttention借鉴操作系统虚拟内存分页：将KV Cache切分为固定大小的block，按需分配，逻辑block通过block table映射到物理block。这样不同请求可共享物理块，显存利用率可从不足50%提升到90%以上，并支持前缀共享和连续批处理。生产环境通常直接使用vLLM：

from vllm import LLM, SamplingParams

llm = LLM(
    model="Qwen/Qwen2.5-7B-Instruct",
    quantization="fp8",
    max_model_len=8192,
    gpu_memory_utilization=0.9,
)
params = SamplingParams(temperature=0.7, top_p=0.9, max_tokens=256)
outputs = llm.generate(["解释PagedAttention的优势"], params)
print(outputs[0].outputs[0].text)

量化是另一条主线。FP16权重和KV Cache各占2字节；FP8降至1字节，显存和带宽减半；AWQ/GPTQ采用4bit权重量化，进一步压缩模型体积，但通常保留KV Cache为FP16以保精度。量化收益主要在显存带宽受限的Decode阶段，因为该阶段是memory-bound。实践中需用困惑度、任务准确率和生成质量做回归评估，不能只看显存下降。

连续批处理（Continuous Batching）解决请求长度异构问题：传统静态batch等最长序列结束才释放，短请求浪费大量算力；连续批处理在每个decode步动态加入新请求、移除已完成请求，显著提升GPU利用率。vLLM、TensorRT-LLM、SGLang均内置该机制。

评估推理系统应关注四个指标：TTFT（首token延迟，反映Prefill性能）、TPOT（每输出token时间，反映Decode速度）、吞吐（tokens/s，反映批处理效率）、显存占用。优化时先定位瓶颈：TTFT高查Prefill计算和排队；TPOT高查KV Cache带宽和量化；吞吐低查批处理策略。

综上，大模型推理优化是系统工程：KV Cache消除重复计算，PagedAttention解决显存碎片，量化降低带宽压力，连续批处理提升并发吞吐。四者叠加，才能把实验室模型变成可承载真实流量的生产服务。代码层面，优先使用vLLM等成熟推理引擎，把精力放在请求调度、量化评估和指标监控上，而不是重复实现底层注意力缓存。