[diffusion] fix: fix using upstream flash_attn on blackwell (#17111)

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

@@ -7,7 +7,10 @@ from typing import Any, List, Optional, Tuple
 import torch
 
 from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context
-from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.platforms import (
+    AttentionBackendEnum,
+    current_platform,
+)
 from sglang.srt.utils.custom_op import register_custom_op
 
 try:
@@ -19,6 +22,10 @@ try:
 except ImportError as e:
     raise e
 
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
 
 def maybe_contiguous(x: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
     return x.contiguous() if x is not None and x.stride(-1) != 1 else x
@@ -294,11 +301,18 @@ def flash_attn_varlen_func_op_lse(
 
 
 try:
-    from flash_attn_interface import (
-        flash_attn_varlen_func as flash_attn_varlen_func_upstream,
-    )
+    if current_platform.is_hopper():
+        from flash_attn_interface import (
+            flash_attn_varlen_func as flash_attn_varlen_func_upstream,
+        )
+    else:
+        flash_attn_varlen_func_upstream = None
+
 except Exception:
     flash_attn_varlen_func_upstream = None
+    logger.warning(
+        "flash_attn 3 package is not installed. It's recommended to install flash_attn3 on hopper, otherwise performance is sub-optimal"
+    )
 
 from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
     AttentionBackend,
@@ -306,9 +320,6 @@ from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend i
     AttentionMetadata,
     AttentionMetadataBuilder,
 )
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
-
-logger = init_logger(__name__)
 
 fa_ver = 3
 

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

@@ -422,7 +422,7 @@ class TextEncoderLoader(ComponentLoader):
         )
         model_config = get_diffusers_component_config(model_path=component_model_path)
         _clean_hf_config_inplace(model_config)
-        logger.info("HF model config: %s", model_config)
+        logger.debug("HF model config: %s", model_config)
 
         def is_not_first_encoder(module_name):
             return "2" in module_name
@@ -549,7 +549,7 @@ class ImageEncoderLoader(TextEncoderLoader):
         with open(os.path.join(component_model_path, "config.json")) as f:
             model_config = json.load(f)
         _clean_hf_config_inplace(model_config)
-        logger.info("HF model config: %s", model_config)
+        logger.debug("HF model config: %s", model_config)
 
         encoder_config = server_args.pipeline_config.image_encoder_config
         encoder_config.update_model_arch(model_config)
@@ -615,7 +615,7 @@ class VAELoader(ComponentLoader):
 
         server_args.model_paths["vae"] = component_model_path
 
-        logger.info("HF model config: %s", config)
+        logger.debug("HF model config: %s", config)
         vae_config = server_args.pipeline_config.vae_config
         vae_config.update_model_arch(config)
 

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

@@ -115,6 +115,13 @@ class Platform:
             return False
         return torch.cuda.get_device_capability()[0] == 10
 
+    @classmethod
+    @lru_cache(maxsize=1)
+    def is_hopper(cls):
+        if not cls.is_cuda_static():
+            return False
+        return torch.cuda.get_device_capability() == (9, 0)
+
     @classmethod
     @lru_cache(maxsize=1)
     def is_sm120(cls):