11 - Diffusion 推理系统优化

Diffusion 的主要开销来自重复 denoising、attention 中间矩阵和 VAE decode，优化位置与 LLM 不同。本章用 T16、T17 的实验数据检查六项具体做法。

1. Prompt Embedding Cache（提示词嵌入缓存）

扩散模型在推理期间的 text encoding 只执行 1 次：输入的 prompt 经 T5/CLIP 编码为固定长度的 embedding，在后续所有 denoising step 中复用。如果用户连续生成多张图像，且 prompt 不变，那 text encoding 就是完全重复的计算。

实验数据（T16）

Cache key 设计：model_id + text_encoder hash + prompt hash + negative_prompt hash + max_sequence_length + dtype + device/offload 策略。这确保了跨模型隔离，1 个 FLUX-schnell 的 cache 不会被 SD3-Medium 误命中。实验中验证了跨模型隔离：10 条共享 prompt 在 2 个模型间全部正确 cache miss。

限制：embedding cache 能省掉 text encoding，但这部分通常只占总推理时间的 1-5%，主要耗时仍在 denoiser loop。实际收益还包括减少重复显存分配和 CPU→GPU 传输。

2. Latent Buffer Manager（潜空间缓冲管理器）

扩散推理的去噪循环中，每步 denoiser forward 的输出（噪声预测/矢量场）都要写回当前 latent。如果每步都用 torch.empty() 创建新 tensor，会产生 GPU 显存碎片。Latent buffer manager 通过预分配、in-place reset 和 ping-pong swap 避免重复分配。

实验数据（T16）

主要结论：in-place reset 节省了 91.8% 的 allocation 次数（5 vs 61），同时每步延迟也低 40%（0.231 vs 0.387 ms）。峰值已分配内存不变（两者都预分配，区别在于是否复用），但 allocation 次数的减少直接转化为更少的 CUDA 同步和更稳定的推理延迟。

3. Scheduler Step Profiling（调度器步数分析）

scheduler 的选择直接影响推理质量与速度的平衡。本 benchmark 对比了 Euler（等步长）和 RectifiedFlow（log-linear 步长）在 4/8/16/28/50 步下的表现。

实验数据（T16）

主要发现：延迟与步数近似线性（\(R^2 \approx 1.000\)）。\(50\) 步 \(\approx 4\) 步 \(\times 12.5\)。Euler 和 RectifiedFlow 每步都执行 1 次 denoiser forward + ODE update，计算量接近，主要差异是步长间距。在这组实验中，减少步数比切换 scheduler 类型更能降低延迟。

4. CFG Batching（分类器无关引导批处理）

CFG（Classifier-Free Guidance）在每步需要跑 2 次 denoiser forward：1 次 conditional（带 prompt embedding），1 次 unconditional（空 prompt embedding）。结果按 \(v_{\mathrm{cond}} + \mathrm{cfg\_scale}\left(v_{\mathrm{cond}} - v_{\mathrm{uncond}}\right)\) 组合。

2 种策略：Sequential CFG 分别执行 cond 和 uncond forward；Batched CFG 将 2 组 latent 拼成双倍 batch，用 1 次 forward 同时处理。

实验数据（T17）

数值差异 0.00e+00（理论上完全一致，微小差异仅来自浮点 accumulate order）。中等显存配置 策略：剩余 > 6GB → Batched CFG；3-6GB → 视具体情况；< 3GB → Sequential CFG。

5. Attention Memory（注意力显存估算）

这是扩散推理中 最大的瓶颈。DiT/MMDiT 的 full attention 复杂度为 \(O(N^2)\)，其中 \(N\) 是 latent token 数。1024×1024 图像 → latent 128×128 → patch_size=2 → 4096 tokens → attention matrix 32 MB (fp16)，可接受。但 \(2048^2\) → 16384 tokens → 512 MB。

实验数据（T17）

\(O(N^2)\) 复杂度验证：\(N \to 2N\) 时，\(N^2 \to 4N^2\)。\(512^2\) 图像用 standard attention 安全，\(1024^2\) 推荐 memory-efficient attention，\(2048^2\) 和长视频 必须 memory-efficient attention。

6. VAE Tiling（VAE 分块解码）

VAE tiling 在应用层把 latent 或像素切成多个 tile，逐块 decode，再用 overlap blending 处理接缝。它属于显式 chunk decode，与 torch.compile、flash-attn 解决的不是同一层问题。

实验数据（T17）

典型开销：overlap=4 latent pixels → 额外 10-20% 计算量。受限显存配置下 VAE tiling 几乎总是必要的。推荐 tile 大小：\(1024^2\) → 64×64 latent tile；\(2048^2\) → 32×32 latent tile。

7. 与 LLM KV Cache 的根本差异

LLM 的 KV cache 依赖 1 个前提：旧 token 保持不变，新 token 只需读取历史。扩散每个 denoising step 都会更新整块 latent，因此上一轮的 K/V 不能直接作为下一轮缓存。两者 1 个追加状态，1 个重写状态。

那哪里相似？在 prompt/text encoder 层面。扩散和 LLM 推理的第 1 步都是 text encode，这个步骤的 cache 策略是通用的：相同的 prompt → 相同的 embedding → 避免重复编码。这 1.2 者都会做。但进入主循环后，优化路径就完全不同了。