Improve CI by trying a warmup before unit tests (#14669)

.github/workflows/pr-test.yml

@@ -457,11 +457,6 @@ jobs:
       - name: Checkout code
         uses: actions/checkout@v4
 
-      - name: Cleanup
-        run: |
-          ls -alh sgl-kernel/dist || true
-          rm -rf sgl-kernel/dist/* || true
-
       - name: Download artifacts
         if: needs.check-changes.outputs.sgl_kernel == 'true'
         uses: actions/download-artifact@v4
@@ -504,11 +499,6 @@ jobs:
       - name: Checkout code
         uses: actions/checkout@v4
 
-      - name: Cleanup
-        run: |
-          ls -alh sgl-kernel/dist || true
-          rm -rf sgl-kernel/dist/* || true
-
       - name: Download artifacts
         if: needs.check-changes.outputs.sgl_kernel == 'true'
         uses: actions/download-artifact@v4
@@ -557,8 +547,8 @@ jobs:
 
         - name: Install dependencies
           run: |
-            pip install "bitsandbytes>=0.44.0"
             CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_dependency.sh
+
         - name: Run test
           timeout-minutes: 30
           run: |
@@ -724,6 +714,12 @@ jobs:
         run: |
           CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_dependency.sh
 
+      - name: Warmup Weights and JIT Compilation
+        timeout-minutes: 20
+        run: |
+          # An example command for testing the warmup. TODO: make this more general and move them to python scripts.
+          python3 -m sglang.compile_deep_gemm --model deepseek-ai/DeepSeek-V3-0324 --tp 8 --trust-remote-code
+
       - name: Run test
         timeout-minutes: 20
         run: |

python/pyproject.toml

@@ -104,6 +104,7 @@ vsa = ["torch", "setuptools"]
 
 test = [
   "accelerate",
+  "bitsandbytes",
   "expecttest",
   "jsonlines",
   "matplotlib",

python/sglang/srt/entrypoints/engine.py

@@ -788,8 +788,8 @@ def _launch_subprocesses(
     """
     # Configure global environment
     configure_logger(server_args)
-    server_args.check_server_args()
     _set_envs_and_config(server_args)
+    server_args.check_server_args()
 
     # Allocate ports for inter-process communications
     if port_args is None:

python/sglang/srt/server_args.py

@@ -4171,6 +4171,9 @@ class ServerArgs:
             assert is_npu(), "MindSpore model impl is only supported on Ascend npu."
 
     def check_torch_2_9_1_cudnn_compatibility(self):
+        if get_bool_env_var("SGLANG_DISABLE_CUDNN_CHECK"):
+            return
+
         if self.get_model_config().is_multimodal:
             import torch
 
@@ -4198,7 +4201,8 @@ class ServerArgs:
                             "Reference: https://github.com/pytorch/pytorch/issues/168167\n\n"
                             "Solution:  You MUST upgrade CuDNN to version 9.15+ to ensure correctness.\n\n"
                             "Run the following command immediately to fix:\n"
-                            "    pip install nvidia-cudnn-cu12==9.16.0.29\n"
+                            "    pip install nvidia-cudnn-cu12==9.16.0.29\n\n"
+                            "Or you can disable this check by setting env var SGLANG_DISABLE_CUDNN_CHECK=1\n"
                             "--------------------------------------------------------------------------------\n"
                             f"{RESET}"
                         )

scripts/ci/ci_install_dependency.sh

@@ -2,6 +2,7 @@
 # Install the dependency in CI.
 set -euxo pipefail
 
+# Set up environment variables
 IS_BLACKWELL=${IS_BLACKWELL:-0}
 CU_VERSION="cu129"
 OPTIONAL_DEPS="${1:-}"
@@ -70,37 +71,37 @@ else
 fi
 
 # Install uv
+pip install --upgrade pip
+
 if [ "$IS_BLACKWELL" = "1" ]; then
     # The blackwell CI runner has some issues with pip and uv,
     # so we can only use pip with `--break-system-packages`
     PIP_CMD="pip"
     PIP_INSTALL_SUFFIX="--break-system-packages"
-    $PIP_CMD install --upgrade pip
+    PIP_UNINSTALL_CMD="pip uninstall -y"
-
+    PIP_UNINSTALL_SUFFIX="--break-system-packages"
-    # Clean up existing installations
-    $PIP_CMD uninstall -y sgl-kernel sglang $PIP_INSTALL_SUFFIX || true
-    $PIP_CMD uninstall -y flashinfer-python flashinfer-cubin flashinfer-jit-cache $PIP_INSTALL_SUFFIX || true
 else
     # In normal cases, we use uv, which is much faster than pip.
-    pip install --upgrade pip
     pip install uv
     export UV_SYSTEM_PYTHON=true
 
     PIP_CMD="uv pip"
     PIP_INSTALL_SUFFIX="--index-strategy unsafe-best-match --prerelease allow"
-
+    PIP_UNINSTALL_CMD="uv pip uninstall"
-    # Clean up existing installations
+    PIP_UNINSTALL_SUFFIX=""
-    $PIP_CMD uninstall sgl-kernel sglang || true
-    $PIP_CMD uninstall flashinfer-python flashinfer-cubin flashinfer-jit-cache || true
 fi
 
+# Clean up existing installations
+$PIP_UNINSTALL_CMD sgl-kernel sglang $PIP_UNINSTALL_SUFFIX || true
+$PIP_UNINSTALL_CMD flashinfer-python flashinfer-cubin flashinfer-jit-cache $PIP_UNINSTALL_SUFFIX || true
+
+# Install the main package
 EXTRAS="dev"
 if [ -n "$OPTIONAL_DEPS" ]; then
     EXTRAS="dev,${OPTIONAL_DEPS}"
 fi
 echo "Installing python extras: [${EXTRAS}]"
 
-# Install the main package
 $PIP_CMD install -e "python[${EXTRAS}]" --extra-index-url https://download.pytorch.org/whl/${CU_VERSION} $PIP_INSTALL_SUFFIX
 
 # Install router for pd-disagg test
@@ -137,6 +138,7 @@ fi
 # Show current packages
 $PIP_CMD list
 
+# Install other python dependencies
 $PIP_CMD install mooncake-transfer-engine==0.3.7.post2 "${NVRTC_SPEC}" py-spy scipy huggingface_hub[hf_xet] pytest $PIP_INSTALL_SUFFIX
 
 if [ "$IS_BLACKWELL" != "1" ]; then
@@ -151,6 +153,7 @@ $PIP_CMD install nvidia-nvshmem-cu12==3.4.5 --force-reinstall $PIP_INSTALL_SUFFI
 # Cudnn with version less than 9.16.0.29 will cause performance regression on Conv3D kernel
 $PIP_CMD install nvidia-cudnn-cu12==9.16.0.29 --force-reinstall $PIP_INSTALL_SUFFIX
 $PIP_CMD uninstall xformers || true
+
 # Show current packages
 $PIP_CMD list
 python3 -c "import torch; print(torch.version.cuda)"
@@ -158,6 +161,7 @@ python3 -c "import torch; print(torch.version.cuda)"
 # Prepare the CI runner (cleanup HuggingFace cache, etc.)
 bash "${SCRIPT_DIR}/prepare_runner.sh"
 
+# Remove flash_attn folder to avoid conflicts with sgl-kernel
 PYTHON_LIB_PATH=$(python3 -c "import site; print(site.getsitepackages()[0])")
 FLASH_ATTN_PATH="${PYTHON_LIB_PATH}/flash_attn"