bump fa4 to official released fa4 pkg (#20303)

docker/Dockerfile

@@ -321,22 +321,21 @@ RUN --mount=type=cache,target=/root/.cache/pip \
     && rm -rf /root/.cargo /root/.rustup target dist ~/.cargo \
     && sed -i '/\.cargo\/env/d' /root/.profile /root/.bashrc 2>/dev/null || true
 
-# Pin nvidia-cutlass-dsl to 4.3.5
 RUN --mount=type=cache,target=/root/.cache/pip \
-   python3 -m pip uninstall nvidia-cutlass-dsl-libs-base; \
-   python3 -m pip install nvidia-cutlass-dsl==4.3.5 --force-reinstall;
+   python3 -m pip install "nvidia-cutlass-dsl>=4.4.1" "nvidia-cutlass-dsl-libs-base>=4.4.1" --force-reinstall --no-deps;
 
 # Patching packages for CUDA 12/13 compatibility
 # TODO: Remove when torch version covers these packages
 RUN --mount=type=cache,target=/root/.cache/pip if [ "${CUDA_VERSION%%.*}" = "12" ]; then \
     python3 -m pip install nvidia-nccl-cu12==2.28.3 --force-reinstall --no-deps ; \
     python3 -m pip install nvidia-cudnn-cu12==9.16.0.29 --force-reinstall --no-deps ; \
+    python3 -m pip install cuda-python==12.9 ; \
 elif [ "${CUDA_VERSION%%.*}" = "13" ]; then \
     python3 -m pip install nvidia-nccl-cu13==2.28.3 --force-reinstall --no-deps ; \
     python3 -m pip install nvidia-cudnn-cu13==9.16.0.29 --force-reinstall --no-deps ; \
     python3 -m pip install nvidia-cublas==13.1.0.3 --force-reinstall --no-deps ; \
     python3 -m pip install nixl-cu13 --no-deps ; \
-    python3 -m pip install cuda-python==13.1.1 ; \
+    python3 -m pip install cuda-python==13.2.0 ; \
 fi
 
 # Install development tools

python/pyproject.toml

@@ -38,7 +38,7 @@ dependencies = [
   "msgspec",
   "ninja",
   "numpy",
-  "nvidia-cutlass-dsl==4.3.5",
+  "nvidia-cutlass-dsl>=4.4.1",
   "nvidia-ml-py",
   "openai-harmony==0.0.4",
   "openai==2.6.1",
@@ -54,12 +54,12 @@ dependencies = [
   "pydantic",
   "python-multipart",
   "pyzmq>=25.1.2",
-  "quack-kernels==0.2.4",
+  "quack-kernels>=0.3.0",
   "requests",
   "scipy",
   "sentencepiece",
   "setproctitle",
-  "sgl-fa4==4.0.5",
+  "flash-attn-4>=4.0.0b4",
   "sglang-kernel==0.4.0",
   "soundfile==0.13.1",
   "tiktoken",
@@ -77,7 +77,6 @@ dependencies = [
   "uvloop",
   "watchfiles",
   "xgrammar==0.1.27",
-
   "smg-grpc-servicer>=0.5.0",
 ]
 

python/sglang/jit_kernel/flash_attention_v4.py

@@ -5,7 +5,7 @@ from typing import Callable, Optional, Tuple, Union
 import torch
 
 try:
-    from sgl_fa4.cute import flash_attn_varlen_func as _flash_attn_varlen_func
+    from flash_attn.cute import flash_attn_varlen_func as _flash_attn_varlen_func
 except Exception as _e:  # pragma: no cover
     _flash_attn_varlen_func = None
     _flash_attn_import_error = _e
@@ -44,7 +44,7 @@ def flash_attn_varlen_func(
     if _flash_attn_varlen_func is None:  # pragma: no cover
         raise ImportError(
             "Vendored FlashAttention CUTE is not available (cannot import "
-            "sgl_fa4.cute). Please check your source tree."
+            "flash_attn.cute). Please check your source tree."
         ) from _flash_attn_import_error
 
     q, k, v = [_maybe_contiguous(t) for t in (q, k, v)]

scripts/ci/cuda/ci_install_dependency.sh

@@ -117,7 +117,7 @@ else
 fi
 
 # Clean up existing installations
-$PIP_UNINSTALL_CMD sgl-kernel sglang-kernel sglang $PIP_UNINSTALL_SUFFIX || true
+$PIP_UNINSTALL_CMD sgl-kernel sglang-kernel sglang sgl-fa4 flash-attn-4 $PIP_UNINSTALL_SUFFIX || true
 
 # Keep flashinfer packages installed if version matches to avoid re-downloading:
 # - flashinfer-cubin: 150+ MB, plus extra cubins from ci_download_flashinfer_cubin.sh
@@ -178,17 +178,6 @@ fi
 # Install router for pd-disagg test
 $PIP_CMD install sglang-router $PIP_INSTALL_SUFFIX
 
-# Remove flash_attn folder to avoid conflicts
-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 '(?<=sglang-kernel==)[0-9A-Za-z\.\-]+' python/pyproject.toml)
@@ -329,9 +318,8 @@ fi
 # Download flashinfer cubins if the local set is incomplete
 bash "${SCRIPT_DIR}/ci_download_flashinfer_cubin.sh"
 
-# Clean nvidia-cutlass-dsl-libs-base for cutedsl lower than 0.4.4
-$PIP_UNINSTALL_CMD nvidia-cutlass-dsl-libs-base $PIP_UNINSTALL_SUFFIX || true
-$PIP_CMD install nvidia-cutlass-dsl==4.3.5 --force-reinstall $PIP_INSTALL_SUFFIX || true
+# Force reinstall nvidia-cutlass-dsl to avoid potential version mismatch issues
+$PIP_CMD install "nvidia-cutlass-dsl>=4.4.1" "nvidia-cutlass-dsl-libs-base>=4.4.1" --no-deps --force-reinstall $PIP_INSTALL_SUFFIX || true
 
 # Show current packages
 $PIP_CMD list