LLM 推理优化全指南：从 KV Cache 到 Speculative Decoding

         显存优化
            ↑
            |
  计算优化 ←──→ 通信优化

模型权重 (FP16):     14 GB
KV Cache (2K ctx):    1 GB
KV Cache (8K ctx):    4 GB
KV Cache (32K ctx):  16 GB  ← 比模型还大！
KV Cache (128K ctx): 64 GB

总显存 (8K ctx):     18 GB → A100 40G 可以跑
总显存 (32K ctx):    30 GB → A100 40G 勉强
总显存 (128K ctx):   78 GB → 需要 A100 80G

def kv_cache_size(
    num_layers: int,      # 层数
    num_heads: int,       # 注意力头数
    head_dim: int,        # 每头维度
    seq_len: int,         # 序列长度
    batch_size: int,      # 批大小
    dtype_size: int = 2,  # FP16 = 2 bytes
):
    """计算 KV Cache 显存"""
    # K 和 V 各一份
    return 2 * num_layers * num_heads * head_dim * seq_len * batch_size * dtype_size

# Qwen3-7B 示例
size = kv_cache_size(
    num_layers=32,
    num_heads=32,
    head_dim=128,
    seq_len=8192,
    batch_size=1,
)
print(f"{size / 1024**3:.2f} GB")  # 4.00 GB

# 权重 INT8：14GB → 7GB
# KV Cache INT8：显存减半
# 性能损失：1-3%

from transformers import AutoModelForCausalLM
import torch

model = AutoModelForCausalLM.from_pretrained(
    "Qwen3-7B-Instruct",
    load_in_8bit=True,        # 权重 INT8
    device_map="auto",
)

# 权重 INT4：14GB → 3.5GB
# 性能损失：3-8%
# 适合消费级 GPU

model = AutoModelForCausalLM.from_pretrained(
    "Qwen3-7B-Instruct",
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_quant_type="nf4",
)

# AWQ：比 INT4 更好的量化方案
# 保护重要权重，精度损失更小
# vLLM 原生支持

# 量化模型
python -m awq.quantize \
  --model-path /models/Qwen3-7B \
  --output-path /models/Qwen3-7B-AWQ \
  --w-bits 4

# vLLM 加载
llm = LLM(model="/models/Qwen3-7B-AWQ", quantization="awq")

# 传统：预分配最大长度的连续显存
# 浪费率：60%（大量预分配空间未使用）

# PagedAttention：按需分配，像操作系统的分页机制
# 浪费率：<5%

# vLLM 默认启用
llm = LLM(
    model="...",
    gpu_memory_utilization=0.9,  # 95% 显存用于 KV Cache
)

# KV Cache 用 INT8 存储，显存减半
llm = LLM(
    model="...",
    kv_cache_dtype="int8",  # 或 "fp8"
)

# 多个请求共享相同的 system prompt KV Cache
# vLLM 的 prefix caching
llm = LLM(
    model="...",
    enable_prefix_caching=True,
)

# 效果：
# Request 1: "You are a helpful assistant. [Q1]" → 生成 KV Cache
# Request 2: "You are a helpful assistant. [Q2]" → 复用 KV Cache
# 节省：system prompt 部分的计算和显存

# Flash Attention：减少 GPU 显存读写
# 标准 Attention：O(n²) 显存
# Flash Attention：O(n) 显存，速度提升 2-4x

# vLLM 默认启用
llm = LLM(
    model="...",
    enforce_eager=False,  # 启用 CUDA Graph + Flash Attention
)

# CUDA Graph：预编译计算图，减少 kernel launch 开销
# 适合：固定形状的推理（批处理）

llm = LLM(
    model="...",
    enforce_eager=False,  # 启用 CUDA Graph
)
# 延迟降低 10-20%

# 传统批处理：等所有请求完成才处理下一批
# Continuous Batching：每个 step 动态加入/移除请求

# vLLM 默认使用 continuous batching
# 效果：吞吐量提升 3-10x

# 原理：
# Time 0: [A, B, C] 开始生成
# Time 5: A 完成，D 加入 → [B, C, D]
# Time 8: B 完成，E 加入 → [C, D, E]
# GPU 始终满载

# 原理：用小模型猜，大模型验
# 小模型生成 5 个 token → 大模型一次验证 5 个
# 如果全部正确：5x 加速
# 如果 3 个正确：3x 加速

llm = LLM(
    model="/models/Qwen3-7B-Instruct",
    speculative_model="/models/Qwen3-0.5B-Instruct",
    num_speculative_tokens=5,
)

# 实测效果：
# 简单任务（翻译、摘要）：2.0x 加速
# 复杂任务（推理、创作）：1.3x 加速
# 代码生成：1.8x 加速

# MHA：每个 head 有独立的 K, V → KV Cache 大
# MQA：所有 head 共享 K, V → KV Cache 小，但质量降
# GQA：折中，g 个 head 共享 K, V

# Qwen3 使用 GQA，KV Cache 减少为 MHA 的 1/4
# 显存节省 75%，质量损失 <1%

# 大模型蒸馏到小模型
# Teacher: Qwen3-72B → Student: Qwen3-7B
# 保留 90% 能力，推理快 10x，显存省 10x

# 适用场景：固定任务，不需要通用能力

import time
import asyncio

async def benchmark(llm, prompts, max_tokens=512):
    start = time.time()
    
    results = await asyncio.gather(*[
        llm.generate(p, SamplingParams(max_tokens=max_tokens))
        for p in prompts
    ])
    
    elapsed = time.time() - start
    total_tokens = sum(len(r.outputs[0].token_ids) for r in results)
    
    return {
        "total_time": elapsed,
        "total_tokens": total_tokens,
        "throughput": total_tokens / elapsed,  # tokens/s
        "avg_latency": elapsed / len(prompts),
    }

# 配置决策树
def recommend_config(gpu, model_size, concurrency):
    if gpu == "A100 80G":
        if model_size <= 7:
            return {"tp": 1, "quant": "fp16", "max_seqs": 256}
        elif model_size <= 72:
            return {"tp": 1, "quant": "awq", "max_seqs": 64}
    elif gpu == "4090 24G":
        if model_size <= 7:
            return {"tp": 1, "quant": "awq", "max_seqs": 32}
        elif model_size <= 72:
            return {"tp": 2, "quant": "awq", "max_seqs": 16}
    
    return {"error": "configuration not supported"}