推理优化

推理的计算特征

LLM 推理和训练的计算特征差异很大。推理时，模型的内存带宽往往比算力更早成为瓶颈。

推理的两阶段：

1. Prefill 阶段：处理输入序列，并行计算所有 token 的注意力，计算密集型
2. Decode 阶段：逐 token 生成输出，每次只生成一个 token，内存带宽密集型

KV Cache

原理

在自注意力计算中，每个 token 需要与之前所有 token 的 Key 和 Value 做注意力计算。如果不缓存，每次生成新 token 都要重新计算之前所有 token 的 KV，时间复杂度是 O(n²)。

KV Cache 将之前 token 的 K 和 V 矩阵缓存在显存中，每次只需计算新 token 的 K 和 V，然后追加到缓存中。

# KV Cache 简化示意
class KVCache:
    def __init__(self, max_batch, max_seq_len, num_layers, num_heads, head_dim):
        self.cache = {
            layer: {
                "k": torch.zeros(max_batch, num_heads, max_seq_len, head_dim),
                "v": torch.zeros(max_batch, num_heads, max_seq_len, head_dim)
            }
            for layer in range(num_layers)
        }
        self.current_len = 0

    def update(self, layer, k, v):
        """追加新的 K, V 到缓存"""
        seq_len = k.shape[-2]
        self.cache[layer]["k"][:, :, self.current_len:self.current_len+seq_len] = k
        self.cache[layer]["v"][:, :, self.current_len:self.current_len+seq_len] = v
        self.current_len += seq_len

    def get(self, layer):
        """返回当前缓存的 K, V"""
        return (self.cache[layer]["k"][:, :, :self.current_len],
                self.cache[layer]["v"][:, :, :self.current_len])

显存占用

KV Cache 是推理时最大的显存消耗源之一。对于一个 7B 模型（32 层，32 头，128 维）：

单 token KV = 2 × 32 × 32 × 128 × 2 bytes (fp16) = 512 KB/token
2048 token 上下文 = 512 KB × 2048 = 1 GB
4096 token 上下文 = 512 KB × 4096 = 2 GB

优化技术

• MQA/GQA：多个注意力头共享 KV，大幅减少缓存量。LLaMA 2 70B 使用 GQA
• PagedAttention：类似操作系统分页，解决 KV 缓存碎片化问题。vLLM 的核心创新
• Prefix Caching：系统提示词等公共前缀的 KV 可跨请求共享

模型量化

量化是将模型权重和激活值从高精度格式（fp16）映射到低精度格式（int8/int4/nf4），以降低显存占用和带宽需求。

量化方法对比

GPTQ

GPTQ 基于后训练量化，使用少量校准数据选择最优量化参数。

# 使用 AutoGPTQ 量化模型
python -m auto_gptq.quantize \
  --model_name Qwen/Qwen2.5-7B \
  --dataset c4 \
  --bits 4 \
  --group_size 128 \
  --desc_act True \
  --output_dir ./qwen-7b-gptq-4bit

GGUF

GGUF 是 llama.cpp 的量化格式，特别适合在 CPU 和 Apple Silicon 上运行。

# 使用 llama.cpp 量化
./quantize \
  ./models/qwen-7b-fp16.gguf \
  ./models/qwen-7b-q4_k_m.gguf \
  Q4_K_M

量化后效果对比：

Qwen 2.5 7B 在不同量化下的显存需求和加速比：
┌──────────┬──────────┬──────────┐
│ 量化格式  │ 显存占用  │ 速度加速  │
├──────────┼──────────┼──────────┤
│ FP16     │ ~14 GB   │ 1.0x     │
│ GPTQ-4bit│ ~4.5 GB  │ 1.3x     │
│ GGUF Q4  │ ~4.2 GB  │ 0.8x*    │
└──────────┴──────────┴──────────┘
* GGUF 在 GPU 上受限于 offloading 开销

部署框架

vLLM

vLLM 是目前最流行的 LLM 推理引擎，核心优势是 PagedAttention 和高效批处理。

from vllm import LLM, SamplingParams

llm = LLM(
    model="Qwen/Qwen2.5-7B",
    tensor_parallel_size=1,
    dtype="half",
    gpu_memory_utilization=0.9,
    max_model_len=8192,
    enable_prefix_caching=True
)

sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.9,
    max_tokens=1024
)

outputs = llm.generate(
    ["用 Python 实现快速排序", "解释一下量子计算"],
    sampling_params
)

for output in outputs:
    print(f"输入：{output.prompt}")
    print(f"输出：{output.outputs[0].text}")

vLLM 也支持 OpenAI 兼容的 API 服务：

python -m vllm.entrypoints.openai.api_server \
  --model Qwen/Qwen2.5-7B \
  --port 8000 \
  --gpu-memory-utilization 0.9

TGI（Text Generation Inference）

TGI 是 Hugging Face 推出的推理服务器，与 Hugging Face 生态深度集成。

docker run --gpus all \
  -p 8080:80 \
  -v ~/.cache/huggingface:/data \
  ghcr.io/huggingface/text-generation-inference:2.0 \
  --model-id Qwen/Qwen2.5-7B \
  --max-total-tokens 8192 \
  --quantize gptq

框架对比