diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py index e49c024af..6d8d5bdc0 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loader.py index a704c4297..d29c03bf3 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loader.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/platforms/interface.py b/python/sglang/multimodal_gen/runtime/platforms/interface.py index 283381ab9..6c7297076 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/interface.py +++ b/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):