Prhub

#22091 [diffusion] Default NVFP4 to CUTLASS and add all-model shape benchmarks

原始 PR 作者 BBuf 合并时间 2026-04-04 16:14 文件变更 7 提交数 5 评论 3 代码增减 +3064 / -37

执行摘要

将扩散模型 NVFP4 后端默认值切换为 CUTLASS 并添加全模型基准测试,优化 Blackwell GPU 性能。

PR body中指出:'On Blackwell, comfy-kitchen is not the best default once we look beyond FLUX.' 基准测试显示CUTLASS在252/265形状中获胜,FLOP加权时间比为1.013x vs comfy的1.095x,因此切换默认后端以优化扩散模型推理性能。

建议工程师精读cuda.pyget_modelopt_fp4_gemm_op函数的变更,学习基于基准测试的后端选择策略;测试团队可参考新增基准测试脚本扩展CI覆盖。

讨论亮点

reviewer mickqian建议:'may consider adding more nvfp4 test to CI later',表明未来可增强测试覆盖,但当前PR无争议或深入讨论。

实现拆解

1) 新增基准测试文件bench_diffusion_nvfp4_scaled_mm.py和形状库diffusion_nvfp4_shapes.json,用于性能评估和形状管理。
2) 修改cuda.py中的get_modelopt_fp4_gemm_op函数,默认返回CUTLASS后端,并新增get_modelopt_flashinfer_fp4_backend函数控制FlashInfer后端。
3) 更新interface.py移除不再使用的方法,简化平台接口。
4) 在envs.py添加环境变量SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND,支持后端回退配置。
5) 调整transformer_load_utils.py中FLUX.2 NVFP4相关逻辑,确保兼容性。

文件 模块 状态 重要度
python/sglang/jit_kernel/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py benchmark/diffusion added 5.0
python/sglang/jit_kernel/benchmark/diffusion/diffusion_nvfp4_shapes.json benchmark/diffusion added 4.0
python/sglang/multimodal_gen/runtime/platforms/cuda.py multimodal_gen/runtime/platforms modified 6.0

关键符号

get_modelopt_fp4_gemm_op get_modelopt_flashinfer_fp4_backend benchmark_provider

分析完成后,这里会展示 LLM 生成的相对完整源码片段和详细注释。

评论区精华

CI 测试扩展建议 测试

reviewer mickqian 建议未来添加更多 NVFP4 测试到 CI,以增强测试覆盖

结论:无明确结论,仅为后续改进建议,未在 PR 中实现 · suggestion

风险与影响

技术风险:默认后端变更可能导致在comfy获胜的12.7个形状上性能下降,但环境变量SGLANG_DIFFUSION_NVFP4_LINEAR_BACKEND=comfy可覆盖。新增基准测试和代码可能引入bug,但已通过python3 -m compileall和实际运行验证。兼容性:保持comfy-kitchen为可选后端,不影响现有显式配置的用户。

对用户影响:使用Blackwell GPU运行扩散模型的用户将获得推理性能提升。系统影响:优化NVFP4量化路径,减少延迟,增强基准测试基础设施。团队影响:提供数据驱动决策范例,支持后续性能调优和测试扩展。

默认后端变更 测试覆盖待增强

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论