# PR #22428 完整报告

- 仓库：`sgl-project/sglang`
- 标题：[AMD] Diffusion - Enabel rocm miopen tuning on vae
- 合并时间：2026-04-11 13:47
- 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/22428

---

# 执行摘要
本 PR 为 ROCm 平台上的扩散模型 VAE 解码阶段引入了 MIOpen 自动调优功能，通过设置 `SGLANG_USE_ROCM_CUDNN_BENCHMARK` 环境变量启用 `torch.backends.cudnn.benchmark`，实现了约 30-35% 的 VAE 解码速度提升，总推理时间改善约 2.4-3.0%。优化针对 Conv3d 密集操作，在保持图像质量达标的前提下，显著提升了 AMD 平台的性能表现。

# 功能与动机
**为什么做？**在 ROCm 平台上，VAE 解码阶段（特别是 Conv3d 层）的性能有优化空间。通过启用 MIOpen 的自动调优（即 `cudnn.benchmark`），系统可以为每个不同的输入形状自动选择最快的卷积算法，从而加速推理。PR body 中提供了详细的基准数据：
> "Enabling `cudnn.benchmark` allows MIOpen to auto-select the fastest convolution algorithm for each distinct input shape, primarily benefiting the Conv3d-heavy VAE decode stage."

测试显示，在 FP8 精度下，VAE 解码时间从 27.86 秒减少到 18.23 秒，提升 34.6%；总速度提升约 2.8%。准确性指标（SSIM、PSNR、LPIPS）均满足项目默认阈值。

# 实现拆解
改动集中在两个文件，按模块拆解如下：

| 文件 | 模块 | 关键变更 | 说明 |
|------|------|----------|------|
| `python/sglang/multimodal_gen/envs.py` | 环境变量管理 | 新增 `SGLANG_USE_ROCM_CUDNN_BENCHMARK` 变量 | 提供用户配置入口，默认禁用，需显式启用。 |
| `python/sglang/multimodal_gen/runtime/platforms/rocm.py` | ROCm 平台优化 | 在 `optimize_vae` 方法中启用 `cudnn.benchmark` | 核心优化逻辑，仅当环境变量为真且未启用时设置全局标志。 |

关键代码逻辑位于 `optimize_vae` 方法：
```python
if envs.SGLANG_USE_ROCM_CUDNN_BENCHMARK and not torch.backends.cudnn.benchmark:
    torch.backends.cudnn.benchmark = True
    logger.info(
        "Enabled cudnn.benchmark (MIOpen auto-tuning) for VAE conv layers"
    )
```

# 评论区精华
review 中，gemini-code-assist[bot] 指出了设计上的一个重要细节：
> "`torch.backends.cudnn.benchmark` is a global flag that affects all convolution operations in the process. It's better to clarify its global scope to avoid confusion if other models (like the Transformer) are affected by the tuning overhead."

该评论建议更新注释和日志，以明确此标志的全局性，避免用户误以为优化仅局限于 VAE 层。PR 作者通过重命名环境变量和合并 main 分支进行了调整，但未直接采纳文本修改建议，这可能留下潜在的误解风险。

# 风险与影响
**技术风险**：
1. **全局配置影响**：`torch.backends.cudnn.benchmark` 是进程级设置，启用后会影响所有卷积操作。如果同一进程中运行其他模型（如 Transformer）且输入形状动态变化，可能导致持续的重新调优，带来性能开销。
2. **文档缺失**：环境变量的全局性未在代码注释或日志中充分说明，用户可能低估其对系统其他部分的影响。

**影响评估**：
- **用户**：ROCm 平台用户可通过设置环境变量获得显著的 VAE 解码加速，但需注意可能影响其他模型组件。
- **系统**：优化仅针对扩散模型的 VAE 阶段，但标志的全局性意味着更广泛的影响。
- **团队**：延续了 AMD 平台的性能优化路线，与近期多个 PR（如 #22228、#22264）共同完善 ROCm 支持。

# 关联脉络
从近期历史 PR 看，本 PR 是 **AMD 平台优化 **和 **扩散模型改进 **两条脉络的交汇点：
- **AMD 平台**：PR#22228（修复 CI 超时）、#22264（升级 Aiter 依赖）等共同提升 ROCm 的稳定性和性能。
- **扩散模型**：PR#22507（改进 CI 可读性）、#22560（修复单元测试）等聚焦于测试和质量保障。

本 PR 通过环境变量控制优化，保持了向后兼容性，符合项目渐进式优化的风格。未来若需扩展自动调优到其他组件，可参考此设计模式。