minor: support flashinfer nightly (#2295)

.github/workflows/pr-test.yml

@@ -12,6 +12,15 @@ on:
       - "python/sglang/**"
       - "test/**"
   workflow_dispatch:
+    inputs:
+      version:
+        description: "FlashInfer version"
+        required: true
+        type: choice
+        default: 'release'
+        options:
+          - 'release'
+          - 'nightly'
 
 concurrency:
   group: pr-test-${{ github.ref }}
@@ -26,6 +35,8 @@ jobs:
         uses: actions/checkout@v3
 
       - name: Install dependencies
+        env:
+          FLASHINFER_REPO: ${{ inputs.version == 'nightly' && 'https://flashinfer.ai/whl/nightly/cu121/torch2.4/' || 'https://flashinfer.ai/whl/cu121/torch2.4/' }}
         run: |
           bash scripts/ci_install_dependency.sh
 
@@ -46,6 +57,8 @@ jobs:
         uses: actions/checkout@v3
 
       - name: Install dependencies
+        env:
+          FLASHINFER_REPO: ${{ inputs.version == 'nightly' && 'https://flashinfer.ai/whl/nightly/cu121/torch2.4/' || 'https://flashinfer.ai/whl/cu121/torch2.4/' }}
         run: |
           bash scripts/ci_install_dependency.sh
 
@@ -66,6 +79,8 @@ jobs:
         uses: actions/checkout@v3
 
       - name: Install dependencies
+        env:
+          FLASHINFER_REPO: ${{ inputs.version == 'nightly' && 'https://flashinfer.ai/whl/nightly/cu121/torch2.4/' || 'https://flashinfer.ai/whl/cu121/torch2.4/' }}
         run: |
           bash scripts/ci_install_dependency.sh
 
@@ -92,6 +107,8 @@ jobs:
         uses: actions/checkout@v3
 
       - name: Install dependencies
+        env:
+          FLASHINFER_REPO: ${{ inputs.version == 'nightly' && 'https://flashinfer.ai/whl/nightly/cu121/torch2.4/' || 'https://flashinfer.ai/whl/cu121/torch2.4/' }}
         run: |
           bash scripts/ci_install_dependency.sh
 
@@ -127,6 +144,8 @@ jobs:
         uses: actions/checkout@v3
 
       - name: Install dependencies
+        env:
+          FLASHINFER_REPO: ${{ inputs.version == 'nightly' && 'https://flashinfer.ai/whl/nightly/cu121/torch2.4/' || 'https://flashinfer.ai/whl/cu121/torch2.4/' }}
         run: |
           bash scripts/ci_install_dependency.sh
 
@@ -156,6 +175,8 @@ jobs:
         uses: actions/checkout@v3
 
       - name: Install dependencies
+        env:
+          FLASHINFER_REPO: ${{ inputs.version == 'nightly' && 'https://flashinfer.ai/whl/nightly/cu121/torch2.4/' || 'https://flashinfer.ai/whl/cu121/torch2.4/' }}
         run: |
           bash scripts/ci_install_dependency.sh
 
@@ -185,6 +206,8 @@ jobs:
         uses: actions/checkout@v3
 
       - name: Install dependencies
+        env:
+          FLASHINFER_REPO: ${{ inputs.version == 'nightly' && 'https://flashinfer.ai/whl/nightly/cu121/torch2.4/' || 'https://flashinfer.ai/whl/cu121/torch2.4/' }}
         run: |
           bash scripts/ci_install_dependency.sh
 
@@ -207,6 +230,8 @@ jobs:
         uses: actions/checkout@v3
 
       - name: Install dependencies
+        env:
+          FLASHINFER_REPO: ${{ inputs.version == 'nightly' && 'https://flashinfer.ai/whl/nightly/cu121/torch2.4/' || 'https://flashinfer.ai/whl/cu121/torch2.4/' }}
         run: |
           bash scripts/ci_install_dependency.sh
 

python/sglang/check_env.py

@@ -9,7 +9,7 @@ from collections import OrderedDict, defaultdict
 
 import torch
 
-# List of packages to check versions for
+# List of packages to check versions
 PACKAGE_LIST = [
     "sglang",
     "flashinfer",

scripts/ci_install_dependency.sh

@@ -1,14 +1,16 @@
-"""
-Install the dependency in CI.
-"""
+# Install the dependency in CI.
 
-./killall_sglang.sh
+# Use repo from environment variable, passed from GitHub Actions
+FLASHINFER_REPO="${FLASHINFER_REPO:-https://flashinfer.ai/whl/cu121/torch2.4}"
+
+SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
+bash "${SCRIPT_DIR}/killall_sglang.sh"
 
 pip install --upgrade pip
 pip install -e "python[all]" --find-links https://flashinfer.ai/whl/cu121/torch2.4/flashinfer/
 
 # Force reinstall flashinfer
-pip install flashinfer -i https://flashinfer.ai/whl/cu121/torch2.4/ --force-reinstall
+pip install flashinfer -i ${FLASHINFER_REPO} --force-reinstall
 
 pip install transformers==4.45.2 sentence_transformers accelerate peft