Stable Diffusion 并行计算优化技巧与 GPU 性能提升指南

结果：batch=7 时 cuda time 降低 34%，显存占用合理。

多线程与异步流水线

1. WebUI 单线程阻塞问题

Automatic1111 的 WebUI 默认在 Gradio 主线程串行执行文本编码→U-Net→VAE，导致请求阻塞。

2. 异步流水线设计

# async_pipe.py
import asyncio, torch, threading, queue
from diffusers import StableDiffusionPipeline

class AsyncSD:
    def __init__(self, model_id="runwayml/stable-diffusion-v1-5"):
        self.pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")
        self.q_in = queue.Queue(maxsize=8)
        self.q_out = queue.Queue(maxsize=8)
        threading.Thread(target=self.worker, daemon=True).start()

    def worker(self):
        while True:
            item = self.q_in.get()
            if item is None:
                break
            prompt, seed, callback = item
            generator = torch.Generator("cuda").manual_seed(seed)
            image = self.pipe(prompt, num_inference_steps=20, generator=generator, callback=callback).images[0]
            self.q_out.put((prompt, image))

    async def generate(self, prompt, seed=42):
        loop = asyncio.get_event_loop()
        fut = loop.create_future()
        def cb(i, t, latents):
            loop.call_soon_threadsafe(fut.set_result, (i, latents))
        self.q_in.put((prompt, seed, cb))
        return await fut

效果：主线程立即返回，后台异步处理，GPU 利用率显著提升。

3. CLIP 多进程预编码

CLIP 文本编码为 CPU 密集型任务，支持多进程 + LRU 缓存。

# clip_cache.py
from multiprocessing import Pool
from functools import lru_cache
from transformers import CLIPTextModel, CLIPTokenizer
import torch

tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14").to("cuda")

@lru_cache(maxsize=1024)
def encode_text(prompt: str):
    tokens = tokenizer(prompt, max_length=77, padding="max_length", truncation=True, return_tensors="pt").input_ids.to("cuda")
    with torch.no_grad():
        embeds = text_encoder(tokens).last_hidden_state
    return embeds.cpu()

if __name__ == "__main__":
    pool = Pool(4)
    prompts = ["a cat"]*100
    list(pool.map(encode_text, prompts))

模型量化与内存复用

1. FP16 与 INT8 量化

官方权重为 FP32，使用 FP16 可减半显存并加速 Tensor Core。INT8 需校准。

# int8_calib.py
import torch, tqdm
from diffusers import StableDiffusionPipeline
from torch.quantization import MinMaxObserver, QConfig

pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to("cuda")

def calibrate(unet, n_samples=200):
    unet.eval()
    for _ in tqdm.trange(n_samples):
        latent = torch.randn(1, 4, 64, 64, dtype=torch.float16, device="cuda")
        t = torch.randint(0, 1000, (1,), device="cuda")
        with torch.no_grad():
            _ = unet(latent, t).sample
    print("calib done")

qconfig = QConfig(
    activation=MinMaxObserver.with_args(qscheme=torch.per_tensor_symmetric),
    weight=MinMaxObserver.with_args(qscheme=torch.per_tensor_symmetric)
)
pipe.unet.qconfig = qconfig
torch.quantization.prepare(pipe.unet, inplace=True)
calibrate(pipe.unet)
torch.quantization.convert(pipe.unet, inplace=True)

2. 内存池复用

针对 U-Net 中间 tensor 频繁申请的问题，使用 CUDA Graph 内存池减少 malloc 开销。

# pool_unet.py
from torch.cuda import graph_pool
import torch, functools

def get_pool(device):
    return graph_pool._graph_pool_handle(device)

def pooled_forward(self, *args, **kwargs):
    with torch.cuda.graph_pool.enable_pool(get_pool(args[0].device)):
        return self._forward(*args, **kwargs)

pipe.unet.forward = pooled_forward.__get__(pipe.unet, type(pipe.unet))

真实落地方案

1. Triton Inference Server 动态批处理

配置 config.pbtxt 启用 dynamic batching，结合 INT8 量化。

name: "stable_diffusion"
platform: "pytorch_libtorch"
max_batch_size: 16
dynamic_batching {
  max_queue_delay_microseconds: 50000
}
instance_group [{ count: 2, kind: KIND_GPU }]

客户端示例：

# client.py
import tritonclient.http as httpclient
from diffusers.utils import load_image
import numpy as np, torch

client = httpclient.InferenceServerClient(url="localhost:8000")
prompt = "a dog in a hat"
inputs = [
    httpclient.InferInput("PROMPT", [1], "BYTES"),
    httpclient.InferInput("LATENT", [1, 4, 64, 64], "FP16")
]
inputs[0].set_data_from_numpy(np.array([prompt.encode()], dtype=object))
inputs[1].set_data_from_numpy(torch.randn(1, 4, 64, 64, dtype=torch.float16).cpu().numpy())
result = client.infer("stable_diffusion", inputs)
image = result.as_numpy("IMAGE")

2. torch.compile 编译加速

PyTorch 2.0 提供编译优化，需关闭 guard。

# compile_pipe.py
import torch, warnings
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to("cuda")
pipe.unet = torch.compile(pipe.unet, mode="max-autotune", fullgraph=True, dynamic=False, disable=True)

3. ONNX Runtime CPU fallback

无 N 卡场景下，U-Net 转 ONNX，VAE 保留 GPU。

# onnx_fallback.py
import torch, onnxruntime as ort
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
pipe.unet.eval()
latent = torch.randn(1, 4, 64, 64, dtype=torch.float32)
t = torch.tensor([500], dtype=torch.int64)

torch.onnx.export(
    pipe.unet, (latent, t), "unet.onnx",
    input_names=["latent", "t"], output_names=["noise"],
    dynamic_axes={"latent": {0: "B"}, "noise": {0: "B"}}, opset_version=14
)

ort_sess = ort.InferenceSession("unet.onnx", providers=["CPUExecutionProvider"])

def ort_unet(latent, t):
    return torch.tensor(ort_sess.run(None, {"latent": latent.cpu().numpy(), "t": t.cpu().numpy()})[0])

pipe.unet.forward = ort_unet

常见问题排查

1. 显存 OOM

检查 attention mask 是否占用过多显存。将 mask 转为 bool 类型并半精度。

mask = mask.to(torch.bool).half()

2. AMP 配置错误

自定义节点可能意外关闭自动混合精度（AMP），导致速度下降。

from torch.cuda.amp import autocast
# 确保 enabled=True
with autocast(enabled=True): output = custom_op(input)

3. CUDA 版本冲突

RuntimeError 通常由 xformers、torch、triton 版本不兼容引起，需降级或升级匹配版本。

优化建议总结

1. Warm up：首次推理前跑空图预热 kernel。
2. medvram 模式：拆分 VAE 切片以适配低显存。
3. Prompt Embed 缓存：Redis 存储文本向量，避免重复编码。
4. FlashAttention-2：替换 Attention 处理器以提升速度并降低显存。

真正的提速依赖于对数据流、控制流及硬件特性的深入理解，而非依赖一键脚本。