From ba12a16a62f1702fbb316bc1969113dfddbe8ef4 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Thu, 6 Aug 2026 08:49:10 +0800 Subject: [PATCH] [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> --- .../runtime/layers/attention/backends/sdpa.py | 33 +++++++- .../component_loaders/transformer_loader.py | 76 +++++++++++++++---- .../runtime/loader/transformer_load_utils.py | 4 + .../multimodal_gen/runtime/platforms/cuda.py | 11 +++ .../test/unit/test_cuda_attention_backend.py | 13 ++++ .../test/unit/test_transformer_quant.py | 37 +++++++++ 6 files changed, 154 insertions(+), 20 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py index 901446c0c..77807ad68 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py index eca50e5f0..83374ba6a 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py index 8128f7cd0..e13be4541 100644 --- a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/platforms/cuda.py b/python/sglang/multimodal_gen/runtime/platforms/cuda.py index 86613223f..397286073 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/cuda.py +++ b/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: diff --git a/python/sglang/multimodal_gen/test/unit/test_cuda_attention_backend.py b/python/sglang/multimodal_gen/test/unit/test_cuda_attention_backend.py index fb13bd493..4dae9e1ba 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cuda_attention_backend.py +++ b/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) diff --git a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py index 8c98efeaf..f0575eb3a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py +++ b/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)