fa4 cleanup (#19727)

python/sglang/srt/layers/attention/nsa_backend.py

@@ -42,8 +42,6 @@ from sglang.srt.layers.dp_attention import get_attention_tp_size
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.utils import is_cuda, is_hip
 
-# from sgl_kernel.flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache
-
 if TYPE_CHECKING:
     from sglang.srt.layers.radix_attention import RadixAttention
     from sglang.srt.model_executor.model_runner import ModelRunner
@@ -1783,8 +1781,6 @@ class NativeSparseAttnBackend(
             )
 
         # Use FA3 for SM90 (Hopper/H200)
-        fa_version = 3
-
         return flash_attn_varlen_func(
             q=q,
             k=k,
@@ -1795,7 +1791,6 @@ class NativeSparseAttnBackend(
             max_seqlen_k=max_seqlen_k,
             softmax_scale=layer.scaling,
             causal=causal,
-            ver=fa_version,
         )
 
     def _forward_tilelang(

python/sglang/srt/layers/attention/vision.py

@@ -35,7 +35,22 @@ _is_hip = is_hip()
 
 if _is_cuda:
     from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache
-    from sgl_kernel.flash_attn import flash_attn_varlen_func
+
+    try:
+        from sgl_kernel.flash_attn import flash_attn_varlen_func
+
+        from sglang.jit_kernel.flash_attention_v4 import (
+            flash_attn_varlen_func as flash_attn_varlen_func_fa4,
+        )
+
+        def flash_attn_func(*args, ver: int = 3, **kwargs):
+            if ver == 4:
+                return flash_attn_varlen_func_fa4(*args, **kwargs)
+            return flash_attn_varlen_func(*args, **kwargs)
+
+    except ImportError as e:
+        raise e
+
 
 if _is_npu:
     import torch_npu
@@ -391,7 +406,7 @@ class VisionFlash3Attention(nn.Module):
         """
         if envs.SGLANG_VIT_ENABLE_CUDA_GRAPH.get():
             max_seqlen = cu_seqlens[1]
-            output = flash_attn_varlen_func(
+            output = flash_attn_func(
                 q,
                 k,
                 v,
@@ -406,7 +421,7 @@ class VisionFlash3Attention(nn.Module):
             seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
             max_seqlen = seq_lens.max().item()
 
-            output = flash_attn_varlen_func(
+            output = flash_attn_func(
                 q,
                 k,
                 v,
@@ -457,7 +472,7 @@ class VisionFlash4Attention(nn.Module):
         seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
         max_seqlen = seq_lens.max().item()
 
-        output = flash_attn_varlen_func(
+        output = flash_attn_func(
             q,
             k,
             v,