执行摘要
本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方法：
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分支进行了调整，但未直接采纳文本修改建议，这可能留下潜在的误解风险。
风险与影响
技术风险：
- 全局配置影响：
torch.backends.cudnn.benchmark是进程级设置，启用后会影响所有卷积操作。如果同一进程中运行其他模型（如Transformer）且输入形状动态变化，可能导致持续的重新调优，带来性能开销。
- 文档缺失：环境变量的全局性未在代码注释或日志中充分说明，用户可能低估其对系统其他部分的影响。
影响评估：
- 用户：ROCm平台用户可通过设置环境变量获得显著的VAE解码加速，但需注意可能影响其他模型组件。
- 系统：优化仅针对扩散模型的VAE阶段，但标志的全局性意味着更广泛的影响。
- 团队：延续了AMD平台的性能优化路线，与近期多个PR（如#22228、#22264）共同完善ROCm支持。
关联脉络
从近期历史PR看，本PR是AMD平台优化和扩散模型改进两条脉络的交汇点：
- AMD平台：PR#22228（修复CI超时）、#22264（升级Aiter依赖）等共同提升ROCm的稳定性和性能。
- 扩散模型：PR#22507（改进CI可读性）、#22560（修复单元测试）等聚焦于测试和质量保障。
本PR通过环境变量控制优化，保持了向后兼容性，符合项目渐进式优化的风格。未来若需扩展自动调优到其他组件，可参考此设计模式。
参与讨论