[diffusion] Prefer cuDNN SDPA over FA4 for dense attention on sm_100 (B200) (#33655)

Co-authored-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>

python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py

@@ -169,6 +169,12 @@ class DynamicCudnnSDPAImpl(SDPAImpl):
 
         self.causal = causal
         self.head_size = head_size
+        self._is_sm100 = (
+            torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10
+        )
+        # Set once cuDNN SDPA raised for this layer; permanently pins the
+        # fail-safe FA path so we do not retry a failing kernel every step.
+        self._cudnn_failed = False
         if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 10:
             set_fa_ver(4)
         self.cudnn_impl = CudnnSDPAImpl(
@@ -193,7 +199,7 @@ class DynamicCudnnSDPAImpl(SDPAImpl):
     def _use_cudnn_sdpa(
         self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor
     ) -> bool:
-        if self.causal:
+        if self.causal or self._cudnn_failed:
             return False
         if query.device.type != "cuda":
             return False
@@ -201,6 +207,11 @@ class DynamicCudnnSDPAImpl(SDPAImpl):
             return False
         if query.shape[2] != key.shape[2]:
             return False
+        if self._is_sm100:
+            # B200/sm_100: cuDNN SDPA measured 1.25-1.5x faster than FA4 CuTe
+            # for dense non-causal diffusion attention, both self-attn
+            # (Sq == Skv, up to S=506K) and cross-attn (Skv = text len).
+            return True
         if query.shape[1] != key.shape[1]:
             return False
         return query.shape[-1] == 64 and query.shape[1] == 1024 and query.shape[0] >= 4
@@ -211,7 +222,21 @@ class DynamicCudnnSDPAImpl(SDPAImpl):
         key: torch.Tensor,
         value: torch.Tensor,
         attn_metadata: AttentionMetadata,
+        **kwargs,
     ) -> torch.Tensor:
-        if self._use_cudnn_sdpa(query, key, value):
-            return self.cudnn_impl.forward(query, key, value, attn_metadata)
-        return self.fa_impl.forward(query, key, value, attn_metadata)
+        # LSE is only produced by the FA impl (e.g. ring attention).
+        if not kwargs.get("return_softmax_lse") and self._use_cudnn_sdpa(
+            query, key, value
+        ):
+            try:
+                return self.cudnn_impl.forward(query, key, value, attn_metadata)
+            except RuntimeError as e:
+                # cuDNN raises "No available kernel" for some shapes; pin the
+                # FA fail-safe path for this layer and keep going.
+                logger.warning(
+                    "cuDNN SDPA failed (%s); falling back to FlashAttention " "for %s.",
+                    e,
+                    type(self).__name__,
+                )
+                self._cudnn_failed = True
+        return self.fa_impl.forward(query, key, value, attn_metadata, **kwargs)

python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py

@@ -1,21 +1,32 @@
 import copy
 import logging
+from contextlib import nullcontext
 from typing import Any
 
 import torch
 
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.layers.attention.selector import (
+    component_attn_backend_context_manager,
+    get_component_forced_attn_backend,
+    get_global_forced_attn_backend,
+)
 from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
     ComponentLoader,
 )
 from sglang.multimodal_gen.runtime.loader.fsdp_load import maybe_load_fsdp_model
 from sglang.multimodal_gen.runtime.loader.transformer_load_utils import (
+    TransformerQuantLoadSpec,
     resolve_transformer_quant_load_spec,
     resolve_transformer_safetensors_to_load,
 )
 from sglang.multimodal_gen.runtime.loader.utils import _normalize_component_type
 from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan
 from sglang.multimodal_gen.runtime.models.registry import ModelRegistry
+from sglang.multimodal_gen.runtime.platforms import (
+    AttentionBackendEnum,
+    current_platform,
+)
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
     get_diffusers_component_config,
@@ -28,6 +39,21 @@ _is_npu = is_npu()
 logger = init_logger(__name__)
 
 
+def _default_quantized_attention_backend(
+    quant_spec: TransformerQuantLoadSpec, server_args: ServerArgs
+) -> AttentionBackendEnum | None:
+    """Preserve stable NVFP4 numerics unless the user selected a backend."""
+    if not current_platform.is_blackwell() or not quant_spec.is_modelopt_fp4:
+        return None
+    if (
+        get_global_forced_attn_backend() is not None
+        or get_component_forced_attn_backend() is not None
+        or server_args.attention_backend is not None
+    ):
+        return None
+    return AttentionBackendEnum.FA
+
+
 def _warn_if_expected_param_dtype_missing(
     model: torch.nn.Module, expected_dtype: torch.dtype | None
 ) -> None:
@@ -178,23 +204,41 @@ class TransformerLoader(ComponentLoader):
             component_cpu_offload=bool(component_server_args.dit_cpu_offload),
         )
 
-        # Load the model using FSDP loader
-        model = maybe_load_fsdp_model(
-            model_cls=model_cls,
-            init_params=init_params,
-            weight_dir_list=safetensors_list,
-            device=local_torch_device,
-            hsdp_replicate_dim=server_args.hsdp_replicate_dim,
-            hsdp_shard_dim=server_args.hsdp_shard_dim,
-            cpu_offload=component_server_args.dit_cpu_offload,
-            pin_cpu_memory=component_server_args.pin_cpu_memory,
-            fsdp_inference=component_server_args.use_fsdp_inference,
-            param_dtype=quant_spec.param_dtype,
-            reduce_dtype=torch.float32,
-            output_dtype=None,
-            strict=False,
-            weight_load_plan=weight_load_plan,
+        quantized_attn_backend = _default_quantized_attention_backend(
+            quant_spec, component_server_args
         )
+        if quantized_attn_backend is not None:
+            logger.info(
+                "Using %s attention for ModelOpt NVFP4 to preserve output precision",
+                quantized_attn_backend.name.lower(),
+            )
+        attn_backend_context = (
+            component_attn_backend_context_manager(
+                quantized_attn_backend, component_name=component_name
+            )
+            if quantized_attn_backend is not None
+            else nullcontext()
+        )
+
+        # Model construction resolves attention implementations, so apply the
+        # quantization-specific default around FSDP initialization and loading.
+        with attn_backend_context:
+            model = maybe_load_fsdp_model(
+                model_cls=model_cls,
+                init_params=init_params,
+                weight_dir_list=safetensors_list,
+                device=local_torch_device,
+                hsdp_replicate_dim=server_args.hsdp_replicate_dim,
+                hsdp_shard_dim=server_args.hsdp_shard_dim,
+                cpu_offload=component_server_args.dit_cpu_offload,
+                pin_cpu_memory=component_server_args.pin_cpu_memory,
+                fsdp_inference=component_server_args.use_fsdp_inference,
+                param_dtype=quant_spec.param_dtype,
+                reduce_dtype=torch.float32,
+                output_dtype=None,
+                strict=False,
+                weight_load_plan=weight_load_plan,
+            )
 
         # post-hooks (e.g., patch scales (nunchaku))
         for post_load_hook in quant_spec.post_load_hooks:

python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py

@@ -131,6 +131,10 @@ class TransformerQuantLoadSpec:
             return self.quant_config
         return self.nunchaku_config
 
+    @property
+    def is_modelopt_fp4(self) -> bool:
+        return _get_quant_config_name(self.quant_config) == "modelopt_fp4"
+
 
 class _TransformerQuantAdapter:
     def prepare(self) -> None:

python/sglang/multimodal_gen/runtime/platforms/cuda.py

@@ -519,6 +519,17 @@ class CudaPlatformBase(Platform):
     ) -> str:
         if selected_backend is None:
             target_backend = cls._resolve_default_attn_backend()
+            if target_backend == AttentionBackendEnum.FA and cls.is_blackwell():
+                # cuDNN SDPA is 1.25-1.5x faster than the FA4 CuTe kernels on
+                # sm_100 for dense diffusion attention; DYNAMIC_CUDNN_SDPA
+                # keeps FA as the fallback for causal/unsupported shapes and
+                # cuDNN runtime errors.
+                fa_cls_str = cls._resolve_flash_attention_backend_cls_str(
+                    target_backend, head_size, dtype
+                )
+                if fa_cls_str == _SDPA_BACKEND_CLS_STR:
+                    return fa_cls_str
+                return _DYNAMIC_CUDNN_SDPA_BACKEND_CLS_STR
         else:
             resolver = _CUDA_ATTENTION_BACKEND_RESOLVERS.get(selected_backend)
             if resolver is None:

python/sglang/multimodal_gen/test/unit/test_cuda_attention_backend.py

@@ -91,6 +91,19 @@ class TestCudaAttentionBackendSelection(unittest.TestCase):
 
         prepare_flash_attention.assert_called_once_with()
 
+    def test_default_backend_prefers_dynamic_cudnn_sdpa_on_blackwell(self):
+        FakeCudaPlatform.is_blackwell_device = True
+
+        with patch.object(
+            FakeCudaPlatform,
+            "_prepare_flash_attention_for_blackwell",
+            return_value=True,
+        ):
+            self.assertEqual(
+                self.resolve(None),
+                "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.DynamicCudnnSDPABackend",
+            )
+
     def test_invalid_backend_raises(self):
         with self.assertRaisesRegex(ValueError, "Invalid attention backend"):
             self.resolve(AttentionBackendEnum.AITER_SAGE)

python/sglang/multimodal_gen/test/unit/test_transformer_quant.py

@@ -57,9 +57,11 @@ from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
 )
 from sglang.multimodal_gen.runtime.loader.component_loaders import transformer_loader
 from sglang.multimodal_gen.runtime.loader.component_loaders.transformer_loader import (
+    _default_quantized_attention_backend,
     _warn_if_expected_param_dtype_missing,
 )
 from sglang.multimodal_gen.runtime.loader.transformer_load_utils import (
+    TransformerQuantLoadSpec,
     _filter_duplicate_precision_variant_safetensors,
     _Flux2Nvfp4FallbackAdapter,
     _needs_device_weight_postprocess,
@@ -69,6 +71,7 @@ from sglang.multimodal_gen.runtime.loader.transformer_load_utils import (
 )
 from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan
 from sglang.multimodal_gen.runtime.models.dits.flux import FluxSingleTransformerBlock
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
 from sglang.multimodal_gen.runtime.utils.quantization_utils import (
     build_nvfp4_config_from_safetensors_list,
     get_quant_config,
@@ -120,6 +123,40 @@ class TestTransformerQuantHelpers(unittest.TestCase):
         defaults.update(overrides)
         return SimpleNamespace(**defaults)
 
+    def test_modelopt_fp4_uses_fa_by_default_on_blackwell(self):
+        quant_spec = TransformerQuantLoadSpec([], _FakeQuantConfig(), None, None)
+        server_args = SimpleNamespace(attention_backend=None)
+
+        with (
+            patch.object(
+                transformer_loader.current_platform, "is_blackwell", return_value=True
+            ),
+            patch.object(
+                transformer_loader,
+                "get_global_forced_attn_backend",
+                return_value=None,
+            ),
+            patch.object(
+                transformer_loader,
+                "get_component_forced_attn_backend",
+                return_value=None,
+            ),
+        ):
+            backend = _default_quantized_attention_backend(quant_spec, server_args)
+
+        self.assertEqual(backend, AttentionBackendEnum.FA)
+
+    def test_modelopt_fp4_preserves_explicit_attention_backend(self):
+        quant_spec = TransformerQuantLoadSpec([], _FakeQuantConfig(), None, None)
+        server_args = SimpleNamespace(attention_backend="dynamic_cudnn_sdpa")
+
+        with patch.object(
+            transformer_loader.current_platform, "is_blackwell", return_value=True
+        ):
+            backend = _default_quantized_attention_backend(quant_spec, server_args)
+
+        self.assertIsNone(backend)
+
     def test_resolve_transformer_safetensors_to_load_uses_single_override_file(self):
         with tempfile.NamedTemporaryFile(suffix=".safetensors") as f:
             server_args = self._make_server_args(transformer_weights_path=f.name)