2024 年我们做一个 AI 视频生成业务 后端用 Stable Diffusion + AnimateDiff 跑文本生成短视频。需求是用户输入一段描述 30 秒内出 5 秒短视频。原型阶段单卡 A100 跑得也挺顺 1 个请求 25 秒出图。但真正上线后我们陆续踩了一堆坑。第一种最让我傻眼 上线第一天 100 个并发请求 GPU 显存爆掉 OOM 整个服务 crash 用户全部 500。第二种最难缠 同一段 prompt 不同时刻跑出的视频质量差距巨大 有时清晰流畅有时模糊跳帧 排查发现是不同模型版本 不同 sampler 不同 steps 没固定 团队成员各自试各的参数。第三种最离谱 月度账单一个月烧 12 万美元 单卡 A100 包月 8000 美元 我们开了 15 张卡平均利用率只有 30% 高峰时还不够 谷峰时全闲着浪费严重。第四种最致命 一个用户上传超长 prompt 1000 字符 我们没做长度限制 模型 forward 时显存翻倍 把整张卡占满 其他用户全部排队 5 分钟才轮上 SLA 完全没法看。第五种最莫名其妙 我们用 K8s 调度 GPU pod 同一个 deployment 同样的 image 在两台 A100 上跑出的速度差 30% 排查发现是云厂商分配的 A100 有不同代次 我们没指定具体型号 调度器随机分。我盯着这一连串问题想了很久才彻底想明白第一版错在一个根本的认知上我以为 GPU 推理就是 把模型 load 进显存 给个请求 forward 一下出结果 可这个认知是错的真正能扛业务的 AI 推理服务是一个 模型加载与显存管理 加 batch 与队列调度 加 模型版本与参数固化 加 GPU 异构与亲和性 加 动态扩缩容 加 限流与降级 加 成本监控 的整套工程方法论 任何一环没做都可能让你的服务挂掉或者亏钱本文从头梳理 AI 推理服务的工程化要点 显存怎么管 batch 怎么调度 模型怎么管理 GPU 怎么调度 扩缩容怎么做 限流怎么设 以及一些把 GPU 推理做扎实要避开的工程坑
问题背景:为什么 GPU 推理不是 model.forward 就完事
很多人对 AI 推理服务的认知是 load 模型 收请求 forward 返回结果 但生产里你会发现 显存爆 batch 不优 模型版本乱 GPU 利用率低 成本失控。问题的根源在于:
- 显存是硬上限:A100 80G 看似很大 但模型权重 KV cache 中间激活 加起来很容易爆 必须算清楚每个请求的显存预算。
- batch 决定吞吐:单请求单 forward GPU 利用率只有 10-20% 必须做 dynamic batching 把多个请求拼一起 forward 吞吐翻 5-10 倍。
- 模型版本要锁死:同一个 prompt 不同模型 weights sampler steps 出来结果可能差天差地 必须固化整套推理参数。
- GPU 异构必须感知:同样标 A100 实际可能是 PCIe 也可能是 SXM 性能差 30% 必须用 nodeSelector 强制分配。
- 队列调度比扩容更重要:GPU 资源稀缺贵 不能随便扩容 必须做请求队列 排队等待 才能在固定卡数内最大化吞吐。
- 输入大小必须有上限:用户超长 prompt 超大图像会让显存翻倍 必须前置校验拒绝超大请求。
一 显存管理:模型加载与卸载
显存是 GPU 推理的硬约束。一张 A100 80G 跑 SD 1.5 大概占 6G 留给 KV cache 中间激活 单请求峰值约 12G batch 4 就 30G 留点 buffer 一张卡最多并发 4-6 个请求。把这账算清楚 才能定 batch size 与并发上限。
import torch
from diffusers import StableDiffusionPipeline
class ModelManager:
"""统一管理 GPU 上的模型加载与显存"""
def __init__(self, device: str = 'cuda:0'):
self.device = device
self.models = {}
# 关键 启用 memory efficient attention 显存降 30%
torch.backends.cuda.matmul.allow_tf32 = True
def load(self, name: str, model_id: str, dtype=torch.float16):
"""按需加载 同一进程内复用"""
if name in self.models:
return self.models[name]
pipe = StableDiffusionPipeline.from_pretrained(
model_id, torch_dtype=dtype,
safety_checker=None, # 关闭 safety checker 省 1G 显存
)
# 启用 xformers attention 显存与速度双提升
pipe.enable_xformers_memory_efficient_attention()
# 模型权重移到 GPU
pipe = pipe.to(self.device)
# vae 用 slicing 减少峰值显存
pipe.enable_vae_slicing()
self.models[name] = pipe
return pipe
def unload(self, name: str):
"""从显存彻底卸载"""
if name in self.models:
del self.models[name]
torch.cuda.empty_cache()
torch.cuda.synchronize()
def gpu_memory_mb(self) -> dict:
return {
'allocated': torch.cuda.memory_allocated(self.device) / 1024 / 1024,
'reserved': torch.cuda.memory_reserved(self.device) / 1024 / 1024,
'free': (torch.cuda.get_device_properties(self.device).total_memory
- torch.cuda.memory_allocated(self.device)) / 1024 / 1024,
}
显存优化的工程经验 xformers 必开 显存降 30% 速度升 20% safety_checker 关掉省 1G vae slicing 让峰值显存降一半 fp16 比 fp32 省一半 这些加起来一张 A100 跑 SD 的能力从 batch 2 提到 batch 8。生产服务必须定期上报 gpu_memory_mb 监控 显存利用率超过 85% 报警 避免突发 OOM 让服务挂掉。
二 Dynamic Batching:吞吐倍增的核心
单请求单 forward GPU 利用率极低 因为 GPU 是 SIMD 单指令多数据 一次处理 1 个样本和处理 8 个样本耗时差不多。Dynamic batching 把短时间内到的请求拼成一个 batch 一起 forward 吞吐能提升 5-10 倍 这是 AI 推理服务的必做优化。
import asyncio
import time
from collections import deque
from dataclasses import dataclass
@dataclass
class InferenceRequest:
prompt: str
params: dict
future: asyncio.Future
enqueued_at: float
class DynamicBatcher:
"""收集请求 凑齐 batch 或超时后统一 forward"""
def __init__(self, model, max_batch: int = 8, max_wait_ms: int = 50):
self.model = model
self.max_batch = max_batch
self.max_wait = max_wait_ms / 1000
self.queue: deque[InferenceRequest] = deque()
self.lock = asyncio.Lock()
self.worker_task = None
async def start(self):
self.worker_task = asyncio.create_task(self._loop())
async def submit(self, prompt: str, params: dict) -> bytes:
fut = asyncio.get_event_loop().create_future()
req = InferenceRequest(prompt, params, fut, time.time())
async with self.lock:
self.queue.append(req)
return await fut
submit 把请求塞进队列就返回 future 真正的 batch forward 在独立的后台 loop 里跑 这样 API 层不会被推理阻塞 高并发也能轻松接受请求。下面是后台 loop 的核心逻辑 每 5ms 检查一次队列 凑够 max_batch 或者最早请求等待超过 max_wait 就触发 forward 在 lock 外异步执行 不阻塞新请求入队。
async def _loop(self):
while True:
await asyncio.sleep(0.005) # 5ms tick
async with self.lock:
if not self.queue:
continue
# 凑够 batch 或者最早请求等了超过 max_wait
first_age = time.time() - self.queue[0].enqueued_at
if len(self.queue) < self.max_batch and first_age < self.max_wait:
continue
batch = []
while self.queue and len(batch) < self.max_batch:
batch.append(self.queue.popleft())
# batch forward 在 lock 外执行 不阻塞新请求入队
prompts = [r.prompt for r in batch]
try:
# 假设所有请求参数兼容 实际可能要按参数分桶
images = await asyncio.to_thread(
self.model, prompts,
num_inference_steps=batch[0].params.get('steps', 30),
)
for req, img in zip(batch, images.images):
req.future.set_result(img)
except Exception as e:
for req in batch:
req.future.set_exception(e)
dynamic batching 的关键参数 max_wait_ms 决定单请求的额外延迟 50ms 是个甜区 用户感知不到但能凑够 batch max_batch 决定显存占用 按显存上限算回去 不要拍脑袋。我们公司 SD 服务上 dynamic batching 后 单卡 QPS 从 0.4 提到 3.2 翻 8 倍 总卡数从 15 降到 4 月成本省 8 万美元。
三 模型版本与参数固化:可重现性
同一段 prompt 不同模型 weights sampler steps cfg_scale seed 跑出来结果可能差天差地。生产服务必须把整套推理参数版本化 用户每次请求都能拿到稳定结果 这对 SLA 与调试都至关重要。
from pydantic import BaseModel, Field
from typing import Literal
import hashlib
class InferenceConfig(BaseModel):
"""所有推理参数固化在这里 任何改动版本号 +1"""
config_version: str = '2024.05.01'
model_id: str = 'runwayml/stable-diffusion-v1-5'
model_revision: str = 'aa9ba505e1973ae5cd05f5aedd345178f52f8e6a' # commit hash
vae_id: str = 'stabilityai/sd-vae-ft-mse'
sampler: Literal['DPMSolverMultistep', 'Euler', 'DDIM'] = 'DPMSolverMultistep'
num_inference_steps: int = 30
guidance_scale: float = 7.5
width: int = 768
height: int = 768
negative_prompt: str = 'low quality, blurry, distorted, watermark'
@property
def hash(self) -> str:
# 整套配置的 hash 用于 cache key 与监控
return hashlib.sha256(self.json().encode()).hexdigest()[:16]
class InferenceRequest(BaseModel):
prompt: str = Field(..., max_length=300)
seed: int | None = None
config_version: str = '2024.05.01' # 客户端可以指定版本 服务端校验
def validate_prompt(self):
# 长度校验 token 数校验 内容审核
if len(self.prompt) > 300:
raise ValueError('prompt too long')
# 实战还要过 PII 过滤 NSFW 检测
return self
有了 InferenceConfig 服务端在 model load 时就锁定整套参数 客户端发请求时如果 config_version 对不上 直接拒绝或者走兼容路径。这样任何参数变化都会触发版本升级 不会有任何隐式漂移。
# 服务端加载与版本校验
class InferenceService:
def __init__(self, config: InferenceConfig):
self.config = config
self.pipe = StableDiffusionPipeline.from_pretrained(
config.model_id, revision=config.model_revision,
torch_dtype=torch.float16,
).to('cuda')
# 替换 sampler
from diffusers import DPMSolverMultistepScheduler
if config.sampler == 'DPMSolverMultistep':
self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(
self.pipe.scheduler.config,
)
def generate(self, req: InferenceRequest):
if req.config_version != self.config.config_version:
raise ValueError(f'config version mismatch '
f'server={self.config.config_version} '
f'client={req.config_version}')
generator = torch.Generator('cuda').manual_seed(req.seed) if req.seed else None
return self.pipe(
prompt=req.prompt,
negative_prompt=self.config.negative_prompt,
num_inference_steps=self.config.num_inference_steps,
guidance_scale=self.config.guidance_scale,
width=self.config.width,
height=self.config.height,
generator=generator,
)
模型版本固化的工程价值 一是可重现性 同一 prompt 同一 seed 永远出同一结果 调试与回滚有依据 二是兼容性 版本号让客户端与服务端协议明确 不会因为后端静默升级让客户端模型行为变化。我们的规范是 任何模型 sampler steps 变化必须 bump config_version 灰度切流 测完才全量切。
四 GPU 调度:Kubernetes 上的 GPU pod
K8s 调度 GPU pod 默认行为不够细致 同样标 A100 可能是 PCIe 也可能是 SXM 性能差 30% 必须用 nodeSelector 与 resource limit 精确指定 让调度器知道你的 pod 需要什么型号的 GPU。
apiVersion: apps/v1
kind: Deployment
metadata:
name: sd-inference
namespace: ai
spec:
replicas: 4
selector:
matchLabels:
app: sd-inference
template:
metadata:
labels:
app: sd-inference
spec:
# 关键 必须 nodeSelector 精确锁定 GPU 型号
nodeSelector:
nvidia.com/gpu.product: NVIDIA-A100-SXM4-80GB
# tolerations 让 pod 能调度到打了 GPU taint 的节点
tolerations:
- key: nvidia.com/gpu
operator: Exists
effect: NoSchedule
containers:
- name: inference
image: registry.mycompany.com/sd-inference:v2024.05.01
ports:
- containerPort: 8080
resources:
requests:
cpu: 4
memory: 16Gi
nvidia.com/gpu: 1 # 显式申请 1 张 GPU
limits:
cpu: 8
memory: 32Gi
nvidia.com/gpu: 1
env:
- name: NVIDIA_VISIBLE_DEVICES
value: all
- name: CUDA_VISIBLE_DEVICES
value: "0"
资源声明与 GPU 申请只是基础 真正决定 GPU pod 行不行的是探针与生命周期 模型 load 一次要 60-90 秒 readinessProbe 必须给足 initialDelay 否则 Service 会把流量过早转发到一个还没 load 完模型的 pod 上 立刻被打挂 优雅退出也要等队列排空再发 SIGTERM 否则正在 forward 的请求会被强制中断。
readinessProbe:
httpGet: { path: /healthz, port: 8080 }
initialDelaySeconds: 60 # 模型 load 慢 给足时间
periodSeconds: 10
livenessProbe:
httpGet: { path: /healthz, port: 8080 }
initialDelaySeconds: 90
periodSeconds: 30
failureThreshold: 3
lifecycle:
preStop:
exec:
command: ['sh', '-c', 'sleep 30'] # 优雅退出 等队列排空
GPU 调度的工程要点 nodeSelector 必须精确到 GPU product 名 不要只标 nvidia.com/gpu 否则你可能被分到 V100 P100 性能差很多 livenessProbe 阈值要宽 模型 load 60-90 秒 不能因为启动慢被反复 kill。我们公司金融业务的 SD 服务 严格指定 A100-SXM4-80GB tolerations + nodeSelector 双重保险 杜绝异构机型带来的 SLA 抖动。
[mermaid]flowchart TD
A[用户请求 prompt] --> B[API Gateway 限流]
B --> C[输入校验 长度 NSFW]
C --> D[配置版本校验]
D --> E[请求入队 DynamicBatcher]
E --> F{凑够 batch 或超时}
F -->|是| G[GPU forward]
F -->|否| H[等待 5ms tick]
H --> F
G --> I[结果分发 future.set]
I --> J[图像后处理]
J --> K[上传 S3 返回 URL]
G --> L[GPU 利用率 显存 监控]
L --> M[HPA 触发扩缩容]
五 扩缩容与限流:GPU 资源稀缺贵
GPU 卡贵 资源稀缺 不能像 CPU 服务那样随便扩容。必须做 智能扩缩容 加 入口限流 加 队列降级 三位一体 才能在固定卡数内最大化吞吐与 SLA。
apiVersion: autoscaling/v2
kind: HorizontalPodAutoscaler
metadata:
name: sd-inference-hpa
spec:
scaleTargetRef:
apiVersion: apps/v1
kind: Deployment
name: sd-inference
minReplicas: 2
maxReplicas: 10
metrics:
# GPU 利用率走 prometheus-adapter 自定义指标
- type: Pods
pods:
metric:
name: gpu_utilization_percent
target:
type: AverageValue
averageValue: "70"
# 队列深度 排队请求数
- type: Pods
pods:
metric:
name: inference_queue_depth
target:
type: AverageValue
averageValue: "10"
behavior:
scaleUp:
stabilizationWindowSeconds: 30 # GPU 加载慢 扩容要快
policies:
- type: Pods
value: 2
periodSeconds: 60
scaleDown:
stabilizationWindowSeconds: 600 # 缩容慢 避免抖动
policies:
- type: Pods
value: 1
periodSeconds: 120
GPU 扩缩容的特殊性 一是新 pod 启动慢 模型 load 60-90 秒 扩容必须比 CPU 服务更激进 二是缩容必须谨慎 GPU 卡按小时计费 缩容后再扩可能等不到 卡 三是必须有降级机制 队列深度超阈值时直接拒绝新请求或排队提示。我们公司 SD 服务峰谷比 5 倍 通过 HPA + 限流 + 排队 在 4-10 张 A100 区间动态伸缩 月成本省 50%。
六 GPU 推理的工程坑:那些文档里学不到的
讲完原理来说几个真实生产里踩过的坑。第一个坑是 CUDA OOM 不会立刻 crash 而是后续所有请求都失败 必须捕获 OOM 后主动 torch.cuda.empty_cache 并重启 worker 否则一个超大请求把卡污染所有后续请求都跪。第二个坑是 模型权重在多个 worker 重复加载 一张 80G 卡跑 4 个 worker 模型权重 6G * 4 = 24G 浪费严重 必须用 torch.multiprocessing 或者共享 cuda IPC 让多 worker 共享同一份权重。第三个坑是 PyTorch 的 fp16 inplace 操作偶尔出 NaN 模型输出全黑图 必须开 torch.set_grad_enabled(False) 加 autocast 否则数值不稳定。第四个坑是 GPU 推理在 K8s 上的就绪检测 模型 load 60 秒期间 Service 不该把流量发过来 必须 readinessProbe initialDelaySeconds 给足时间 不然新 pod 启动一上线就被打挂。第五个坑是 多模型混部一张 GPU 显存碎片化严重 不同模型大小不同 显存无法对齐 实际可用容量比理论值少 30% 生产推荐一张 GPU 跑一个模型多 worker 不要混部。
关键概念速查
| 概念 | 含义 | 工程价值 |
|---|---|---|
| fp16 | 半精度浮点 | 显存与速度双省 |
| xformers | 高效 attention | 显存降 30% |
| vae slicing | 分块解码 | 峰值显存降半 |
| Dynamic Batching | 请求拼批 | 吞吐提升 5-10 倍 |
| config_version | 推理参数版本 | 结果可重现 |
| nodeSelector GPU product | 精确 GPU 型号 | 避免异构性能抖动 |
| readinessProbe 60s | 就绪时长 | 模型 load 完才接流量 |
| GPU HPA | 按利用率扩缩 | 稀缺资源动态分配 |
| OOM recovery | 显存清理重启 | 污染恢复 |
| cuda IPC | 多 worker 共享权重 | 显存利用提升 |
避坑清单
- 模型加载必须开 xformers fp16 vae slicing 显存优化三连 不然 batch 上不去吞吐就上不去。
- 必做 dynamic batching max_wait 50ms 用户感知不到 但吞吐能翻 5-10 倍 这是性价比最高的优化。
- 整套推理参数必须固化版本 config_version 任何 sampler steps weight 改动都 bump 版本 客户端服务端校验。
- K8s 调度 GPU pod 必须 nodeSelector 锁 GPU product 名 不要只标 nvidia.com/gpu 异构 GPU 性能差 30%。
- readinessProbe initialDelaySeconds 给 60-90 秒 模型 load 慢 不要让启动期被反复 kill。
- HPA 扩容用 GPU 利用率与队列深度双指标 不要用 CPU GPU 服务 CPU 永远低 用 CPU 指标永远不扩。
- 限流必须在入口做 GPU 资源跟 CPU 不同 不能靠堆扩容 必须排队加拒绝。
- OOM 必须有恢复机制 捕获后 empty_cache 重启 worker 不然一个超大请求污染所有后续请求。
- 多 worker 共享模型权重 用 cuda IPC 不要每个 worker 各 load 一份 显存浪费严重。
- 一张 GPU 跑一个模型多 worker 不要混部不同模型 显存碎片化让实际可用容量少 30%。
总结
GPU 推理这事 很多人的直觉是 load 模型 给 prompt forward 出结果 就跟调一个本地函数一样简单 可这其实是把 我会跑 model.generate 和 我能在生产用 GPU 扛住高并发低延迟低成本 SLA 混为一谈。前者是会调 API 后者是懂 GPU 推理工程。中间隔着的是 显存管理 dynamic batching 模型版本 GPU 调度 扩缩容 限流 整整一套工程方法论。
从原型到生产 你需要做的事远不止 跑通推理。你要懂 显存怎么算 batch 怎么拼 模型怎么版本化 GPU 怎么精确调度 HPA 怎么按 GPU 指标走 OOM 怎么恢复 多 worker 怎么共享权重。每一项单独看都不复杂 但它们组合在一起 才是一个能在生产扛得住的 AI 推理服务。少任何一项 都可能让你的服务挂掉或者亏钱。
我经常用一个比喻来理解 GPU 推理服务 它有点像一个高级餐厅的厨房。GPU 是厨师 模型是菜谱 prompt 是客人下单 dynamic batching 是把同时下单的客人凑一桌一起做 显存是厨房工作台面积 大菜会占整张台子 GPU 调度是把不同档次的厨师配给不同档次的客人 HPA 是高峰期加请临时厨师 限流是入口领位 满桌了请客人排队不能让厨房乱套。你不能因为有了厨师就觉得餐厅能开得好 还要管厨房动线 拼桌策略 工作台调度 高峰加人 入口排队 这才是一整套餐厅运营。
这套架构最难的地方在于 它的复杂度在小流量 demo 时几乎完全暴露不了。你单卡跑一个 demo 一切都很顺 觉得 GPU 推理真好用。但真正放到生产 几百并发 长 prompt 异构 GPU 多模型混部 高峰扩缩容 你才发现 99% 的复杂度都在 那 1% 的极端 case 里 显存爆 OOM batch 不优 模型版本漂 GPU 性能抖动。建议任何想做 AI 推理服务上规模的团队 上线前一定要做 GPU 压测 故意打到 OOM 故意混部不同模型 故意拉满 batch 看系统能不能扛住 千万别等真实流量来教你 那时候账单可能已经烧到老板要砍项目的程度了。
—— 别看了 · 2026