大模型推理优化与加速完全教程
一、概述
大语言模型（LLM）的推理性能直接影响用户体验和服务成本。一个70B参数的模型，如果未经优化，在普通GPU上可能需要数秒才能生成一个完整的回答。这对于实时对话、代码补全等低延迟场景是不可接受的。
本教程将深入讲解大模型推理的性能瓶颈，系统介绍各种优化技术，包括量化、KV Cache优化、投机解码、批处理优化、模型并行等，并提供实际的代码示例和部署方案。无论你是要部署一个7B的小模型还是数百B的大模型，本教程都能为你提供有价值的指导。
1.1 推理性能的关键指标
在开始优化之前，我们需要明确衡量推理性能的关键指标：
首Token延迟（Time to First Token, TTFT）：从请求发送到生成第一个token的时间。这主要取决于prefill阶段的计算量。对于对话应用，TTFT直接影响用户感知的响应速度。
吞吐量（Throughput）：单位时间内生成的token数量，通常用tokens/s表示。高吞吐量意味着服务器能同时服务更多用户。
每Token延迟（Time Per Output Token, TPOT）：生成每个token所需的时间。这决定了用户看到文字"流出"的速度。
每请求延迟（Latency）：完成整个请求的总时间，等于TTFT + TPOT × 输出token数。
并发数（Concurrency）：系统能同时处理的请求数量。
成本效率：每百万token的成本，通常用$/M tokens表示。
1.2 推理的两个阶段
LLM推理分为两个截然不同的阶段：
Prefill阶段（预填充）：处理输入prompt的所有token，计算注意力并生成KV Cache。这个阶段是计算密集型（compute-bound），可以高度并行化。
Decode阶段（解码）：逐个生成输出token，每步只处理一个新token，但需要访问之前所有token的KV Cache。这个阶段是内存带宽密集型（memory-bound），因为需要频繁读取大量的KV Cache数据。
理解这两个阶段的特性对于优化策略的选择至关重要。
二、推理性能瓶颈分析
2.1 计算瓶颈
大模型推理的计算瓶颈主要来自矩阵乘法运算。对于一个标准的Transformer层，主要计算包括：
- 注意力计算：Q、K、V的线性投影和注意力分数计算
- 前馈网络（FFN）：两个线性层和激活函数
- 输出投影：将隐藏状态映射到词表大小
import torch
import time
def profile_transformer_layer(hidden_size=4096, num_heads=32, ffn_size=11008, seq_len=2048, batch_size=1):
"""分析Transformer层的计算量"""
# QKV投影: 3 * (hidden_size * hidden_size) 次乘法
qkv_flops = 3 * 2 * batch_size * seq_len * hidden_size * hidden_size
# 注意力分数: batch * heads * seq_len * seq_len * head_dim
head_dim = hidden_size // num_heads
attention_flops = 2 * batch_size * num_heads * seq_len * seq_len * head_dim
# FFN: 两个线性层
ffn_flops = 2 * 2 * batch_size * seq_len * hidden_size * ffn_size
total_flops = qkv_flops + attention_flops + ffn_flops
print(f"序列长度: {seq_len}")
print(f"QKV投影: {qkv_flops/1e9:.2f} GFLOPS")
print(f"注意力计算: {attention_flops/1e9:.2f} GFLOPS")
print(f"FFN计算: {ffn_flops/1e9:.2f} GFLOPS")
print(f"总计算量: {total_flops/1e9:.2f} GFLOPS")
# 对比不同序列长度
for sl in [512, 1024, 2048, 4096, 8192]:
attn = 2 * batch_size * num_heads * sl * sl * head_dim
total = qkv_flops * (sl/seq_len) + attn + ffn_flops * (sl/seq_len)
print(f" 序列长度 {sl}: 总计算量 {total/1e9:.2f} GFLOPS")
profile_transformer_layer()
2.2 内存瓶颈
内存瓶颈是大模型推理的主要限制因素，包括：
模型权重：一个70B参数的模型，以FP16存储需要约140GB显存。
KV Cache：随着序列长度和批大小的增加，KV Cache占用的显存快速增长。
激活值：中间计算结果需要临时存储。
def estimate_memory_requirements(
model_params_b: float,
precision_bytes: int = 2, # FP16
seq_len: int = 2048,
batch_size: int = 1,
num_layers: int = 80,
hidden_size: int = 8192,
num_kv_heads: int = 8,
):
"""估算模型推理的显存需求"""
# 模型权重
model_memory_gb = model_params_b * precision_bytes / 1.0 # GB
# KV Cache (每个token的KV cache大小)
# 每层: 2(K+V) * hidden_size * precision_bytes
kv_per_token = 2 * num_layers * hidden_size * precision_bytes # bytes
kv_cache_gb = batch_size * seq_len * kv_per_token / 1e9
# 激活值（粗略估计）
activation_gb = batch_size * seq_len * hidden_size * num_layers * precision_bytes * 4 / 1e9
# 运行时开销
runtime_overhead_gb = 1.0
total_gb = model_memory_gb + kv_cache_gb + activation_gb + runtime_overhead_gb
print(f"模型参数: {model_params_b}B")
print(f"精度: FP{precision_bytes * 8}")
print(f"序列长度: {seq_len}, 批大小: {batch_size}")
print(f"---")
print(f"模型权重: {model_memory_gb:.1f} GB")
print(f"KV Cache: {kv_cache_gb:.1f} GB")
print(f"激活值: {activation_gb:.1f} GB")
print(f"运行时开销: {runtime_overhead_gb:.1f} GB")
print(f"---")
print(f"总计: {total_gb:.1f} GB")
return total_gb
# 不同模型的显存需求
print("=== 7B 模型 ===")
estimate_memory_requirements(7, num_layers=32, hidden_size=4096, num_kv_heads=32)
print("\n=== 70B 模型 ===")
estimate_memory_requirements(70, num_layers=80, hidden_size=8192, num_kv_heads=8)
print("\n=== 70B 模型 (INT8量化) ===")
estimate_memory_requirements(70, precision_bytes=1, num_layers=80, hidden_size=8192, num_kv_heads=8)
2.3 IO瓶颈
IO瓶颈主要体现在：
- GPU显存带宽：decode阶段需要频繁读取KV Cache，受限于显存带宽
- CPU-GPU数据传输：批处理调度和结果返回需要PCIe传输
- 磁盘IO：模型加载和日志写入
def analyze_bandwidth_bottleneck(
model_params_b: float,
batch_size: int,
precision_bytes: int = 2,
gpu_bandwidth_gbps: float = 2048, # A100: 2TB/s
):
"""分析带宽瓶颈"""
# 每生成一个token需要读取的权重数据量
weight_read_gb = model_params_b * precision_bytes / 1e3 # GB
# 理论最大吞吐量 = 带宽 / 每token数据量
max_tokens_per_sec = gpu_bandwidth_gbps * 1e9 / (weight_read_gb * 1e9 / batch_size)
print(f"模型大小: {model_params_b}B, 精度: FP{precision_bytes*8}")
print(f"批大小: {batch_size}")
print(f"每token需读取: {weight_read_gb:.1f} GB")
print(f"GPU带宽: {gpu_bandwidth_gbps} GB/s")
print(f"理论最大吞吐: {max_tokens_per_sec:.0f} tokens/s")
return max_tokens_per_sec
print("=== 不同配置的带宽瓶颈分析 ===")
for bs in [1, 4, 16, 64]:
analyze_bandwidth_bottleneck(70, bs)
print()
三、量化技术
量化是降低模型精度以减少内存占用和计算量的技术，是大模型推理优化最常用的手段之一。
3.1 量化基础
import torch
import torch.nn as nn
import numpy as np
class QuantizationDemo:
"""量化技术演示"""
@staticmethod
def symmetric_quantize(tensor: torch.Tensor, bits: int = 8) -> tuple:
"""对称量化"""
qmin = -(2 ** (bits - 1))
qmax = 2 ** (bits - 1) - 1
# 计算缩放因子
abs_max = tensor.abs().max()
scale = abs_max / qmax
# 量化
quantized = torch.round(tensor / scale).clamp(qmin, qmax).to(torch.int8)
return quantized, scale
@staticmethod
def asymmetric_quantize(tensor: torch.Tensor, bits: int = 8) -> tuple:
"""非对称量化"""
qmin = 0
qmax = 2 ** bits - 1
# 计算缩放因子和零点
min_val = tensor.min()
max_val = tensor.max()
scale = (max_val - min_val) / (qmax - qmin)
zero_point = torch.round(-min_val / scale).clamp(qmin, qmax).to(torch.int)
# 量化
quantized = torch.round(tensor / scale + zero_point).clamp(qmin, qmax).to(torch.uint8)
return quantized, scale, zero_point
@staticmethod
def per_channel_quantize(weight: torch.Tensor, bits: int = 8, axis: int = 0) -> tuple:
"""分通道量化 - 更精确的量化方式"""
qmin = -(2 ** (bits - 1))
qmax = 2 ** (bits - 1) - 1
# 按通道计算缩放因子
shape = [1] * weight.ndim
shape[axis] = weight.shape[axis]
abs_max = weight.abs().amax(dim=[i for i in range(weight.ndim) if i != axis], keepdim=True)
scale = abs_max / qmax
quantized = torch.round(weight / scale).clamp(qmin, qmax).to(torch.int8)
return quantized, scale
@staticmethod
def dequantize(quantized: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor = None) -> torch.Tensor:
"""反量化"""
if zero_point is not None:
return (quantized.float() - zero_point) * scale
return quantized.float() * scale
@staticmethod
def calculate_quantization_error(original: torch.Tensor, quantized: torch.Tensor, scale: torch.Tensor) -> dict:
"""计算量化误差"""
dequantized = quantized.float() * scale
mse = ((original - dequantized) ** 2).mean().item()
max_error = (original - dequantized).abs().max().item()
cos_sim = torch.nn.functional.cosine_similarity(
original.flatten().unsqueeze(0),
dequantized.flatten().unsqueeze(0)
).item()
return {
"mse": mse,
"max_error": max_error,
"cosine_similarity": cos_sim,
}
# 演示不同量化方式的效果
demo = QuantizationDemo()
original = torch.randn(1024, 1024)
# INT8对称量化
q8, s8 = demo.symmetric_quantize(original, bits=8)
error8 = demo.calculate_quantization_error(original, q8, s8)
print(f"INT8量化误差: MSE={error8['mse']:.6f}, 余弦相似度={error8['cosine_similarity']:.6f}")
# INT4量化
q4, s4 = demo.symmetric_quantize(original, bits=4)
error4 = demo.calculate_quantization_error(original, q4, s4)
print(f"INT4量化误差: MSE={error4['mse']:.6f}, 余弦相似度={error4['cosine_similarity']:.6f}")
# 分通道INT8量化
q8_ch, s8_ch = demo.per_channel_quantize(original, bits=8)
error8_ch = demo.calculate_quantization_error(original, q8_ch, s8_ch)
print(f"分通道INT8量化误差: MSE={error8_ch['mse']:.6f}, 余弦相似度={error8_ch['cosine_similarity']:.6f}")
3.2 GPTQ量化
GPTQ是一种基于近似二阶信息的训练后量化方法，能够在INT4精度下保持较高的模型质量。
# GPTQ量化示例（使用auto-gptq库）
# pip install auto-gptq
from transformers import AutoTokenizer
# from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
def gptq_quantize_model(
model_name: str,
output_dir: str,
bits: int = 4,
group_size: int = 128,
desc_act: bool = True,
):
"""
使用GPTQ量化模型
参数:
model_name: 原始模型名称或路径
output_dir: 量化后模型保存路径
bits: 量化位数 (2, 3, 4, 8)
group_size: 分组大小 (32, 64, 128, -1表示不分组)
desc_act: 是否按激活值降序排列
"""
# 加载tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
# 配置量化参数
# quantize_config = BaseQuantizeConfig(
# bits=bits,
# group_size=group_size,
# desc_act=desc_act,
# damp_percent=0.01,
# )
# 加载模型
# model = AutoGPTQForCausalLM.from_pretrained(
# model_name,
# quantize_config=quantize_config,
# trust_remote_code=True,
# )
# 准备校准数据（使用少量数据即可）
# calibration_data = load_calibration_data(tokenizer, num_samples=128)
# 执行量化
# model.quantize(calibration_data)
# 保存量化后的模型
# model.save_quantized(output_dir)
# tokenizer.save_pretrained(output_dir)
print(f"模型已量化并保存到: {output_dir}")
print(f"量化配置: bits={bits}, group_size={group_size}")
# GPTQ量化的关键参数说明:
# - bits: 4bit是最常用的设置，在精度和性能间取得良好平衡
# - group_size: 越小越精确但速度越慢，128是常用值
# - desc_act: True通常能获得更好的量化质量
# - damp_percent: 量化过程中的阻尼系数，处理奇异矩阵
def load_calibration_data(tokenizer, num_samples=128):
"""加载校准数据"""
# 实际使用时应加载有代表性的数据
# 例如: WikiText, C4, 或者业务相关的数据
samples = []
for _ in range(num_samples):
text = "这是一段用于校准的示例文本。" * 10
tokens = tokenizer(text, return_tensors="pt", truncation=True, max_length=2048)
samples.append(tokens)
return samples
3.3 AWQ量化
AWQ（Activation-aware Weight Quantization）通过保护重要权重通道来提高量化质量。
# AWQ量化示例
# pip install awq
def awq_quantize_model(
model_name: str,
output_dir: str,
quant_config: dict = None,
):
"""
使用AWQ量化模型
AWQ的核心思想:
1. 识别对模型输出影响最大的权重通道（通过分析激活值分布）
2. 对这些重要通道进行特殊处理（缩放后再量化）
3. 从而在低比特量化下保持更好的精度
"""
if quant_config is None:
quant_config = {
"zero_point": True,
"q_group_size": 128,
"w_bit": 4,
"version": "GEMM", # GEMM或GEMV
}
# from awq import AutoAWQForCausalLM
# from transformers import AutoTokenizer
# 加载模型和tokenizer
# model = AutoAWQForCausalLM.from_pretrained(model_name)
# tokenizer = AutoTokenizer.from_pretrained(model_name)
# 准备校准数据
# calibration_data = prepare_calibration_data(tokenizer)
# 执行量化
# model.quantize(tokenizer, quant_config=quant_config)
# 保存
# model.save_quantized(output_dir)
# tokenizer.save_pretrained(output_dir)
print(f"AWQ量化完成: {output_dir}")
print(f"量化配置: {quant_config}")
# AWQ vs GPTQ 对比:
# AWQ优势:
# - 量化速度更快（不需要逐层Hessian计算）
# - 对outlier处理更好
# - 推理速度通常更快
#
# GPTQ优势:
# - 在某些模型上精度更高
# - 社区支持更广泛
# - 生态更成熟
3.4 GGUF格式与llama.cpp
GGUF是llama.cpp使用的量化格式，特别适合CPU推理和边缘设备部署。
# GGUF量化参数说明
GGUF_QUANT_TYPES = {
"Q2_K": {"bits": 2.5, "description": "2-bit量化，最小体积，质量损失明显"},
"Q3_K_S": {"bits": 3.0, "description": "3-bit量化（小），体积小"},
"Q3_K_M": {"bits": 3.5, "description": "3-bit量化（中），平衡选择"},
"Q3_K_L": {"bits": 3.75, "description": "3-bit量化（大），质量稍好"},
"Q4_0": {"bits": 4.0, "description": "4-bit量化（旧版），速度快"},
"Q4_K_S": {"bits": 4.0, "description": "4-bit量化（小），推荐"},
"Q4_K_M": {"bits": 4.5, "description": "4-bit量化（中），最推荐"},
"Q5_0": {"bits": 5.0, "description": "5-bit量化（旧版）"},
"Q5_K_S": {"bits": 5.0, "description": "5-bit量化（小）"},
"Q5_K_M": {"bits": 5.5, "description": "5-bit量化（中），高质量"},
"Q6_K": {"bits": 6.0, "description": "6-bit量化，接近原始质量"},
"Q8_0": {"bits": 8.0, "description": "8-bit量化，几乎无损"},
}
def recommend_quantization(model_size_b: float, available_memory_gb: float, use_case: str) -> str:
"""根据场景推荐量化方案"""
# 估算不同量化下的内存需求
for quant_type, info in GGUF_QUANT_TYPES.items():
memory_needed = model_size_b * info["bits"] / 8 # GB
if memory_needed <= available_memory_gb * 0.85: # 留15%余量
if use_case == "quality":
# 质量优先：选择最高精度
if info["bits"] >= 5.0:
return f"{quant_type} (内存需求: {memory_needed:.1f}GB, 质量优先)"
elif use_case == "speed":
# 速度优先：选择4-bit
if 4.0 <= info["bits"] <= 4.5:
return f"{quant_type} (内存需求: {memory_needed:.1f}GB, 速度优先)"
elif use_case == "balanced":
# 平衡：Q4_K_M
if quant_type == "Q4_K_M":
return f"{quant_type} (内存需求: {memory_needed:.1f}GB, 平衡选择)"
return "可用显存不足，建议使用更小的模型或更多GPU"
# 使用示例
print("=== 7B模型，16GB显存 ===")
print(recommend_quantization(7, 16, "balanced"))
print("\n=== 70B模型，48GB显存 ===")
print(recommend_quantization(70, 48, "balanced"))
print("\n=== 70B模型，80GB显存 ===")
print(recommend_quantization(70, 80, "quality"))
3.5 动态量化与混合精度
import torch
from torch.quantization import quantize_dynamic
class HybridPrecisionManager:
"""混合精度管理器 - 对不同层使用不同精度"""
def __init__(self, model, config: dict = None):
self.model = model
self.config = config or self._default_config()
def _default_config(self) -> dict:
"""默认精度配置"""
return {
"attention": "int8", # 注意力层使用INT8
"ffn_up": "int4", # FFN上投影使用INT4
"ffn_down": "int4", # FFN下投影使用INT4
"ffn_gate": "int4", # FFN门控使用INT4
"embeddings": "fp16", # 嵌入层保持FP16
"lm_head": "fp16", # 输出头保持FP16
"layer_norm": "fp32", # LayerNorm保持FP32
}
def apply_hybrid_precision(self):
"""应用混合精度量化"""
for name, module in self.model.named_modules():
precision = self._get_layer_precision(name)
if precision == "int8":
self._quantize_int8(module)
elif precision == "int4":
self._quantize_int4(module)
elif precision == "fp16":
module.half()
# fp32保持不变
def _get_layer_precision(self, layer_name: str) -> str:
"""根据层名确定精度"""
for pattern, precision in self.config.items():
if pattern in layer_name:
return precision
return "fp16" # 默认FP16
def _quantize_int8(self, module):
"""INT8量化"""
if isinstance(module, torch.nn.Linear):
# 使用动态量化
quantized = torch.quantization.quantize_dynamic(
module, {torch.nn.Linear}, dtype=torch.qint8
)
return quantized
def _quantize_int4(self, module):
"""INT4量化（需要专门的INT4实现）"""
# INT4量化通常需要自定义CUDA kernel
pass
def demonstrate_memory_savings():
"""演示不同精度的内存节省"""
sizes = {
"FP32": 4,
"FP16": 2,
"INT8": 1,
"INT4": 0.5,
"INT2": 0.25,
}
model_params_b = 70
print(f"模型参数量: {model_params_b}B")
print("-" * 40)
for precision, bytes_per_param in sizes.items():
memory_gb = model_params_b * bytes_per_param
savings = (1 - bytes_per_param / 4) * 100
print(f"{precision:6s}: {memory_gb:6.1f} GB (节省 {savings:.0f}%)")
demonstrate_memory_savings()
四、KV Cache优化
KV Cache是LLM推理中最重要的优化技术之一。在自回归生成过程中，每个新token的生成都需要访问之前所有token的Key和Value，KV Cache避免了重复计算。
4.1 KV Cache基础
import torch
import torch.nn as nn
import math
class KVCache:
"""基础KV Cache实现"""
def __init__(self, num_layers: int, num_heads: int, head_dim: int, max_seq_len: int, dtype=torch.float16):
self.num_layers = num_layers
self.num_heads = num_heads
self.head_dim = head_dim
self.max_seq_len = max_seq_len
# 预分配KV Cache
self.k_cache = torch.zeros(
num_layers, 1, num_heads, max_seq_len, head_dim, dtype=dtype
)
self.v_cache = torch.zeros(
num_layers, 1, num_heads, max_seq_len, head_dim, dtype=dtype
)
self.current_len = 0
def update(self, layer_idx: int, new_k: torch.Tensor, new_v: torch.Tensor):
"""更新KV Cache"""
seq_len = new_k.shape[2]
self.k_cache[layer_idx, :, :, self.current_len:self.current_len + seq_len] = new_k
self.v_cache[layer_idx, :, :, self.current_len:self.current_len + seq_len] = new_v
if layer_idx == self.num_layers - 1:
self.current_len += seq_len
def get(self, layer_idx: int) -> tuple:
"""获取指定层的KV Cache"""
return (
self.k_cache[layer_idx, :, :, :self.current_len],
self.v_cache[layer_idx, :, :, :self.current_len]
)
def memory_usage(self) -> float:
"""计算KV Cache的显存占用（GB）"""
# 每个元素的字节数
element_size = 2 if self.k_cache.dtype == torch.float16 else 4
total_elements = self.k_cache.numel() + self.v_cache.numel()
return total_elements * element_size / 1e9
# KV Cache内存分析
def analyze_kv_cache_memory():
"""分析不同配置下的KV Cache内存占用"""
configs = [
{"name": "7B", "layers": 32, "heads": 32, "head_dim": 128, "kv_heads": 32},
{"name": "13B", "layers": 40, "heads": 40, "head_dim": 128, "kv_heads": 40},
{"name": "70B", "layers": 80, "heads": 64, "head_dim": 128, "kv_heads": 8},
]
seq_lengths = [512, 1024, 2048, 4096, 8192]
batch_sizes = [1, 4, 16, 64]
for config in configs:
print(f"\n=== {config['name']} 模型 ===")
for bs in batch_sizes:
for sl in seq_lengths:
# KV Cache大小 = 2(K+V) * layers * batch * seq_len * kv_heads * head_dim * bytes
kv_size_gb = 2 * config['layers'] * bs * sl * config['kv_heads'] * config['head_dim'] * 2 / 1e9
print(f" batch={bs:3d}, seq_len={sl:5d}: KV Cache = {kv_size_gb:6.2f} GB")
analyze_kv_cache_memory()
4.2 PagedAttention（vLLM核心优化）
PagedAttention是vLLM引入的核心优化技术，通过类似操作系统虚拟内存分页的机制来管理KV Cache，大幅减少内存碎片和浪费。
import torch
from typing import List, Optional, Tuple
class PagedKVCache:
"""
PagedAttention的KV Cache实现
核心思想:
1. 将KV Cache分成固定大小的"页"（block）
2. 使用页表（block table）管理逻辑到物理的映射
3. 支持非连续存储，减少内存碎片
4. 支持copy-on-write，优化beam search等场景
"""
def __init__(
self,
num_layers: int,
num_kv_heads: int,
head_dim: int,
block_size: int = 16,
num_blocks: int = 1000,
dtype=torch.float16,
):
self.num_layers = num_layers
self.num_kv_heads = num_kv_heads
self.head_dim = head_dim
self.block_size = block_size
self.num_blocks = num_blocks
self.dtype = dtype
# 物理块池: [num_layers, 2, num_blocks, block_size, num_kv_heads, head_dim]
# 2 表示 K 和 V
self.block_pool = torch.zeros(
num_layers, 2, num_blocks, block_size, num_kv_heads, head_dim,
dtype=dtype
)
# 空闲块列表
self.free_blocks = list(range(num_blocks))
# 每个序列的页表: seq_id -> list of physical block indices
self.block_tables = {}
# 每个序列的当前长度
self.seq_lengths = {}
def allocate_sequence(self, seq_id: int) -> List[int]:
"""为新序列分配初始块"""
if not self.free_blocks:
raise RuntimeError("没有可用的空闲块")
block_idx = self.free_blocks.pop(0)
self.block_tables[seq_id] = [block_idx]
self.seq_lengths[seq_id] = 0
return [block_idx]
def append_token(self, seq_id: int, layer_idx: int, k: torch.Tensor, v: torch.Tensor):
"""
追加一个token的KV到缓存
参数:
seq_id: 序列ID
layer_idx: 层索引
k: [1, num_kv_heads, 1, head_dim] 的Key张量
v: [1, num_kv_heads, 1, head_dim] 的Value张量
"""
if seq_id not in self.block_tables:
self.allocate_sequence(seq_id)
current_len = self.seq_lengths[seq_id]
block_idx_in_seq = current_len // self.block_size
offset_in_block = current_len % self.block_size
# 需要新块
if block_idx_in_seq >= len(self.block_tables[seq_id]):
if not self.free_blocks:
raise RuntimeError("没有可用的空闲块")
new_block = self.free_blocks.pop(0)
self.block_tables[seq_id].append(new_block)
# 写入KV
physical_block = self.block_tables[seq_id][block_idx_in_seq]
self.block_pool[layer_idx, 0, physical_block, offset_in_block] = k.squeeze(0).squeeze(1)
self.block_pool[layer_idx, 1, physical_block, offset_in_block] = v.squeeze(0).squeeze(1)
if layer_idx == self.num_layers - 1:
self.seq_lengths[seq_id] += 1
def get_kv(self, seq_id: int, layer_idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
"""获取序列的完整KV Cache"""
block_table = self.block_tables[seq_id]
seq_len = self.seq_lengths[seq_id]
# 收集所有块的KV
k_list = []
v_list = []
for i, physical_block in enumerate(block_table):
start = i * self.block_size
end = min(start + self.block_size, seq_len)
actual_len = end - start
k_list.append(self.block_pool[layer_idx, 0, physical_block, :actual_len])
v_list.append(self.block_pool[layer_idx, 1, physical_block, :actual_len])
k = torch.cat(k_list, dim=0) # [seq_len, num_kv_heads, head_dim]
v = torch.cat(v_list, dim=0)
# 调整形状为 [1, num_kv_heads, seq_len, head_dim]
k = k.unsqueeze(0).transpose(1, 2)
v = v.unsqueeze(0).transpose(1, 2)
return k, v
def free_sequence(self, seq_id: int):
"""释放序列的KV Cache"""
if seq_id in self.block_tables:
self.free_blocks.extend(self.block_tables[seq_id])
del self.block_tables[seq_id]
del self.seq_lengths[seq_id]
def memory_usage(self) -> dict:
"""内存使用统计"""
total_blocks = self.num_blocks
used_blocks = total_blocks - len(self.free_blocks)
# 每个块的字节数
block_bytes = self.block_size * self.num_kv_heads * self.head_dim * 2 * 2 # 2(K+V) * 2(FP16)
return {
"total_blocks": total_blocks,
"used_blocks": used_blocks,
"free_blocks": len(self.free_blocks),
"utilization": used_blocks / total_blocks * 100,
"total_memory_gb": total_blocks * block_bytes * self.num_layers / 1e9,
"used_memory_gb": used_blocks * block_bytes * self.num_layers / 1e9,
}
# 使用示例
cache = PagedKVCache(
num_layers=32,
num_kv_heads=32,
head_dim=128,
block_size=16,
num_blocks=500,
)
# 模拟多个序列
for seq_id in range(5):
cache.allocate_sequence(seq_id)
for token_idx in range(100):
for layer in range(32):
k = torch.randn(1, 32, 1, 128, dtype=torch.float16)
v = torch.randn(1, 32, 1, 128, dtype=torch.float16)
cache.append_token(seq_id, layer, k, v)
print("内存使用统计:", cache.memory_usage())
4.3 Prefix Caching
Prefix Caching（前缀缓存）通过缓存公共前缀的KV Cache来加速多轮对话和共享system prompt的场景。
import hashlib
from typing import Dict, List, Optional, Tuple
class PrefixCacheManager:
"""
前缀缓存管理器
适用场景:
1. 多轮对话共享system prompt
2. 批量请求共享相同的前缀
3. RAG场景中共享检索结果的KV Cache
"""
def __init__(self, kv_cache: PagedKVCache):
self.kv_cache = kv_cache
# 前缀哈希 -> 物理块列表
self.prefix_cache: Dict[str, List[int]] = {}
# 前缀哈希 -> 前缀长度
self.prefix_lengths: Dict[str, int] = {}
# 引用计数
self.ref_counts: Dict[str, int] = {}
def compute_prefix_hash(self, token_ids: List[int]) -> str:
"""计算token序列的哈希"""
return hashlib.sha256(str(token_ids).encode()).hexdigest()[:16]
def get_cached_prefix(self, token_ids: List[int]) -> Optional[Tuple[List[int], int]]:
"""
查找已缓存的前缀
返回: (物理块列表, 缓存的token数量) 或 None
"""
# 尝试从最长前缀开始查找
for length in range(len(token_ids), 0, -1):
prefix_hash = self.compute_prefix_hash(token_ids[:length])
if prefix_hash in self.prefix_cache:
self.ref_counts[prefix_hash] = self.ref_counts.get(prefix_hash, 0) + 1
return self.prefix_cache[prefix_hash], length
return None
def cache_prefix(self, token_ids: List[int], block_table: List[int]):
"""缓存前缀的KV Cache"""
prefix_hash = self.compute_prefix_hash(token_ids)
# 复制块表（copy-on-write语义）
self.prefix_cache[prefix_hash] = block_table.copy()
self.prefix_lengths[prefix_hash] = len(token_ids)
self.ref_counts[prefix_hash] = 1
def release_prefix(self, token_ids: List[int]):
"""释放前缀缓存引用"""
prefix_hash = self.compute_prefix_hash(token_ids)
if prefix_hash in self.ref_counts:
self.ref_counts[prefix_hash] -= 1
if self.ref_counts[prefix_hash] <= 0:
# 可以选择保留或释放物理块
pass
class MultiTurnConversation:
"""多轮对话管理器 - 利用Prefix Caching优化"""
def __init__(self, prefix_manager: PrefixCacheManager):
self.prefix_manager = prefix_manager
self.conversations: Dict[str, dict] = {}
def start_conversation(self, conv_id: str, system_prompt_tokens: List[int]):
"""开始新对话，缓存system prompt"""
cached = self.prefix_manager.get_cached_prefix(system_prompt_tokens)
self.conversations[conv_id] = {
"system_tokens": system_prompt_tokens,
"cached_blocks": cached[0] if cached else None,
"cached_length": cached[1] if cached else 0,
"history": [],
}
if not cached:
# 需要计算并缓存system prompt的KV
# 这里是伪代码，实际需要通过模型计算
pass
def add_turn(self, conv_id: str, user_tokens: List[int], assistant_tokens: List[int]):
"""添加一轮对话"""
self.conversations[conv_id]["history"].append({
"user": user_tokens,
"assistant": assistant_tokens,
})
def get_context_tokens(self, conv_id: str) -> Tuple[List[int], int]:
"""
获取完整的上下文token序列
返回: (token列表, 已缓存的token数量)
"""
conv = self.conversations[conv_id]
# 系统提示
tokens = conv["system_tokens"].copy()
cached_len = conv["cached_length"]
# 历史对话
for turn in conv["history"]:
tokens.extend(turn["user"])
tokens.extend(turn["assistant"])
return tokens, cached_len
def estimate_speedup(self, conv_id: str) -> float:
"""估算Prefix Caching带来的加速比"""
tokens, cached_len = self.get_context_tokens(conv_id)
total_len = len(tokens)
if total_len == 0:
return 1.0
# Prefill时间与token数成正比
speedup = total_len / (total_len - cached_len) if cached_len < total_len else float('inf')
return speedup
# 使用示例
# 假设system prompt有500个token，每轮对话约200个token
system_tokens = list(range(500))
conv_tokens = list(range(200))
print("=== Prefix Caching 加速效果分析 ===")
for turn in range(1, 11):
total_context = 500 + turn * 200 * 2 # system + user + assistant
cached = 500 + (turn - 1) * 200 * 2 # 前面的都已缓存
speedup = total_context / (total_context - cached)
print(f"第{turn}轮对话: 上下文{total_context}tokens, 缓存{cached}tokens, 加速{speedup:.1f}x")
4.4 Multi-Query Attention与Grouped-Query Attention
import torch
import torch.nn as nn
import math
class MultiHeadAttention(nn.Module):
"""标准多头注意力"""
def __init__(self, hidden_size, num_heads):
super().__init__()
self.num_heads = num_heads
self.head_dim = hidden_size // num_heads
self.q_proj = nn.Linear(hidden_size, hidden_size)
self.k_proj = nn.Linear(hidden_size, hidden_size)
self.v_proj = nn.Linear(hidden_size, hidden_size)
self.o_proj = nn.Linear(hidden_size, hidden_size)
def forward(self, x, kv_cache=None):
B, L, D = x.shape
q = self.q_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
k = self.k_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
v = self.v_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
# KV Cache更新
if kv_cache is not None:
k = torch.cat([kv_cache[0], k], dim=2)
v = torch.cat([kv_cache[1], v], dim=2)
# 注意力计算
attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
attn = torch.softmax(attn, dim=-1)
out = torch.matmul(attn, v)
out = out.transpose(1, 2).contiguous().view(B, L, -1)
return self.o_proj(out), (k, v)
class GroupedQueryAttention(nn.Module):
"""
分组查询注意力 (GQA)
多个查询头共享一组KV头，减少KV Cache大小
例如: 32个Q头, 8个KV头 -> KV Cache减少4倍
"""
def __init__(self, hidden_size, num_q_heads, num_kv_heads):
super().__init__()
self.num_q_heads = num_q_heads
self.num_kv_heads = num_kv_heads
self.num_groups = num_q_heads // num_kv_heads
self.head_dim = hidden_size // num_q_heads
self.q_proj = nn.Linear(hidden_size, num_q_heads * self.head_dim)
self.k_proj = nn.Linear(hidden_size, num_kv_heads * self.head_dim)
self.v_proj = nn.Linear(hidden_size, num_kv_heads * self.head_dim)
self.o_proj = nn.Linear(hidden_size, hidden_size)
def forward(self, x, kv_cache=None):
B, L, D = x.shape
q = self.q_proj(x).view(B, L, self.num_q_heads, self.head_dim).transpose(1, 2)
k = self.k_proj(x).view(B, L, self.num_kv_heads, self.head_dim).transpose(1, 2)
v = self.v_proj(x).view(B, L, self.num_kv_heads, self.head_dim).transpose(1, 2)
# KV Cache更新
if kv_cache is not None:
k = torch.cat([kv_cache[0], k], dim=2)
v = torch.cat([kv_cache[1], v], dim=2)
# 扩展KV以匹配Q的头数
k_expanded = k.repeat_interleave(self.num_groups, dim=1)
v_expanded = v.repeat_interleave(self.num_groups, dim=1)
# 注意力计算
attn = torch.matmul(q, k_expanded.transpose(-2, -1)) / math.sqrt(self.head_dim)
attn = torch.softmax(attn, dim=-1)
out = torch.matmul(attn, v_expanded)
out = out.transpose(1, 2).contiguous().view(B, L, -1)
return self.o_proj(out), (k, v) # 返回未扩展的KV用于缓存
class MultiQueryAttention(nn.Module):
"""
多查询注意力 (MQA)
所有Q头共享一个KV头，KV Cache最小
"""
def __init__(self, hidden_size, num_heads):
super().__init__()
self.num_heads = num_heads
self.head_dim = hidden_size // num_heads
self.q_proj = nn.Linear(hidden_size, hidden_size)
self.k_proj = nn.Linear(hidden_size, self.head_dim) # 只有一个KV头
self.v_proj = nn.Linear(hidden_size, self.head_dim)
self.o_proj = nn.Linear(hidden_size, hidden_size)
def forward(self, x, kv_cache=None):
B, L, D = x.shape
q = self.q_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2)
k = self.k_proj(x).view(B, L, 1, self.head_dim).transpose(1, 2)
v = self.v_proj(x).view(B, L, 1, self.head_dim).transpose(1, 2)
if kv_cache is not None:
k = torch.cat([kv_cache[0], k], dim=2)
v = torch.cat([kv_cache[1], v], dim=2)
# 扩展KV
k_expanded = k.expand(-1, self.num_heads, -1, -1)
v_expanded = v.expand(-1, self.num_heads, -1, -1)
attn = torch.matmul(q, k_expanded.transpose(-2, -1)) / math.sqrt(self.head_dim)
attn = torch.softmax(attn, dim=-1)
out = torch.matmul(attn, v_expanded)
out = out.transpose(1, 2).contiguous().view(B, L, -1)
return self.o_proj(out), (k, v)
# KV Cache大小对比
def compare_kv_cache_sizes(hidden_size=4096, num_heads=32, seq_len=4096, batch_size=1):
"""对比不同注意力机制的KV Cache大小"""
head_dim = hidden_size // num_heads
# MHA: 每个头都有独立的KV
mha_kv = 2 * batch_size * num_heads * seq_len * head_dim * 2 # FP16
# GQA (8个KV头)
gqa_kv = 2 * batch_size * 8 * seq_len * head_dim * 2
# MQA: 只有一个KV头
mqa_kv = 2 * batch_size * 1 * seq_len * head_dim * 2
print(f"序列长度: {seq_len}, 批大小: {batch_size}")
print(f"MHA KV Cache: {mha_kv/1e6:.1f} MB")
print(f"GQA KV Cache (8 KV头): {gqa_kv/1e6:.1f} MB (节省 {(1-gqa_kv/mha_kv)*100:.0f}%)")
print(f"MQA KV Cache: {mqa_kv/1e6:.1f} MB (节省 {(1-mqa_kv/mha_kv)*100:.0f}%)")
compare_kv_cache_sizes()
五、投机解码（Speculative Decoding）
投机解码是一种通过使用小型"草稿"模型快速生成候选token，再由大模型并行验证的技术，可以在不损失质量的情况下显著提升推理速度。
5.1 基本原理与实现
import torch
import torch.nn.functional as F
from typing import List, Tuple, Optional
class SpeculativeDecoder:
"""
投机解码器
原理:
1. 使用小模型（draft model）快速生成K个候选token
2. 使用大模型（target model）并行验证这K个token
3. 接受概率匹配的token，拒绝的token从修正分布中重新采样
4. 理论上可以保证输出分布与单独使用大模型完全一致
"""
def __init__(self, target_model, draft_model, tokenizer, K: int = 5):
self.target_model = target_model
self.draft_model = draft_model
self.tokenizer = tokenizer
self.K = K # 每次投机的token数量
def generate(self, prompt: str, max_tokens: int = 100, temperature: float = 0.0) -> str:
"""使用投机解码生成文本"""
input_ids = self.tokenizer.encode(prompt)
generated_ids = []
total_target_calls = 0
total_draft_calls = 0
total_accepted = 0
while len(generated_ids) < max_tokens:
# 第一步：使用draft model快速生成K个候选token
draft_tokens, draft_probs = self._draft_generate(
input_ids + generated_ids, self.K, temperature
)
total_draft_calls += 1
# 第二步：使用target model并行验证
target_probs = self._target_verify(
input_ids + generated_ids + draft_tokens
)
total_target_calls += 1
# 第三步：逐个验证并接受/拒绝
accepted_tokens = []
for i in range(self.K):
draft_token = draft_tokens[i]
draft_prob = draft_probs[i][draft_token]
target_prob = target_probs[i][draft_token]
# 接受概率
if temperature == 0:
# 贪心解码：如果target也选择这个token则接受
target_best = target_probs[i].argmax().item()
if draft_token == target_best:
accepted_tokens.append(draft_token)
total_accepted += 1
else:
# 使用target的token
accepted_tokens.append(target_best)
break
else:
# 采样解码：按概率接受
accept_prob = min(1.0, target_prob / draft_prob)
if torch.rand(1).item() < accept_prob:
accepted_tokens.append(draft_token)
total_accepted += 1
else:
# 从修正分布中采样
corrected_probs = F.relu(target_probs[i] - draft_probs[i])
corrected_probs = corrected_probs / corrected_probs.sum()
new_token = torch.multinomial(corrected_probs, 1).item()
accepted_tokens.append(new_token)
break
generated_ids.extend(accepted_tokens)
# 检查是否生成了结束符
if self.tokenizer.eos_token_id in accepted_tokens:
break
# 计算加速比
avg_accepted = total_accepted / total_draft_calls if total_draft_calls > 0 else 0
speedup = (total_accepted + total_draft_calls) / (total_target_calls + total_draft_calls)
result = self.tokenizer.decode(generated_ids)
print(f"加速统计: 平均每次接受 {avg_accepted:.1f} 个token, 加速比 {speedup:.2f}x")
return result
def _draft_generate(self, input_ids: List[int], K: int, temperature: float) -> Tuple[List[int], List[torch.Tensor]]:
"""使用draft model生成K个token"""
tokens = []
probs = []
current_ids = input_ids.copy()
for _ in range(K):
with torch.no_grad():
outputs = self.draft_model(torch.tensor([current_ids]))
logits = outputs.logits[:, -1, :]
if temperature > 0:
logits = logits / temperature
prob = F.softmax(logits, dim=-1)
token = torch.multinomial(prob, 1).item()
probs.append(prob.squeeze(0))
else:
token = logits.argmax(dim=-1).item()
prob = F.softmax(logits, dim=-1)
probs.append(prob.squeeze(0))
tokens.append(token)
current_ids.append(token)
return tokens, probs
def _target_verify(self, input_ids: List[int]) -> List[torch.Tensor]:
"""使用target model并行验证"""
with torch.no_grad():
outputs = self.target_model(torch.tensor([input_ids]))
logits = outputs.logits
# 获取每个位置的概率分布
# input_ids[K:] 对应的logits在位置 [K-1:-1]
start_pos = len(input_ids) - self.K - 1
probs = F.softmax(logits[:, start_pos:start_pos + self.K, :], dim=-1)
return [probs[0, i] for i in range(self.K)]
# 投机解码的理论分析
def analyze_speculative_speedup(draft_acceptance_rate: float, K: int, draft_speed_ratio: float = 10.0):
"""
分析投机解码的理论加速比
参数:
draft_acceptance_rate: draft model的token被接受的平均概率
K: 每次投机的token数
draft_speed_ratio: draft model相对target model的速度比
"""
# 期望接受的token数
expected_accepted = sum(draft_acceptance_rate ** i for i in range(1, K + 1))
# 每轮的成本（以target model的前向传播为单位）
cost_per_round = 1.0 + K / draft_speed_ratio # 1次target验证 + K次draft生成
# 加速比
speedup = (expected_accepted + 1) / cost_per_round # +1 因为target也会生成1个token
print(f"K={K}, 接受率={draft_acceptance_rate:.0%}, draft速度比={draft_speed_ratio}x")
print(f" 期望接受: {expected_accepted:.1f} tokens")
print(f" 每轮成本: {cost_per_round:.2f}")
print(f" 加速比: {speedup:.2f}x")
return speedup
print("=== 投机解码理论加速比分析 ===")
for acceptance_rate in [0.5, 0.7, 0.8, 0.9]:
for K in [3, 5, 7, 10]:
analyze_speculative_speedup(acceptance_rate, K)
print()
5.2 自投机解码（Self-Speculative Decoding）
class SelfSpeculativeDecoder:
"""
自投机解码 - 使用模型自身的子集作为draft model
方法:
1. 跳过某些层（Layer Skipping）
2. 使用浅层作为draft，深层作为target
3. 不需要额外的draft model
"""
def __init__(self, model, num_layers: int, skip_layers: List[int], K: int = 5):
self.model = model
self.num_layers = num_layers
self.skip_layers = skip_layers # 跳过的层索引
self.K = K
def draft_forward(self, input_ids: torch.Tensor) -> torch.Tensor:
"""使用跳过层的方式进行draft前向传播"""
# 这是伪代码，实际实现需要修改模型的forward
# 跳过指定层以加速
hidden_states = self.model.embed_tokens(input_ids)
for i, layer in enumerate(self.model.layers):
if i not in self.skip_layers:
hidden_states = layer(hidden_states)
logits = self.model.lm_head(hidden_states)
return logits
def target_forward(self, input_ids: torch.Tensor) -> torch.Tensor:
"""使用完整模型进行target前向传播"""
return self.model(input_ids).logits
六、批处理优化
6.1 Continuous Batching（连续批处理）
传统批处理需要等待所有请求完成才能处理下一批，Continuous Batching允许在生成过程中动态添加和移除请求，大幅提高GPU利用率。
import asyncio
from dataclasses import dataclass, field
from typing import List, Dict, Optional
from collections import deque
import time
@dataclass
class InferenceRequest:
request_id: str
prompt_tokens: List[int]
max_new_tokens: int
temperature: float = 0.0
generated_tokens: List[int] = field(default_factory=list)
is_finished: bool = False
arrival_time: float = 0.0
first_token_time: float = 0.0
class ContinuousBatchScheduler:
"""
连续批处理调度器
核心思想:
1. 不等待整个batch完成
2. 当某个请求完成时，立即用新请求替换
3. 始终保持GPU满载
"""
def __init__(self, max_batch_size: int = 32, max_seq_len: int = 4096):
self.max_batch_size = max_batch_size
self.max_seq_len = max_seq_len
self.waiting_queue: deque = deque()
self.running_batch: Dict[str, InferenceRequest] = {}
self.completed: List[InferenceRequest] = []
self.kv_cache_manager = None # 实际的KV Cache管理器
def add_request(self, request: InferenceRequest):
"""添加新请求到等待队列"""
request.arrival_time = time.time()
self.waiting_queue.append(request)
def schedule(self) -> List[InferenceRequest]:
"""
调度下一批请求
策略:
1. 优先保持当前batch运行
2. 有空位时从等待队列添加新请求
3. 请求完成时释放资源
"""
# 移除已完成的请求
finished = [rid for rid, req in self.running_batch.items() if req.is_finished]
for rid in finished:
self.completed.append(self.running_batch.pop(rid))
# 填充空位
while len(self.running_batch) < self.max_batch_size and self.waiting_queue:
new_request = self.waiting_queue.popleft()
# 检查序列长度限制
if len(new_request.prompt_tokens) < self.max_seq_len:
self.running_batch[new_request.request_id] = new_request
return list(self.running_batch.values())
def step(self, model_output: Dict[str, int]):
"""
执行一步推理
参数:
model_output: {request_id: next_token_id}
"""
for request_id, next_token in model_output.items():
if request_id in self.running_batch:
request = self.running_batch[request_id]
request.generated_tokens.append(next_token)
if request.first_token_time == 0:
request.first_token_time = time.time()
# 检查是否完成
if (next_token == 0 or # EOS token
len(request.generated_tokens) >= request.max_new_tokens):
request.is_finished = True
def get_stats(self) -> dict:
"""获取调度统计"""
return {
"waiting": len(self.waiting_queue),
"running": len(self.running_batch),
"completed": len(self.completed),
"batch_utilization": len(self.running_batch) / self.max_batch_size * 100,
}
class AsyncLLMEngine:
"""异步LLM推理引擎"""
def __init__(self, model, tokenizer, max_batch_size: int = 32):
self.model = model
self.tokenizer = tokenizer
self.scheduler = ContinuousBatchScheduler(max_batch_size=max_batch_size)
self.is_running = False
async def generate(self, prompt: str, max_tokens: int = 100, temperature: float = 0.0) -> str:
"""异步生成接口"""
request_id = f"req_{id(prompt)}"
tokens = self.tokenizer.encode(prompt)
request = InferenceRequest(
request_id=request_id,
prompt_tokens=tokens,
max_new_tokens=max_tokens,
temperature=temperature,
)
self.scheduler.add_request(request)
# 等待完成
while not request.is_finished:
await asyncio.sleep(0.01)
return self.tokenizer.decode(request.generated_tokens)
async def run_loop(self):
"""主推理循环"""
self.is_running = True
while self.is_running:
batch = self.scheduler.schedule()
if batch:
# 准备batch输入
input_ids = self._prepare_batch_input(batch)
# 模型推理
with torch.no_grad():
outputs = self.model(input_ids)
# 处理输出
next_tokens = self._process_batch_output(outputs, batch)
self.scheduler.step(next_tokens)
else:
# 没有请求，等待
await asyncio.sleep(0.001)
def _prepare_batch_input(self, batch: List[InferenceRequest]):
"""准备batch输入（需要padding）"""
# 实际实现需要处理不同长度的序列
pass
def _process_batch_output(self, outputs, batch: List[InferenceRequest]):
"""处理batch输出"""
# 实际实现需要从logits中采样下一个token
pass
6.2 Prefill-Decode分离
class ChunkedPrefillScheduler:
"""
分块Prefill调度器
将长prompt的prefill分成多个chunk，
避免单个长请求阻塞整个batch
"""
def __init__(self, chunk_size: int = 512, max_batch_tokens: int = 8192):
self.chunk_size = chunk_size
self.max_batch_tokens = max_batch_tokens
def schedule_prefill(self, requests: List[InferenceRequest]) -> List[List[int]]:
"""
将prefill请求分块调度
策略:
1. 优先处理短请求（SJF - 最短任务优先）
2. 将长请求分成chunk
3. 保证总token数不超过max_batch_tokens
"""
# 按prompt长度排序
sorted_requests = sorted(requests, key=lambda r: len(r.prompt_tokens))
chunks = []
current_chunk = []
current_tokens = 0
for request in sorted_requests:
prompt_len = len(request.prompt_tokens)
if prompt_len <= self.chunk_size:
# 短请求，直接加入当前chunk
if current_tokens + prompt_len <= self.max_batch_tokens:
current_chunk.append(request)
current_tokens += prompt_len
else:
# 当前chunk已满
chunks.append(current_chunk)
current_chunk = [request]
current_tokens = prompt_len
else:
# 长请求，需要分块
for i in range(0, prompt_len, self.chunk_size):
chunk_tokens = min(self.chunk_size, prompt_len - i)
if current_tokens + chunk_tokens > self.max_batch_tokens:
chunks.append(current_chunk)
current_chunk = []
current_tokens = 0
current_tokens += chunk_tokens
current_chunk.append(request)
if current_chunk:
chunks.append(current_chunk)
return chunks
七、模型并行
7.1 张量并行（Tensor Parallelism）
import torch
import torch.nn as nn
import torch.distributed as dist
class TensorParallelLinear(nn.Module):
"""
张量并行线性层
将权重矩阵按列或行切分到多个GPU
"""
def __init__(self, in_features: int, out_features: int,
tp_size: int, tp_rank: int, split_mode: str = "column"):
super().__init__()
self.tp_size = tp_size
self.tp_rank = tp_rank
self.split_mode = split_mode
if split_mode == "column":
# 列切分：每个GPU持有 out_features/tp_size 列
assert out_features % tp_size == 0
self.local_out_features = out_features // tp_size
self.linear = nn.Linear(in_features, self.local_out_features, bias=False)
elif split_mode == "row":
# 行切分：每个GPU持有 in_features/tp_size 行
assert in_features % tp_size == 0
self.local_in_features = in_features // tp_size
self.linear = nn.Linear(self.local_in_features, out_features, bias=False)
def forward(self, x: torch.Tensor) -> torch.Tensor:
if self.split_mode == "column":
# 列切分：每个GPU计算部分输出，无需通信
return self.linear(x)
elif self.split_mode == "row":
# 行切分：每个GPU计算部分结果，需要AllReduce
local_output = self.linear(x)
dist.all_reduce(local_output, op=dist.ReduceOp.SUM)
return local_output
class TensorParallelMLP(nn.Module):
"""张量并行的MLP层"""
def __init__(self, hidden_size: int, ffn_size: int, tp_size: int, tp_rank: int):
super().__init__()
# gate和up投影使用列切分
self.gate_proj = TensorParallelLinear(hidden_size, ffn_size, tp_size, tp_rank, "column")
self.up_proj = TensorParallelLinear(hidden_size, ffn_size, tp_size, tp_rank, "column")
# down投影使用行切分
self.down_proj = TensorParallelLinear(ffn_size, hidden_size, tp_size, tp_rank, "row")
self.act_fn = nn.SiLU()
def forward(self, x: torch.Tensor) -> torch.Tensor:
gate = self.act_fn(self.gate_proj(x))
up = self.up_proj(x)
return self.down_proj(gate * up)
class TensorParallelAttention(nn.Module):
"""张量并行的注意力层"""
def __init__(self, hidden_size: int, num_heads: int, tp_size: int, tp_rank: int):
super().__init__()
self.num_heads = num_heads
self.tp_size = tp_size
self.tp_rank = tp_rank
self.head_dim = hidden_size // num_heads
self.local_heads = num_heads // tp_size
# QKV投影使用列切分
self.q_proj = TensorParallelLinear(hidden_size, hidden_size, tp_size, tp_rank, "column")
self.k_proj = TensorParallelLinear(hidden_size, hidden_size, tp_size, tp_rank, "column")
self.v_proj = TensorParallelLinear(hidden_size, hidden_size, tp_size, tp_rank, "column")
self.o_proj = TensorParallelLinear(hidden_size, hidden_size, tp_size, tp_rank, "row")
def forward(self, x: torch.Tensor, kv_cache=None):
B, L, D = x.shape
q = self.q_proj(x).view(B, L, self.local_heads, self.head_dim).transpose(1, 2)
k = self.k_proj(x).view(B, L, self.local_heads, self.head_dim).transpose(1, 2)
v = self.v_proj(x).view(B, L, self.local_heads, self.head_dim).transpose(1, 2)
# 注意力计算
attn = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)
attn = torch.softmax(attn, dim=-1)
out = torch.matmul(attn, v)
out = out.transpose(1, 2).contiguous().view(B, L, -1)
return self.o_proj(out)
7.2 流水线并行（Pipeline Parallelism）
class PipelineParallelModel:
"""
流水线并行模型
将模型的不同层分配到不同GPU
通过流水线调度减少气泡时间
"""
def __init__(self, layers: List[nn.Module], num_stages: int, stage_id: int):
self.layers = layers
self.num_stages = num_stages
self.stage_id = stage_id
# 当前stage负责的层范围
layers_per_stage = len(layers) // num_stages
self.start_layer = stage_id * layers_per_stage
self.end_layer = (stage_id + 1) * layers_per_stage if stage_id < num_stages - 1 else len(layers)
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
"""执行当前stage的层"""
for i in range(self.start_layer, self.end_layer):
hidden_states = self.layers[i](hidden_states)
return hidden_states
class GPipeScheduler:
"""
GPipe流水线调度器
将一个batch分成多个micro-batch，
交错执行以减少流水线气泡
"""
def __init__(self, num_stages: int, num_micro_batches: int):
self.num_stages = num_stages
self.num_micro_batches = num_micro_batches
def schedule(self) -> List[List[tuple]]:
"""
生成调度计划
返回: 每个时间步每个stage执行的micro-batch索引
"""
schedule = []
total_steps = self.num_stages + self.num_micro_batches - 1
for step in range(total_steps):
step_schedule = []
for stage in range(self.num_stages):
micro_batch = step - stage
if 0 <= micro_batch < self.num_micro_batches:
step_schedule.append((stage, micro_batch))
else:
step_schedule.append(None) # 空闲
schedule.append(step_schedule)
return schedule
def calculate_bubble_ratio(self) -> float:
"""计算气泡比例"""
total_steps = self.num_stages + self.num_micro_batches - 1
useful_steps = self.num_micro_batches * self.num_stages
total_slots = total_steps * self.num_stages
return 1 - useful_steps / total_slots
# 流水线并行效率分析
def analyze_pipeline_efficiency():
"""分析不同配置的流水线效率"""
for num_stages in [2, 4, 8]:
for num_micro in [4, 8, 16, 32]:
scheduler = GPipeScheduler(num_stages, num_micro)
bubble = scheduler.calculate_bubble_ratio()
print(f"Stages={num_stages}, Micro-batches={num_micro:2d}: "
f"效率={((1-bubble)*100):.1f}%, 气泡={bubble*100:.1f}%")
analyze_pipeline_efficiency()
7.3 专家并行（Expert Parallelism - MoE）
class ExpertParallelMoE(nn.Module):
"""
专家并行的混合专家层
MoE模型（如Mixtral）使用专家并行：
- 不同专家分布在不同GPU上
- 通过All-to-All通信路由token到对应专家
"""
def __init__(self, hidden_size: int, ffn_size: int,
num_experts: int, top_k: int, tp_size: int, tp_rank: int):
super().__init__()
self.num_experts = num_experts
self.top_k = top_k
self.tp_size = tp_size
self.tp_rank = tp_rank
# 每个GPU持有 num_experts/tp_size 个专家
self.local_num_experts = num_experts // tp_size
self.experts_start = tp_rank * self.local_num_experts
# 路由器（所有GPU共享）
self.gate = nn.Linear(hidden_size, num_experts, bias=False)
# 本地专家
self.experts = nn.ModuleList([
nn.Sequential(
nn.Linear(hidden_size, ffn_size, bias=False),
nn.SiLU(),
nn.Linear(ffn_size, hidden_size, bias=False),
)
for _ in range(self.local_num_experts)
])
def forward(self, x: torch.Tensor) -> torch.Tensor:
B, L, D = x.shape
x_flat = x.view(-1, D)
# 计算路由权重
router_logits = self.gate(x_flat)
router_probs = torch.softmax(router_logits, dim=-1)
# 选择top-k专家
top_k_probs, top_k_indices = torch.topk(router_probs, self.top_k, dim=-1)
top_k_probs = top_k_probs / top_k_probs.sum(dim=-1, keepdim=True)
# All-to-All通信：将token路由到对应专家所在的GPU
# 这里简化处理，实际需要复杂的通信逻辑
# 计算输出
output = torch.zeros_like(x_flat)
for expert_idx in range(self.local_num_experts):
global_expert_idx = self.experts_start + expert_idx
# 找到路由到当前专家的token
for k in range(self.top_k):
mask = (top_k_indices[:, k] == global_expert_idx)
if mask.any():
expert_input = x_flat[mask]
expert_output = self.experts[expert_idx](expert_input)
output[mask] += top_k_probs[mask, k].unsqueeze(-1) * expert_output
# All-to-All通信：收集结果
# dist.all_reduce(output)
return output.view(B, L, D)
八、推理框架对比
8.1 主流推理框架概览
"""
主流大模型推理框架对比:
1. vLLM
- 开发者: UC Berkeley
- 核心技术: PagedAttention, Continuous Batching
- 优势: 高吞吐量, 易用性好, 社区活跃
- 适用场景: 通用推理服务
2. TGI (Text Generation Inference)
- 开发者: Hugging Face
- 核心技术: Continuous Batching, Flash Attention, 量化
- 优势: 与HF生态集成好, 支持模型多
- 适用场景: HuggingFace模型部署
3. SGLang
- 开发者: UC Berkeley
- 核心技术: RadixAttention, 前端DSL
- 优势: 复杂推理流程优化, 前缀共享
- 适用场景: 复杂prompt, 多轮对话
4. TensorRT-LLM
- 开发者: NVIDIA
- 核心技术: TensorRT优化, 量化, 自定义kernel
- 优势: NVIDIA GPU性能最优, 延迟最低
- 适用场景: NVIDIA GPU专用
"""
FRAMEWORK_COMPARISON = {
"vLLM": {
"throughput": "高",
"latency": "中",
"ease_of_use": "简单",
"model_support": "广泛",
"quantization": ["AWQ", "GPTQ", "FP8", "INT8"],
"distributed": ["TP", "PP"],
"best_for": "通用推理服务",
"install": "pip install vllm",
},
"TGI": {
"throughput": "高",
"latency": "中",
"ease_of_use": "中等",
"model_support": "广泛（HF模型）",
"quantization": ["GPTQ", "AWQ", "EETQ", "bitsandbytes"],
"distributed": ["TP"],
"best_for": "HF模型部署",
"install": "docker pull ghcr.io/huggingface/text-generation-inference",
},
"SGLang": {
"throughput": "很高",
"latency": "低",
"ease_of_use": "中等",
"model_support": "主流模型",
"quantization": ["AWQ", "GPTQ", "FP8"],
"distributed": ["TP", "DP"],
"best_for": "复杂推理流程",
"install": "pip install sglang",
},
"TensorRT-LLM": {
"throughput": "很高",
"latency": "最低",
"ease_of_use": "复杂",
"model_support": "主流模型",
"quantization": ["INT8", "INT4", "FP8", "SmoothQuant"],
"distributed": ["TP", "PP"],
"best_for": "NVIDIA GPU极致性能",
"install": "需要编译安装",
},
}
def recommend_framework(requirements: dict) -> str:
"""根据需求推荐推理框架"""
if requirements.get("gpu_vendor") != "nvidia":
# 非NVIDIA GPU，排除TensorRT-LLM
if requirements.get("priority") == "throughput":
return "vLLM"
return "vLLM"
if requirements.get("priority") == "latency":
return "TensorRT-LLM (最低延迟)"
elif requirements.get("priority") == "throughput":
if requirements.get("complex_prompts"):
return "SGLang (复杂prompt优化)"
return "vLLM"
elif requirements.get("priority") == "ease_of_use":
return "vLLM"
elif requirements.get("hf_ecosystem"):
return "TGI"
return "vLLM (通用推荐)"
8.2 vLLM部署示例
# vLLM部署示例代码
# 方法1: 使用Python API
from vllm import LLM, SamplingParams
def vllm_serve_example():
"""vLLM Python API示例"""
# 初始化模型
llm = LLM(
model="meta-llama/Llama-3-8B-Instruct",
tensor_parallel_size=1, # GPU数量
gpu_memory_utilization=0.9, # GPU显存利用率
max_model_len=4096, # 最大序列长度
quantization="awq", # 量化方式 (可选)
dtype="half", # 数据类型
)
# 设置采样参数
sampling_params = SamplingParams(
temperature=0.7,
top_p=0.9,
max_tokens=512,
repetition_penalty=1.1,
)
# 批量推理
prompts = [
"什么是大语言模型？",
"请解释量子计算的基本原理。",
"写一首关于春天的诗。",
]
outputs = llm.generate(prompts, sampling_params)
for output in outputs:
print(f"Prompt: {output.prompt[:50]}...")
print(f"Generated: {output.outputs[0].text}")
print(f"Tokens: {len(output.outputs[0].token_ids)}")
print("---")
# 方法2: 启动API服务器
"""
# 命令行启动:
python -m vllm.entrypoints.openai.api_server \
--model meta-llama/Llama-3-8B-Instruct \
--tensor-parallel-size 1 \
--gpu-memory-utilization 0.9 \
--max-model-len 4096 \
--port 8000
# 然后可以使用OpenAI兼容的API调用:
# curl http://localhost:8000/v1/chat/completions \
# -H "Content-Type: application/json" \
# -d '{
# "model": "meta-llama/Llama-3-8B-Instruct",
# "messages": [{"role": "user", "content": "Hello!"}],
# "max_tokens": 100
# }'
"""
# 方法3: 使用Docker部署
"""
docker run --gpus all \
-v ~/.cache/huggingface:/root/.cache/huggingface \
-p 8000:8000 \
--ipc=host \
vllm/vllm-openai:latest \
--model meta-llama/Llama-3-8B-Instruct \
--tensor-parallel-size 1
"""
8.3 SGLang部署示例
# SGLang部署示例
# 方法1: Python API
import sglang as sgl
def sglang_example():
"""SGLang示例"""
# 定义推理函数
@sgl.function
def multi_turn_chat(s, question1, question2):
s += sgl.system("你是一个有帮助的AI助手。")
s += sgl.user(question1)
s += sgl.assistant(sgl.gen("answer1", max_tokens=256))
s += sgl.user(question2)
s += sgl.assistant(sgl.gen("answer2", max_tokens=256))
# 运行
state = multi_turn_chat.run(
question1="什么是机器学习？",
question2="它和深度学习有什么区别？",
)
print(state["answer1"])
print(state["answer2"])
# 方法2: 启动API服务器
"""
python -m sglang.launch_server \
--model-path meta-llama/Llama-3-8B-Instruct \
--tp 1 \
--port 30000
"""
# SGLang的RadixAttention优势:
# - 自动缓存公共前缀的KV Cache
# - 对多轮对话特别有效
# - 支持复杂的推理流程编排
九、硬件选型
9.1 GPU选型指南
"""
主流GPU推理性能对比
"""
GPU_SPECS = {
"NVIDIA A100 80GB": {
"memory_gb": 80,
"bandwidth_gbps": 2048,
"fp16_tflops": 312,
"int8_tflops": 624,
"price_usd": 15000,
"interconnect": "NVLink 3.0 (600 GB/s)",
"best_for": "训练+推理，通用选择",
},
"NVIDIA H100 80GB": {
"memory_gb": 80,
"bandwidth_gbps": 3350,
"fp16_tflops": 990, # 使用FP8
"int8_tflops": 1979,
"price_usd": 30000,
"interconnect": "NVLink 4.0 (900 GB/s)",
"best_for": "高性能推理，FP8支持",
},
"NVIDIA L40S": {
"memory_gb": 48,
"bandwidth_gbps": 864,
"fp16_tflops": 362,
"int8_tflops": 733,
"price_usd": 7000,
"interconnect": "PCIe Gen4",
"best_for": "推理优化，性价比高",
},
"NVIDIA A10G": {
"memory_gb": 24,
"bandwidth_gbps": 600,
"fp16_tflops": 125,
"int8_tflops": 250,
"price_usd": 2000,
"interconnect": "PCIe Gen4",
"best_for": "小模型推理，预算有限",
},
}
def estimate_model_gpu_requirements(model_params_b: float, quantization: str = "fp16"):
"""估算模型需要的GPU配置"""
bytes_per_param = {
"fp32": 4, "fp16": 2, "int8": 1, "int4": 0.5, "awq4": 0.5, "gptq4": 0.5
}
model_memory = model_params_b * bytes_per_param.get(quantization, 2)
kv_cache_memory = model_params_b * 0.1 # 粗略估计
total_memory = model_memory + kv_cache_memory
print(f"\n模型: {model_params_b}B, 量化: {quantization}")
print(f"模型权重: {model_memory:.1f} GB")
print(f"KV Cache估计: {kv_cache_memory:.1f} GB")
print(f"总显存需求: {total_memory:.1f} GB")
for gpu_name, specs in GPU_SPECS.items():
num_gpus = -(-total_memory // specs["memory_gb"]) # 向上取整
if num_gpus <= 8: # 单机最多8卡
cost = num_gpus * specs["price_usd"]
print(f" {gpu_name}: {num_gpus}卡, 成本 ${cost:,}")
# 分析不同模型的GPU需求
for model_size in [7, 13, 34, 70, 110]:
estimate_model_gpu_requirements(model_size, "fp16")
estimate_model_gpu_requirements(model_size, "awq4")
9.2 内存带宽分析
def memory_bandwidth_analysis():
"""内存带宽瓶颈分析"""
# Decode阶段是memory-bound
# 每生成一个token需要读取所有模型权重
models = {
"7B-FP16": 14, # GB
"7B-INT4": 3.5,
"70B-FP16": 140,
"70B-INT4": 35,
"70B-INT8": 70,
}
gpus = {
"A100": 2048, # GB/s
"H100": 3350,
"L40S": 864,
"4090": 1008,
}
print("=== 理论最大吞吐量 (tokens/s) ===\n")
print(f"{'模型':<15}", end="")
for gpu_name in gpus:
print(f"{gpu_name:>10}", end="")
print()
print("-" * 55)
for model_name, model_size in models.items():
print(f"{model_name:<15}", end="")
for gpu_name, bandwidth in gpus.items():
max_tps = bandwidth / model_size
print(f"{max_tps:>10.0f}", end="")
print()
memory_bandwidth_analysis()
十、性能基准测试方法
10.1 基准测试框架
import time
import statistics
import asyncio
from dataclasses import dataclass
from typing import List, Callable
import json
@dataclass
class BenchmarkResult:
test_name: str
total_requests: int
successful_requests: int
failed_requests: int
total_tokens: int
total_time_s: float
ttft_avg_ms: float # 首token延迟（平均）
ttft_p50_ms: float
ttft_p99_ms: float
tpot_avg_ms: float # 每token延迟（平均）
tpot_p50_ms: float
tpot_p99_ms: float
throughput_tps: float # 吞吐量（tokens/s）
requests_per_sec: float
avg_input_tokens: float
avg_output_tokens: float
class LLMBenchmark:
"""LLM推理基准测试框架"""
def __init__(self, api_base: str, model: str):
self.api_base = api_base
self.model = model
self.results = []
async def single_request(self, prompt: str, max_tokens: int = 100) -> dict:
"""发送单个请求并测量延迟"""
import aiohttp
ttft = None
tokens = []
start_time = time.time()
async with aiohttp.ClientSession() as session:
async with session.post(
f"{self.api_base}/v1/chat/completions",
json={
"model": self.model,
"messages": [{"role": "user", "content": prompt}],
"max_tokens": max_tokens,
"stream": True,
},
) as response:
async for line in response.content:
if line.startswith(b"data: "):
data = line[6:].decode()
if data.strip() == "[DONE]":
break
chunk = json.loads(data)
if "choices" in chunk and chunk["choices"]:
delta = chunk["choices"][0].get("delta", {})
if "content" in delta:
if ttft is None:
ttft = (time.time() - start_time) * 1000
tokens.append(delta["content"])
end_time = time.time()
return {
"ttft_ms": ttft or 0,
"tpot_ms": (end_time - start_time - (ttft or 0) / 1000) / max(len(tokens), 1) * 1000,
"total_time_ms": (end_time - start_time) * 1000,
"num_tokens": len(tokens),
"success": True,
}
async def run_concurrent_benchmark(
self,
prompts: List[str],
concurrency: int = 10,
max_tokens: int = 100,
) -> BenchmarkResult:
"""运行并发基准测试"""
semaphore = asyncio.Semaphore(concurrency)
results = []
async def bounded_request(prompt):
async with semaphore:
return await self.single_request(prompt, max_tokens)
start_time = time.time()
tasks = [bounded_request(p) for p in prompts]
results = await asyncio.gather(*tasks, return_exceptions=True)
total_time = time.time() - start_time
# 过滤成功的请求
successful = [r for r in results if isinstance(r, dict) and r.get("success")]
failed = [r for r in results if not isinstance(r, dict) or not r.get("success")]
# 计算统计指标
ttft_values = [r["ttft_ms"] for r in successful]
tpot_values = [r["tpot_ms"] for r in successful if r["tpot_ms"] > 0]
total_tokens = sum(r["num_tokens"] for r in successful)
return BenchmarkResult(
test_name=f"concurrency_{concurrency}",
total_requests=len(prompts),
successful_requests=len(successful),
failed_requests=len(failed),
total_tokens=total_tokens,
total_time_s=total_time,
ttft_avg_ms=statistics.mean(ttft_values) if ttft_values else 0,
ttft_p50_ms=statistics.median(ttft_values) if ttft_values else 0,
ttft_p99_ms=sorted(ttft_values)[int(len(ttft_values)*0.99)] if ttft_values else 0,
tpot_avg_ms=statistics.mean(tpot_values) if tpot_values else 0,
tpot_p50_ms=statistics.median(tpot_values) if tpot_values else 0,
tpot_p99_ms=sorted(tpot_values)[int(len(tpot_values)*0.99)] if tpot_values else 0,
throughput_tps=total_tokens / total_time,
requests_per_sec=len(successful) / total_time,
avg_input_tokens=statistics.mean([len(p) for p in prompts]),
avg_output_tokens=statistics.mean([r["num_tokens"] for r in successful]) if successful else 0,
)
def generate_report(self, results: List[BenchmarkResult]) -> str:
"""生成测试报告"""
report = "# LLM推理性能基准测试报告\n\n"
for result in results:
report += f"## {result.test_name}\n\n"
report += f"- 总请求数: {result.total_requests}\n"
report += f"- 成功: {result.successful_requests}, 失败: {result.failed_requests}\n"
report += f"- 总token数: {result.total_tokens}\n"
report += f"- 总耗时: {result.total_time_s:.2f}s\n\n"
report += "### 延迟指标\n"
report += f"- TTFT (平均/P50/P99): {result.ttft_avg_ms:.0f}ms / {result.ttft_p50_ms:.0f}ms / {result.ttft_p99_ms:.0f}ms\n"
report += f"- TPOT (平均/P50/P99): {result.tpot_avg_ms:.0f}ms / {result.tpot_p50_ms:.0f}ms / {result.tpot_p99_ms:.0f}ms\n\n"
report += "### 吞吐量指标\n"
report += f"- Token吞吐: {result.throughput_tps:.1f} tokens/s\n"
report += f"- 请求吞吐: {result.requests_per_sec:.2f} requests/s\n\n"
return report
# 使用示例
async def run_benchmark():
benchmark = LLMBenchmark("http://localhost:8000", "llama-3-8b")
# 准备测试数据
prompts = [
"请详细解释什么是Transformer架构。",
"写一个Python快速排序算法。",
"解释量子纠缠现象。",
] * 10 # 30个请求
# 测试不同并发度
results = []
for concurrency in [1, 5, 10, 20]:
result = await benchmark.run_concurrent_benchmark(
prompts[:concurrency*3],
concurrency=concurrency,
max_tokens=200,
)
results.append(result)
print(benchmark.generate_report(results))
10.2 使用标准benchmark工具
# 使用vLLM自带的benchmark
# 吞吐量测试
python -m vllm.entrypoints.openai.api_server_benchmark \
--backend openai \
--base-url http://localhost:8000 \
--model meta-llama/Llama-3-8B-Instruct \
--dataset ShareGPT \
--num-prompts 1000 \
--request-rate 10
# 延迟测试
python -m vllm.entrypoints.openai.api_server_benchmark \
--backend openai \
--base-url http://localhost:8000 \
--model meta-llama/Llama-3-8B-Instruct \
--dataset random \
--input-len 128 \
--output-len 128 \
--num-prompts 100
十一、成本优化策略
11.1 成本分析模型
from dataclasses import dataclass
@dataclass
class CostAnalysis:
"""推理服务成本分析"""
# 硬件成本
gpu_cost_per_hour: float # GPU租赁成本 ($/hour)
num_gpus: int
# 性能指标
throughput_tokens_per_sec: float
gpu_utilization: float # 0-1
# 计算成本
@property
def hourly_cost(self) -> float:
return self.gpu_cost_per_hour * self.num_gpus
@property
def tokens_per_dollar(self) -> float:
"""每美元可处理的token数"""
if self.hourly_cost == 0:
return float('inf')
tokens_per_hour = self.throughput_tokens_per_sec * 3600
return tokens_per_hour / self.hourly_cost
@property
def cost_per_million_tokens(self) -> float:
"""每百万token的成本"""
if self.tokens_per_dollar == 0:
return float('inf')
return 1_000_000 / self.tokens_per_dollar
def cost_optimization_analysis():
"""成本优化分析"""
scenarios = [
# 场景1: A100 + FP16
{"name": "A100-FP16", "gpu_cost": 3.0, "gpus": 1, "tps": 30, "util": 0.7},
# 场景2: A100 + INT4
{"name": "A100-INT4", "gpu_cost": 3.0, "gpus": 1, "tps": 50, "util": 0.8},
# 场景3: L40S + INT4
{"name": "L40S-INT4", "gpu_cost": 1.5, "gpus": 1, "tps": 35, "util": 0.75},
# 场景4: H100 + FP8
{"name": "H100-FP8", "gpu_cost": 5.0, "gpus": 1, "tps": 80, "util": 0.85},
]
print("=== 推理成本对比 ===\n")
print(f"{'场景':<15} {'每小时成本':>12} {'吞吐量':>12} {'$/M tokens':>12}")
print("-" * 55)
for s in scenarios:
analysis = CostAnalysis(
gpu_cost_per_hour=s["gpu_cost"],
num_gpus=s["gpus"],
throughput_tokens_per_sec=s["tps"],
gpu_utilization=s["util"],
)
print(f"{s['name']:<15} ${analysis.hourly_cost:>10.2f} "
f"{s['tps']:>10} t/s ${analysis.cost_per_million_tokens:>10.4f}")
cost_optimization_analysis()
11.2 动态扩缩容策略
import time
from collections import deque
class AutoScaler:
"""自动扩缩容管理器"""
def __init__(
self,
min_replicas: int = 1,
max_replicas: int = 10,
target_utilization: float = 0.7,
scale_up_threshold: float = 0.8,
scale_down_threshold: float = 0.3,
cooldown_seconds: int = 300,
):
self.min_replicas = min_replicas
self.max_replicas = max_replicas
self.target_utilization = target_utilization
self.scale_up_threshold = scale_up_threshold
self.scale_down_threshold = scale_down_threshold
self.cooldown_seconds = cooldown_seconds
self.current_replicas = min_replicas
self.last_scale_time = 0
self.utilization_history = deque(maxlen=60) # 保留60个数据点
def record_utilization(self, utilization: float):
"""记录当前利用率"""
self.utilization_history.append({
"time": time.time(),
"utilization": utilization,
})
def should_scale(self) -> tuple:
"""
判断是否需要扩缩容
返回: (should_scale, direction, target_replicas)
"""
if len(self.utilization_history) < 10:
return False, None, self.current_replicas
# 冷却期检查
if time.time() - self.last_scale_time < self.cooldown_seconds:
return False, None, self.current_replicas
# 计算平均利用率
recent = list(self.utilization_history)[-10:]
avg_util = sum(r["utilization"] for r in recent) / len(recent)
if avg_util > self.scale_up_threshold and self.current_replicas < self.max_replicas:
# 扩容
target = min(
self.max_replicas,
int(self.current_replicas * (avg_util / self.target_utilization)) + 1
)
return True, "up", target
elif avg_util < self.scale_down_threshold and self.current_replicas > self.min_replicas:
# 缩容
target = max(
self.min_replicas,
int(self.current_replicas * (avg_util / self.target_utilization))
)
return True, "down", target
return False, None, self.current_replicas
def scale(self, target_replicas: int):
"""执行扩缩容"""
self.current_replicas = target_replicas
self.last_scale_time = time.time()
print(f"扩缩容: {self.current_replicas} -> {target_replicas} replicas")
十二、实战案例：高性能推理服务搭建
12.1 完整的vLLM部署方案
"""
完整的vLLM推理服务部署方案
包含:
1. 模型加载与优化配置
2. API服务搭建
3. 监控与日志
4. 负载均衡
5. 成本优化
"""
# docker-compose.yml
DOCKER_COMPOSE = """
version: '3.8'
services:
vllm-server:
image: vllm/vllm-openai:latest
runtime: nvidia
environment:
- NVIDIA_VISIBLE_DEVICES=all
volumes:
- ~/.cache/huggingface:/root/.cache/huggingface
ports:
- "8000:8000"
ipc: host
command: >
--model meta-llama/Llama-3-8B-Instruct
--tensor-parallel-size 1
--gpu-memory-utilization 0.9
--max-model-len 4096
--quantization awq
--dtype half
--enable-prefix-caching
--max-num-seqs 32
--port 8000
deploy:
resources:
reservations:
devices:
- driver: nvidia
count: 1
capabilities: [gpu]
nginx:
image: nginx:latest
ports:
- "80:80"
volumes:
- ./nginx.conf:/etc/nginx/nginx.conf
depends_on:
- vllm-server
prometheus:
image: prom/prometheus:latest
ports:
- "9090:9090"
volumes:
- ./prometheus.yml:/etc/prometheus/prometheus.yml
grafana:
image: grafana/grafana:latest
ports:
- "3000:3000"
environment:
- GF_SECURITY_ADMIN_PASSWORD=admin
"""
# nginx.conf
NGINX_CONF = """
events {
worker_connections 1024;
}
http {
upstream vllm {
server vllm-server:8000;
}
server {
listen 80;
location /v1/ {
proxy_pass http://vllm;
proxy_set_header Host $host;
proxy_set_header X-Real-IP $remote_addr;
proxy_read_timeout 300s;
proxy_send_timeout 300s;
}
location /health {
proxy_pass http://vllm/health;
}
}
}
"""
# 监控脚本
import psutil
import time
class InferenceMonitor:
"""推理服务监控"""
def __init__(self, api_base: str):
self.api_base = api_base
self.metrics_history = []
def collect_metrics(self) -> dict:
"""收集系统指标"""
import requests
metrics = {
"timestamp": time.time(),
"cpu_percent": psutil.cpu_percent(),
"memory_percent": psutil.virtual_memory().percent,
}
# GPU指标（需要nvidia-smi或pynvml）
try:
import pynvml
pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
gpu_info = pynvml.nvmlDeviceGetUtilizationRates(handle)
gpu_memory = pynvml.nvmlDeviceGetMemoryInfo(handle)
metrics["gpu_utilization"] = gpu_info.gpu
metrics["gpu_memory_used_gb"] = gpu_memory.used / 1e9
metrics["gpu_memory_total_gb"] = gpu_memory.total / 1e9
except:
pass
# API健康检查
try:
response = requests.get(f"{self.api_base}/health", timeout=5)
metrics["api_status"] = "healthy" if response.status_code == 200 else "unhealthy"
except:
metrics["api_status"] = "unreachable"
self.metrics_history.append(metrics)
return metrics
def check_alerts(self, metrics: dict) -> list:
"""检查告警条件"""
alerts = []
if metrics.get("gpu_utilization", 0) > 95:
alerts.append("GPU利用率过高")
if metrics.get("gpu_memory_used_gb", 0) / metrics.get("gpu_memory_total_gb", 1) > 0.95:
alerts.append("GPU显存接近上限")
if metrics.get("api_status") != "healthy":
alerts.append("API服务异常")
return alerts
12.2 性能调优检查清单
PERFORMANCE_CHECKLIST = """
# 大模型推理性能调优检查清单
## 模型层面
- [ ] 选择合适的量化方式（AWQ/GPTQ/INT8）
- [ ] 使用GQA/MQA减少KV Cache
- [ ] 启用Flash Attention
- [ ] 选择合适的精度（FP16/BF16/FP8）
## 系统层面
- [ ] 启用PagedAttention（vLLM默认开启）
- [ ] 配置合适的max_num_seqs
- [ ] 启用Prefix Caching（多轮对话场景）
- [ ] 使用Continuous Batching
- [ ] 调整gpu_memory_utilization
## 部署层面
- [ ] 使用Tensor Parallelism充分利用多GPU
- [ ] 配置负载均衡
- [ ] 设置合理的超时时间
- [ ] 启用请求队列
## 监控层面
- [ ] 监控GPU利用率和显存
- [ ] 监控TTFT和TPOT
- [ ] 监控请求队列长度
- [ ] 设置告警阈值
## 成本优化
- [ ] 选择性价比最优的GPU
- [ ] 使用量化降低硬件需求
- [ ] 实施动态扩缩容
- [ ] 监控和优化GPU利用率
"""
print(PERFORMANCE_CHECKLIST)
十三、常见问题解答
Q1: 如何选择量化方案？
A: 量化方案选择建议：
- AWQ: 推理速度快，质量损失小，推荐首选
- GPTQ: 质量稍好，速度略慢
- INT8: 几乎无损，但内存节省有限
- GGUF: 适合CPU推理和边缘设备
Q2: 投机解码的加速比一般是多少？
A: 投机解码的加速比取决于：
- Draft模型与Target模型的一致性：一致性越高，接受率越高
- Draft模型的速度：越快越好
- 典型加速比：1.5x - 3x
Q3: vLLM和TGI如何选择？
A:
- vLLM: 通用场景推荐，吞吐量高，API兼容性好
- TGI: 使用HuggingFace模型时推荐，与HF生态集成好
Q4: 如何估算需要多少GPU？
A:
- 计算模型权重大小：参数量 × 每参数字节数
- 加上KV Cache开销：约占模型大小的10-20%
- 除以单GPU显存容量
- 留出10-15%余量
Q5: 多轮对话如何优化？
A:
- 启用Prefix Caching缓存system prompt
- 使用KV Cache避免重复计算
- 合理设置max_model_len
十四、总结
大模型推理优化是一个系统工程，需要从模型、算法、系统、硬件等多个层面综合考虑。关键要点：
- 量化是基础：选择合适的量化方案可以大幅降低硬件需求和成本
- KV Cache是关键：PagedAttention和Prefix Caching是提升效率的核心技术
- 批处理是核心：Continuous Batching充分利用GPU计算资源
- 投机解码是加速器：在不损失质量的前提下提升速度
- 并行是扩展手段：Tensor/Pipeline/Expert Parallelism支持更大模型
- 监控是保障：持续监控性能指标，及时发现和解决问题
选择合适的优化策略组合，可以将推理成本降低数倍甚至数十倍，同时保持良好的用户体验。
本教程内容持续更新中，欢迎反馈和建议。
