fa4 cleanup (#19727)

sgl-kernel/CMakeLists.txt

@@ -570,44 +570,3 @@ install(DIRECTORY "${repo-triton_SOURCE_DIR}/python/triton_kernels/triton_kernel
         DESTINATION "triton_kernels"
         PATTERN ".git*" EXCLUDE
         PATTERN "__pycache__" EXCLUDE)
-
-# ============================ Extra Install: FA4 ============================= #
-# TODO: find a better install condition.
-if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.8" OR SGL_KERNEL_ENABLE_SM100A)
-
-    set(FLASH_ATTN_CUTE_SRC "${repo-flash-attention_SOURCE_DIR}/flash_attn/cute")
-    set(FLASH_ATTN_CUTE_DST "${CMAKE_CURRENT_BINARY_DIR}/flash_attn_origin/cute")
-
-    file(MAKE_DIRECTORY "${FLASH_ATTN_CUTE_DST}")
-
-    file(COPY "${FLASH_ATTN_CUTE_SRC}/"
-         DESTINATION "${FLASH_ATTN_CUTE_DST}"
-         PATTERN ".git*" EXCLUDE
-         PATTERN "__pycache__" EXCLUDE)
-
-    file(GLOB_RECURSE FLASH_ATTN_CUTE_DST_PY
-         "${FLASH_ATTN_CUTE_DST}/*.py")
-
-    foreach(FILE_PATH IN LISTS FLASH_ATTN_CUTE_DST_PY)
-        file(READ "${FILE_PATH}" FILE_CONTENT)
-
-        set(MODIFIED_CONTENT "${FILE_CONTENT}")
-
-        # The main goal is to avoid using "flash_attn" so that other libraries (such as transformers) do not mistakenly assume that "flash_attn" is already installed.
-
-        string(REPLACE "flash_attn.cute"
-                       "flash_attn_origin.cute"
-                       MODIFIED_CONTENT "${MODIFIED_CONTENT}")
-
-        if (NOT FILE_CONTENT STREQUAL MODIFIED_CONTENT)
-            file(WRITE "${FILE_PATH}" "${MODIFIED_CONTENT}")
-            message(STATUS "  - [FA4 Patch] Patched: ${FILE_PATH}")
-        endif()
-    endforeach()
-
-    install(DIRECTORY "${FLASH_ATTN_CUTE_DST}/"
-            DESTINATION "flash_attn_origin/cute"
-            PATTERN ".git*" EXCLUDE
-            PATTERN "__pycache__" EXCLUDE)
-
-endif()

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(