diff --git a/scripts/ci/ci_install_dependency.sh b/scripts/ci/ci_install_dependency.sh index 7d0735602..09fdb4104 100755 --- a/scripts/ci/ci_install_dependency.sh +++ b/scripts/ci/ci_install_dependency.sh @@ -106,6 +106,16 @@ $PIP_CMD install -e "python[${EXTRAS}]" --extra-index-url https://download.pytor # Install router for pd-disagg test $PIP_CMD install sglang-router $PIP_INSTALL_SUFFIX +PYTHON_LIB_PATH=$(python3 -c "import site; print(site.getsitepackages()[0])") +FLASH_ATTN_PATH="${PYTHON_LIB_PATH}/flash_attn" + +if [ -d "$FLASH_ATTN_PATH" ]; then + echo "Directory $FLASH_ATTN_PATH exists. Removing..." + rm -rf "$FLASH_ATTN_PATH" +else + echo "Directory $FLASH_ATTN_PATH does not exist." +fi + # Install sgl-kernel SGL_KERNEL_VERSION_FROM_KERNEL=$(grep -Po '(?<=^version = ")[^"]*' sgl-kernel/pyproject.toml) SGL_KERNEL_VERSION_FROM_SRT=$(grep -Po -m1 '(?<=sgl-kernel==)[0-9A-Za-z\.\-]+' python/pyproject.toml) @@ -147,3 +157,14 @@ python3 -c "import torch; print(torch.version.cuda)" # Prepare the CI runner (cleanup HuggingFace cache, etc.) bash "${SCRIPT_DIR}/prepare_runner.sh" + +PYTHON_LIB_PATH=$(python3 -c "import site; print(site.getsitepackages()[0])") +FLASH_ATTN_PATH="${PYTHON_LIB_PATH}/flash_attn" + +if [ -d "$FLASH_ATTN_PATH" ]; then + echo "Directory $FLASH_ATTN_PATH exists. Removing..." + rm -rf "$FLASH_ATTN_PATH" + echo "error: this should not happen" +else + echo "Directory $FLASH_ATTN_PATH does not exist." +fi diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt index 079eb2f45..9a438591a 100644 --- a/sgl-kernel/CMakeLists.txt +++ b/sgl-kernel/CMakeLists.txt @@ -593,9 +593,40 @@ install(DIRECTORY "${repo-triton_SOURCE_DIR}/python/triton_kernels/triton_kernel # ============================ Extra Install: FA4 ============================= # # TODO: find a better install condition. if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.8" OR SGL_KERNEL_ENABLE_SM100A) - # flash_attn/cute - install(DIRECTORY "${repo-flash-attention_SOURCE_DIR}/flash_attn/cute/" - DESTINATION "flash_attn/cute" + + 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() + +endif() diff --git a/sgl-kernel/python/sgl_kernel/_fa4_interface.py b/sgl-kernel/python/sgl_kernel/_fa4_interface.py index 5411cfe16..ffd0563f5 100644 --- a/sgl-kernel/python/sgl_kernel/_fa4_interface.py +++ b/sgl-kernel/python/sgl_kernel/_fa4_interface.py @@ -19,9 +19,9 @@ import cutlass import cutlass.cute as cute import torch from cutlass.cute.runtime import from_dlpack -from flash_attn.cute import utils -from flash_attn.cute.flash_fwd import FlashAttentionForwardSm90 -from flash_attn.cute.flash_fwd_sm100 import FlashAttentionForwardSm100 +from flash_attn_origin.cute import utils +from flash_attn_origin.cute.flash_fwd import FlashAttentionForwardSm90 +from flash_attn_origin.cute.flash_fwd_sm100 import FlashAttentionForwardSm100 def maybe_contiguous(x):