From f8f501f2e8738e3f0bcee17443446f0a292f5b4b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=8E=8B=E9=B9=A4=E7=94=B7?= Date: Sat, 29 Aug 2026 14:42:49 +0800 Subject: [PATCH] [diffusion] fix: fix nvfp4 diffusion models on sm_120 (RTX PRO 6000 / RTX 50xx) (#35739) Co-authored-by: Mick --- .../layers/quantization/modelopt_quant.py | 16 ++---- .../multimodal_gen/runtime/platforms/cuda.py | 5 +- .../test/unit/test_modelopt_fp4_backend.py | 56 +++++++++++++++++++ 3 files changed, 65 insertions(+), 12 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_modelopt_fp4_backend.py diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py index 38e8e4560..1b703f50d 100755 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py @@ -678,18 +678,12 @@ class ModelOptFp4LinearMethod(LinearMethodBase): padded_scales = torch.zeros((B, M_padded, K_padded), dtype=scales.dtype) padded_scales[:B, :M, :K] = scales - _, flashinfer_backend = _get_fp4_gemm_op() - uses_flux1_scale_layout = not getattr( - self.quant_config, "checkpoint_uses_packed_qkv", False - ) and getattr(layer, "prefix", "").startswith( - ("transformer_blocks.", "single_transformer_blocks.") + # Every FP4 GEMM reachable here reads block scales in the 128x4 TMA layout; + # trtllm, the one backend wanting its own shuffled layout, returned above. + padded_scales = padded_scales.reshape( + B, M_padded // 128, 4, 32, K_padded // 4, 4 ) - if flashinfer_backend is None or uses_flux1_scale_layout: - # CUTLASS and FLUX.1 CUDNN paths need the TMA scale layout. - padded_scales = padded_scales.reshape( - B, M_padded // 128, 4, 32, K_padded // 4, 4 - ) - padded_scales = padded_scales.permute(0, 1, 4, 3, 2, 5) + padded_scales = padded_scales.permute(0, 1, 4, 3, 2, 5) padded_scales = padded_scales.contiguous().cuda() padded_scales = ( diff --git a/python/sglang/multimodal_gen/runtime/platforms/cuda.py b/python/sglang/multimodal_gen/runtime/platforms/cuda.py index 2dc303448..590cb6cf8 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/cuda.py +++ b/python/sglang/multimodal_gen/runtime/platforms/cuda.py @@ -443,7 +443,10 @@ class CudaPlatformBase(Platform): @lru_cache(maxsize=1) def get_modelopt_flashinfer_fp4_backend(cls) -> str: backend = envs.SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND - default_backend = "trtllm" + # flashinfer.mm_fp4 rejects backend="trtllm" on sm_120 ("does not support + # backend 'trtllm' with capability 120"); "auto" resolves to its sm_12x + # NVFP4 kernel there. + default_backend = "auto" if cls.is_sm120() else "trtllm" if backend is None: return default_backend diff --git a/python/sglang/multimodal_gen/test/unit/test_modelopt_fp4_backend.py b/python/sglang/multimodal_gen/test/unit/test_modelopt_fp4_backend.py new file mode 100644 index 000000000..db3a9109c --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_modelopt_fp4_backend.py @@ -0,0 +1,56 @@ +import unittest +from unittest.mock import patch + +from sglang.multimodal_gen.runtime.platforms.cuda import CudaPlatform + +ENV_PATH = ( + "sglang.multimodal_gen.runtime.platforms.cuda.envs." + "SGLANG_DIFFUSION_FLASHINFER_FP4_GEMM_BACKEND" +) + + +def _backend(env_value, *, sm120): + CudaPlatform.get_modelopt_flashinfer_fp4_backend.cache_clear() + try: + with patch(ENV_PATH, env_value), patch.object( + CudaPlatform, "is_sm120", classmethod(lambda cls: sm120) + ): + return CudaPlatform.get_modelopt_flashinfer_fp4_backend() + finally: + CudaPlatform.get_modelopt_flashinfer_fp4_backend.cache_clear() + + +class TestModeloptFp4BackendDefault(unittest.TestCase): + """The default FP4 GEMM backend must exist on the device. + + flashinfer.mm_fp4 raises "does not support backend 'trtllm' with capability + 120" on sm_120, so defaulting to trtllm there makes every NVFP4 model fail to + run at all. + """ + + def test_sm120_does_not_default_to_a_backend_it_lacks(self): + self.assertEqual(_backend(None, sm120=True), "auto") + + def test_other_gpus_keep_the_trtllm_default(self): + self.assertEqual(_backend(None, sm120=False), "trtllm") + + def test_explicit_backends_are_passed_through(self): + for env_value, expected in [ + ("trtllm", "trtllm"), + ("flashinfer_trtllm", "trtllm"), + ("cutlass", "cutlass"), + ("flashinfer_cutlass", "cutlass"), + ("cudnn", "cudnn"), + ("flashinfer_cudnn", "cudnn"), + ("auto", "auto"), + ]: + with self.subTest(env=env_value): + self.assertEqual(_backend(env_value, sm120=True), expected) + + def test_unknown_backend_falls_back_to_the_platform_default(self): + self.assertEqual(_backend("nonsense", sm120=True), "auto") + self.assertEqual(_backend("nonsense", sm120=False), "trtllm") + + +if __name__ == "__main__": + unittest.main()