fa4 cleanup (#19727)

sgl-kernel/python/sgl_kernel/flash_attn.py

@@ -10,11 +10,6 @@ except:
         "Can not import FA3 in sgl_kernel. Please check your installation."
     )
 
-try:
-    from ._fa4_interface import flash_attn_varlen_func as flash_attn_varlen_func_v4
-except ImportError:
-    flash_attn_varlen_func_v4 = None
-
 
 @lru_cache(maxsize=1)
 def is_fa3_supported(device=None) -> bool:
@@ -160,45 +155,6 @@ def flash_attn_with_kvcache(
             logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
             normalization factor).
     """
-    if ver == 4:
-        assert (
-            flash_attn_varlen_func_v4 is not None
-        ), "FA4 is not available, please check your installation."
-        # Using `(-1, -1)` as no sliding window causes correctness issues for FA4.
-        assert (
-            k is None and v is None
-        ), "FA4 does not support updating KV cache in-place."
-        assert (
-            rotary_cos is None and rotary_sin is None and rotary_seqlens is None
-        ), "FA4 does not support rotary embedding."
-        assert (
-            cache_batch_idx is None and cache_leftpad is None
-        ), "FA4 does not support non-consecutive batch indices or left padding."
-        assert (
-            q_descale is None and k_descale is None and v_descale is None
-        ), "FA4 does not support descale."
-
-        if window_size == (-1, -1):
-            window_size = (None, None)
-
-        return flash_attn_varlen_func_v4(
-            q=q,
-            k=k_cache,
-            v=v_cache,
-            cu_seqlens_q=cu_seqlens_q,
-            seqused_k=cache_seqlens,
-            softmax_scale=softmax_scale,
-            causal=causal,
-            window_size=window_size,
-            softcap=softcap,
-            num_splits=num_splits,
-            pack_gqa=pack_gqa,
-            return_softmax_lse=return_softmax_lse,
-            learnable_sink=sinks,
-            page_table=page_table,
-            score_mod=score_mod,
-            aux_tensors=aux_tensors,
-        )
 
     assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension"
     assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension"
@@ -298,32 +254,6 @@ def flash_attn_varlen_func(
     aux_tensors=None,
     ver=3,
 ):
-    if ver == 4:
-        assert (
-            flash_attn_varlen_func_v4 is not None
-        ), "FA4 is not available, please check your installation."
-        # Using `(-1, -1)` as no sliding window causes correctness issues for FA4.
-        if window_size == (-1, -1):
-            window_size = (None, None)
-        return flash_attn_varlen_func_v4(
-            q,
-            k,
-            v,
-            cu_seqlens_q=cu_seqlens_q,
-            cu_seqlens_k=cu_seqlens_k,
-            seqused_q=seqused_q,
-            seqused_k=seqused_k,
-            page_table=page_table,
-            softmax_scale=softmax_scale,
-            causal=causal,
-            window_size=window_size,
-            softcap=softcap,
-            pack_gqa=pack_gqa,
-            learnable_sink=sinks,
-            return_softmax_lse=return_softmax_lse,
-            score_mod=score_mod,
-            aux_tensors=aux_tensors,
-        )
 
     if not is_fa3_supported():
         raise NotImplementedError(