diff --git a/python/sglang/jit_kernel/flash_attention_v3.py b/python/sglang/jit_kernel/flash_attention_v3.py index 00524d5b4..26781b062 100644 --- a/python/sglang/jit_kernel/flash_attention_v3.py +++ b/python/sglang/jit_kernel/flash_attention_v3.py @@ -7,6 +7,7 @@ import torch from sglang.jit_kernel.utils import cache_once from sglang.kernel_api_logging import debug_kernel_api from sglang.srt.environ import envs +from sglang.srt.utils import get_device_capability, is_musa logger = logging.getLogger(__name__) @@ -89,12 +90,12 @@ 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 - ) + major, minor = get_device_capability() + if is_musa(): + return major >= 3 + if torch.version.cuda is not None and torch.version.cuda >= "12.3": + return major == 9 or major == 8 + return False @debug_kernel_api @@ -211,31 +212,30 @@ def flash_attn_varlen_func( "flash_attn at sgl-kernel is only supported on sm90 and above" ) - return _call_fa3_kernel( - _load_fa3_kernels()["flash_attn_varlen_func"], - q, - k, - v, - cu_seqlens_q, - cu_seqlens_k, - max_seqlen_q, - max_seqlen_k, - seqused_q, - seqused_k, - page_table, - softmax_scale, - causal, - qv, - q_descale, - k_descale, - v_descale, - window_size, - attention_chunk, - softcap, - num_splits, - pack_gqa, - sm_margin, - return_softmax_lse, - sinks, + return _load_fa3_kernels()["flash_attn_varlen_func"]( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + seqused_q=seqused_q, + seqused_k=seqused_k, + page_table=page_table, + softmax_scale=softmax_scale, + causal=causal, + qv=qv, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + window_size=window_size, + attention_chunk=attention_chunk, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + sm_margin=sm_margin, + return_softmax_lse=return_softmax_lse, + sinks=sinks, out=out, ) diff --git a/python/sglang/jit_kernel/tests/test_flash_attention_3.py b/python/sglang/jit_kernel/tests/test_flash_attention_3.py index e4687da9c..3f8447008 100644 --- a/python/sglang/jit_kernel/tests/test_flash_attention_3.py +++ b/python/sglang/jit_kernel/tests/test_flash_attention_3.py @@ -10,6 +10,7 @@ from einops import rearrange, repeat apply_rotary_emb = None +from sglang.jit_kernel.flash_attention_v3 import _is_fa3_supported from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=120, suite="stage-b-kernel-unit-1-gpu-large") @@ -21,21 +22,6 @@ def is_hopper(): return torch.cuda.get_device_properties(0).major == 9 -def is_fa3_supported(device=None) -> bool: - # There some fa3 FYI - # FA3 can fail without a enough shared memory for a some shapes, such as higher - # hidden_dim or some special cases. - # Right now, fa3 is supported for sm80/sm87 and sm86/sm89. The main different - # Between sm80/sm87 and sm86/sm89 is the shared memory size. you can follow the link below for more information - # 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. - return (torch.version.cuda >= "12.3") and ( - torch.cuda.get_device_capability(device)[0] == 9 - or torch.cuda.get_device_capability(device)[0] == 8 - ) - - DISABLE_BACKWARD = True # For CI test, we close them to True. # DISABLE_SPLIT = os.getenv("FLASH_ATTENTION_DISABLE_SPLIT", "FALSE") == "TRUE" @@ -467,8 +453,8 @@ def generate_qkv( @pytest.mark.skipif( - not is_fa3_supported(), - reason="flash_attn at sgl-kernel is only supported on sm90 or sm80", + not _is_fa3_supported(), + reason="flash_attn at sgl-kernel is only supported on CUDA sm90, sm80 or MUSA >= mp31", ) # @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float8_e4m3fn]) @pytest.mark.parametrize( @@ -1039,8 +1025,8 @@ def _generate_block_kvcache( @pytest.mark.skipif( - not is_fa3_supported(), - reason="flash_attn at sgl-kernel is only supported on sm90 or sm80", + not _is_fa3_supported(), + reason="flash_attn at sgl-kernel is only supported on CUDA sm90, sm80 or MUSA >= mp31", ) # @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float8_e4m3fn]) @pytest.mark.parametrize(