[MUSA][Diffusion] Fix fa3 API on MT MUSA (#23646)

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:
+    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
-    return (torch.version.cuda >= "12.3") and (
-        torch.cuda.get_device_capability(device)[0] == 9
-        or torch.cuda.get_device_capability(device)[0] == 8
-    )
 
 
 @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(
+    return _load_fa3_kernels()["flash_attn_varlen_func"](
-        _load_fa3_kernels()["flash_attn_varlen_func"],
+        q=q,
-        q,
+        k=k,
-        k,
+        v=v,
-        v,
+        cu_seqlens_q=cu_seqlens_q,
-        cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
-        cu_seqlens_k,
+        max_seqlen_q=max_seqlen_q,
-        max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
-        max_seqlen_k,
+        seqused_q=seqused_q,
-        seqused_q,
+        seqused_k=seqused_k,
-        seqused_k,
+        page_table=page_table,
-        page_table,
+        softmax_scale=softmax_scale,
-        softmax_scale,
+        causal=causal,
-        causal,
+        qv=qv,
-        qv,
+        q_descale=q_descale,
-        q_descale,
+        k_descale=k_descale,
-        k_descale,
+        v_descale=v_descale,
-        v_descale,
+        window_size=window_size,
-        window_size,
+        attention_chunk=attention_chunk,
-        attention_chunk,
+        softcap=softcap,
-        softcap,
+        num_splits=num_splits,
-        num_splits,
+        pack_gqa=pack_gqa,
-        pack_gqa,
+        sm_margin=sm_margin,
-        sm_margin,
+        return_softmax_lse=return_softmax_lse,
-        return_softmax_lse,
+        sinks=sinks,
-        sinks,
         out=out,
     )

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(),
+    not _is_fa3_supported(),
-    reason="flash_attn at sgl-kernel is only supported on sm90 or sm80",
+    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(),
+    not _is_fa3_supported(),
-    reason="flash_attn at sgl-kernel is only supported on sm90 or sm80",
+    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(