From 6b96f8341d52552d1d68744c59f80f882c1b94bc Mon Sep 17 00:00:00 2001 From: Bingxu Chen Date: Thu, 9 Apr 2026 13:32:49 +0800 Subject: [PATCH] [AMD] Fix multimodal diffusion test crash on ROCm by falling back to SDPA (#22335) --- python/sglang/jit_kernel/flash_attention_v3.py | 2 ++ .../multimodal_gen/runtime/platforms/rocm.py | 17 +++++++++++++---- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/python/sglang/jit_kernel/flash_attention_v3.py b/python/sglang/jit_kernel/flash_attention_v3.py index 23018961d..d9b4ba01a 100644 --- a/python/sglang/jit_kernel/flash_attention_v3.py +++ b/python/sglang/jit_kernel/flash_attention_v3.py @@ -78,6 +78,8 @@ def _is_fa3_supported(device=None) -> bool: # https://docs.nvidia.com/cuda/cuda-c-programming-guide/#shared-memory-8-x # And for sgl-kernel right now, we can build fa3 on sm80/sm86/sm89/sm90a. # That means if you use A100/A*0/L20/L40/L40s/4090 you can use fa3. + if torch.version.cuda is None: + return False return (torch.version.cuda >= "12.3") and ( torch.cuda.get_device_capability(device)[0] == 9 or torch.cuda.get_device_capability(device)[0] == 8 diff --git a/python/sglang/multimodal_gen/runtime/platforms/rocm.py b/python/sglang/multimodal_gen/runtime/platforms/rocm.py index 4f8952dc7..2ed9a86c3 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/rocm.py +++ b/python/sglang/multimodal_gen/runtime/platforms/rocm.py @@ -149,17 +149,26 @@ class RocmPlatform(Platform): try: import flash_attn # noqa: F401 + from sglang.jit_kernel.flash_attention_v3 import _is_fa3_supported from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import ( # noqa: F401 FlashAttentionBackend, ) - supported_sizes = FlashAttentionBackend.get_supported_head_sizes() - if head_size not in supported_sizes: + if not _is_fa3_supported(): logger.info( - "Cannot use FlashAttention-2 backend for head size %d.", - head_size, + "FlashAttention backend now dispatches through FA3 " + "(CUDA-only). Using Torch SDPA backend on ROCm." ) target_backend = AttentionBackendEnum.TORCH_SDPA + + if target_backend == AttentionBackendEnum.FA: + supported_sizes = FlashAttentionBackend.get_supported_head_sizes() + if head_size not in supported_sizes: + logger.info( + "Cannot use FlashAttention-2 backend for head size %d.", + head_size, + ) + target_backend = AttentionBackendEnum.TORCH_SDPA except ImportError: logger.info( "Cannot use FlashAttention backend because the "