diff --git a/docker/Dockerfile b/docker/Dockerfile index 09953031a..2c5ec107a 100644 --- a/docker/Dockerfile +++ b/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 diff --git a/python/pyproject.toml b/python/pyproject.toml index 3abeddc08..17d4bcfbf 100755 --- a/python/pyproject.toml +++ b/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", ] diff --git a/python/sglang/jit_kernel/flash_attention_v4.py b/python/sglang/jit_kernel/flash_attention_v4.py index 8958ae1c0..dc5e7e4cf 100644 --- a/python/sglang/jit_kernel/flash_attention_v4.py +++ b/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)] diff --git a/scripts/ci/cuda/ci_install_dependency.sh b/scripts/ci/cuda/ci_install_dependency.sh index c4690bc47..9411890bf 100755 --- a/scripts/ci/cuda/ci_install_dependency.sh +++ b/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