[diffusion] fix: fix nvfp4 diffusion models on sm_120 (RTX PRO 6000 / RTX 50xx) (#35739)

Co-authored-by: Mick <mickjagger19@icloud.com>

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 = (

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
 

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()