gahow/agentic-pd-hybrid
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: vendor sglang v0.5.10 snapshot

Browse Source
This commit is contained in:
Gahow Wang
2026-04-24 12:29:36 +00:00
parent 78f0d15221
commit bded08301f
4308 changed files with 1200894 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
third_party/sglang/scripts/ci/cuda/cache_nvidia_wheels.sh vendored Executable file
View File

@@ -0,0 +1,24 @@
#!/bin/bash
# Cache and pre-install nvidia wheels that torch pins.
#
# pypi.nvidia.com returns Cache-Control: no-store, so pip re-downloads
# cudnn (~707 MB) and nvshmem (~125 MB) on every CI run. This script
# caches the wheels locally and installs them so that the subsequent
# `pip install -e "python[dev]"` sees "Requirement already satisfied".
#
# Integrity: uses `unzip -t` to detect partial/corrupt downloads.
#
# Usage: source scripts/ci/cuda/cache_nvidia_wheels.sh
NVIDIA_WHEEL_CACHE="/root/.cache/nvidia-wheels"
mkdir -p "$NVIDIA_WHEEL_CACHE"
for url in \
"https://pypi.nvidia.com/nvidia-cudnn-cu12/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl" \
"https://pypi.nvidia.com/nvidia-nvshmem-cu12/nvidia_nvshmem_cu12-3.3.20-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl"; do
whl="$NVIDIA_WHEEL_CACHE/$(basename "$url")"
[ -f "$whl" ] && unzip -tq "$whl" &>/dev/null || curl -fL -o "$whl" "$url"
done
pip install --no-deps "$NVIDIA_WHEEL_CACHE"/nvidia_cudnn_cu12-*.whl \
"$NVIDIA_WHEEL_CACHE"/nvidia_nvshmem_cu12-*.whl 2>/dev/null || true

62
third_party/sglang/scripts/ci/cuda/ci_download_flashinfer_cubin.sh vendored Executable file
View File

@@ -0,0 +1,62 @@
#!/bin/bash
# Download flashinfer cubins if the local set is incomplete.
#
# The flashinfer-cubin pip package may not include cubins for newer architectures
# (e.g. sm_100, sm_120) due to PyPI size limits. This script checks the local
# cubin status against the flashinfer artifact repository and downloads any
# missing files.
#
# This script is best-effort: if the status check or download times out (e.g.
# due to a GPU in error state blocking CUDA init), we warn and continue.
# The pip package already includes cubins for common architectures (sm_80, sm_90).
set -uxo pipefail
# Early exit: the pip package already includes cubins for sm_80 and sm_90.
# Only sm_100+ (Blackwell) needs extra cubins downloaded. Skip the expensive
# Python status check entirely if no such GPU is present.
if COMPUTE_CAPS=$(timeout 10 nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null); then
NEEDS_EXTRA_CUBINS=false
while IFS= read -r cap; do
major="${cap%%.*}"
if [ "$major" -ge 10 ] 2>/dev/null; then
NEEDS_EXTRA_CUBINS=true
break
fi
done <<< "$COMPUTE_CAPS"
if [ "$NEEDS_EXTRA_CUBINS" = false ]; then
echo "All GPUs are sm_9x or older (compute caps: $(echo $COMPUTE_CAPS | tr '\n' ' ')), pip cubins sufficient — skipping download"
exit 0
fi
fi
# Use timeout to prevent hangs when GPUs are in error state (the flashinfer
# import can trigger CUDA init which blocks on bad GPUs).
CUBIN_STATUS=$(timeout 60 python3 -c "
import os
os.environ.setdefault('CUDA_VISIBLE_DEVICES', '')
from flashinfer.artifacts import get_artifacts_status
status = get_artifacts_status()
total = len(status)
downloaded = sum(1 for _, exists in status if exists)
print(f'{downloaded}/{total}')
" 2>/dev/null) || CUBIN_STATUS="unknown"
echo "Flashinfer cubin status: ${CUBIN_STATUS}"
if echo "$CUBIN_STATUS" | grep -qE '^[0-9]+/[0-9]+$'; then
CUBIN_DOWNLOADED="${CUBIN_STATUS%/*}"
CUBIN_TOTAL="${CUBIN_STATUS#*/}"
if [ "$CUBIN_DOWNLOADED" = "$CUBIN_TOTAL" ] && [ "$CUBIN_TOTAL" != "0" ]; then
echo "All flashinfer cubins already present (${CUBIN_STATUS}), skipping download"
else
echo "Cubins incomplete (${CUBIN_STATUS}), downloading..."
if ! timeout 300 env FLASHINFER_LOGGING_LEVEL=warning python3 -m flashinfer --download-cubin; then
echo "WARNING: flashinfer cubin download failed or timed out, continuing with existing cubins"
fi
fi
else
echo "Could not determine cubin status (status check timed out or failed), attempting download..."
if ! timeout 300 env FLASHINFER_LOGGING_LEVEL=warning python3 -m flashinfer --download-cubin; then
echo "WARNING: flashinfer cubin download failed or timed out, continuing with existing cubins"
fi
fi

69
third_party/sglang/scripts/ci/cuda/ci_download_flashinfer_jit_cache.sh vendored Executable file
View File

@@ -0,0 +1,69 @@
#!/bin/bash
# Install flashinfer-jit-cache with caching and retry logic (flashinfer.ai can have transient DNS issues).
# The jit-cache wheel is 1.2+ GB, so we skip the download entirely if already installed.
#
# Required environment (caller must export or set):
# UNINSTALL_JIT_CACHE — literal true/false (skip download when false)
# FLASHINFER_PYTHON_REQUIRED — e.g. from python/pyproject.toml (flashinfer_python)
# CU_VERSION — e.g. cu129
# PIP_CMD — e.g. "pip" or "uv pip"
# PIP_INSTALL_SUFFIX — extra pip args for this runner
set -euxo pipefail
: "${UNINSTALL_JIT_CACHE:?must be set}"
: "${FLASHINFER_PYTHON_REQUIRED:?must be set}"
: "${CU_VERSION:?must be set}"
: "${PIP_CMD:?must be set}"
FLASHINFER_JIT_CACHE_INSTALLED=false
if [ "$UNINSTALL_JIT_CACHE" = false ]; then
FLASHINFER_JIT_CACHE_INSTALLED=true
echo "flashinfer-jit-cache already at correct version, skipping download"
fi
if [ "$FLASHINFER_JIT_CACHE_INSTALLED" = false ]; then
FLASHINFER_CACHE_DIR="${HOME}/.cache/flashinfer-wheels"
mkdir -p "${FLASHINFER_CACHE_DIR}"
FLASHINFER_WHEEL_PATTERN="flashinfer_jit_cache-${FLASHINFER_PYTHON_REQUIRED}*.whl"
CACHED_WHEEL=$(find "${FLASHINFER_CACHE_DIR}" -name "${FLASHINFER_WHEEL_PATTERN}" -type f 2>/dev/null | head -n 1)
if [ -n "$CACHED_WHEEL" ] && [ -f "$CACHED_WHEEL" ]; then
echo "Found cached flashinfer wheel: $CACHED_WHEEL"
if $PIP_CMD install "$CACHED_WHEEL" $PIP_INSTALL_SUFFIX; then
FLASHINFER_JIT_CACHE_INSTALLED=true
echo "Successfully installed flashinfer-jit-cache from cache"
else
echo "Failed to install from cache, will try downloading..."
rm -f "$CACHED_WHEEL"
fi
fi
if [ "$FLASHINFER_JIT_CACHE_INSTALLED" = false ]; then
for i in {1..5}; do
# Download wheel to cache directory (use pip directly as uv pip doesn't support download)
if timeout 600 pip download "flashinfer-jit-cache==${FLASHINFER_PYTHON_REQUIRED}" \
--index-url "https://flashinfer.ai/whl/${CU_VERSION}" \
-d "${FLASHINFER_CACHE_DIR}"; then
CACHED_WHEEL=$(find "${FLASHINFER_CACHE_DIR}" -name "${FLASHINFER_WHEEL_PATTERN}" -type f 2>/dev/null | head -n 1)
if [ -n "$CACHED_WHEEL" ] && [ -f "$CACHED_WHEEL" ]; then
if $PIP_CMD install "$CACHED_WHEEL" $PIP_INSTALL_SUFFIX; then
FLASHINFER_JIT_CACHE_INSTALLED=true
echo "Successfully downloaded and installed flashinfer-jit-cache"
break
fi
else
echo "Warning: Download succeeded but wheel file not found"
fi
fi
echo "Attempt $i to download flashinfer-jit-cache failed, retrying in 10 seconds..."
sleep 10
done
fi
fi
if [ "$FLASHINFER_JIT_CACHE_INSTALLED" = false ]; then
echo "ERROR: Failed to install flashinfer-jit-cache after 5 attempts"
exit 1
fi

119
third_party/sglang/scripts/ci/cuda/ci_install_deepep.sh vendored Executable file
View File

@@ -0,0 +1,119 @@
#!/bin/bash
# Install the dependency in CI.
set -euxo pipefail
bash scripts/ci/cuda/ci_install_dependency.sh
export GDRCOPY_HOME=/usr/src/gdrdrv-2.5.1/
export CUDA_HOME=/usr/local/cuda
GRACE_BLACKWELL=${GRACE_BLACKWELL:-0}
# Detect architecture
ARCH=$(uname -m)
if [ "$ARCH" != "x86_64" ] && [ "$ARCH" != "aarch64" ]; then
echo "Unsupported architecture: $ARCH"
exit 1
fi
if python3 -c "import deep_ep" >/dev/null 2>&1; then
echo "deep_ep is already installed or importable. Skipping installation."
exit 0
fi
# Install system dependencies
# Use fallback logic in case apt fails due to unrelated broken packages on the runner
DEEPEP_SYSTEM_DEPS="curl wget git sudo rdma-core infiniband-diags openssh-server perftest libibumad3 libibverbs-dev libibverbs1 ibverbs-providers ibverbs-utils libnl-3-200 libnl-route-3-200 librdmacm1 build-essential cmake"
apt-get install -y --no-install-recommends $DEEPEP_SYSTEM_DEPS || {
echo "Warning: apt-get install failed, checking if required packages are available..."
for pkg in $DEEPEP_SYSTEM_DEPS; do
if ! dpkg -l "$pkg" 2>/dev/null | grep -q "^ii"; then
echo "ERROR: Required package $pkg is not installed and apt-get failed"
exit 1
fi
done
echo "All required packages are already installed, continuing..."
}
# Install GDRCopy
rm -rf /opt/gdrcopy && mkdir -p /opt/gdrcopy
cd /opt/gdrcopy
git clone https://github.com/NVIDIA/gdrcopy.git .
git checkout v2.5.1
apt-get update || true # May fail due to unrelated broken packages
GDRCOPY_DEPS_1="nvidia-dkms-580"
GDRCOPY_DEPS_2="build-essential devscripts debhelper fakeroot pkg-config dkms"
GDRCOPY_DEPS_3="check libsubunit0 libsubunit-dev python3-venv"
for deps_group in "$GDRCOPY_DEPS_1" "$GDRCOPY_DEPS_2" "$GDRCOPY_DEPS_3"; do
apt-get install -y --no-install-recommends $deps_group || {
echo "Warning: apt-get install failed for '$deps_group', checking if packages are available..."
for pkg in $deps_group; do
if ! dpkg -l "$pkg" 2>/dev/null | grep -q "^ii"; then
echo "ERROR: Required package $pkg is not installed and apt-get failed"
exit 1
fi
done
echo "All required packages from '$deps_group' are already installed, continuing..."
}
done
cd packages
CUDA=/usr/local/cuda ./build-deb-packages.sh
dpkg -i gdrdrv-dkms_*.deb
dpkg -i libgdrapi_*.deb
dpkg -i gdrcopy-tests_*.deb
dpkg -i gdrcopy_*.deb
# Set up library paths based on architecture
LIB_PATH="/usr/lib/$ARCH-linux-gnu"
if [ ! -e "$LIB_PATH/libmlx5.so" ]; then
ln -s $LIB_PATH/libmlx5.so.1 $LIB_PATH/libmlx5.so
fi
apt-get update || true
apt-get install -y --no-install-recommends libfabric-dev || {
if ! dpkg -l libfabric-dev 2>/dev/null | grep -q "^ii"; then
echo "ERROR: Required package libfabric-dev is not installed and apt-get failed"
exit 1
fi
echo "libfabric-dev is already installed, continuing..."
}
# Install DeepEP
DEEPEP_DIR=/root/.cache/deepep
rm -rf ${DEEPEP_DIR}
if [ "$GRACE_BLACKWELL" = "1" ]; then
# We use Tom's DeepEP fork for GB200 for now, which supports fp4 dispatch.
GRACE_BLACKWELL_DEEPEP_BRANCH=gb200_blog_part_2
git clone https://github.com/fzyzcjy/DeepEP.git ${DEEPEP_DIR} && \
pushd ${DEEPEP_DIR} && \
git checkout ${GRACE_BLACKWELL_DEEPEP_BRANCH} && \
sed -i 's/#define NUM_CPU_TIMEOUT_SECS 100/#define NUM_CPU_TIMEOUT_SECS 1000/' csrc/kernels/configs.cuh && \
popd
else
git clone https://github.com/deepseek-ai/DeepEP.git ${DEEPEP_DIR} && \
pushd ${DEEPEP_DIR} && \
git checkout 9af0e0d0e74f3577af1979c9b9e1ac2cad0104ee && \
popd
fi
cd ${DEEPEP_DIR}
if [ "$GRACE_BLACKWELL" = "1" ]; then
CUDA_VERSION=$(nvidia-smi | grep "CUDA Version" | head -n1 | awk '{print $9}')
if [ "$CUDA_VERSION" = "12.8" ]; then
CHOSEN_TORCH_CUDA_ARCH_LIST='10.0'
elif awk -v ver="$CUDA_VERSION" 'BEGIN {exit !(ver > 12.8)}'; then
# With cuda > 12.8, the compiler supports 10.3, so we should use
# CHOSEN_TORCH_CUDA_ARCH_LIST='10.0;10.3'
#
# However, our CI machine has a weird setup and nvidia-smi reports wrong CUDA version in the container.
# The container is actually cuda 12.8, but nvidia-smi reports 13.0, leading to compilation errors. so we
# drop 10.3.
CHOSEN_TORCH_CUDA_ARCH_LIST='10.0'
else
echo "Unsupported CUDA version for Grace Blackwell: $CUDA_VERSION" && exit 1
fi && \
if [ "${CUDA_VERSION%%.*}" = "13" ]; then \
sed -i "/^ include_dirs = \['csrc\/'\]/a\ include_dirs.append('${CUDA_HOME}/include/cccl')" setup.py; \
fi
TORCH_CUDA_ARCH_LIST="${CHOSEN_TORCH_CUDA_ARCH_LIST}" pip install --no-build-isolation .
else
python3 setup.py install
fi

384
third_party/sglang/scripts/ci/cuda/ci_install_dependency.sh vendored Executable file
View File

@@ -0,0 +1,384 @@
#!/bin/bash
# Install the dependency in CI.
#
# Structure (see section banners below):
# - Configuration & timing
# - Host / runner detection (arch, Blackwell, pip vs uv)
# - Kill existing processes
# - Install apt packages
# - Python package site hygiene & install protoc
# - Pip / uv toolchain & stale package cleanup
# - Uninstall Flashinfer
# - Install main package
# - Install sglang-kernel
# - Install sglang-router
# - Download flashinfer artifacts
# - Install extra dependency
# - Fix other dependencies
# - Prepare runner
# - Verify imports
set -euxo pipefail
# ------------------------------------------------------------------------------
# Configuration & timing
# ------------------------------------------------------------------------------
# Set up environment variables
CU_VERSION="cu129"
# Nvidia package versions we override (torch pins older versions).
# Used both as pip constraints during install and for post-install verification.
NVIDIA_CUDNN_VERSION="9.16.0.29"
NVIDIA_NVSHMEM_VERSION="3.4.5"
OPTIONAL_DEPS="${1:-}"
SECONDS=0
_CI_MARK_PREV=${SECONDS}
mark_step_done() {
local label=$1
local now=${SECONDS}
local step=$((now - _CI_MARK_PREV))
printf '\n[STEP DONE] %s, step: %ss, total: %ss, date: %s\n' \
"${label}" "${step}" "${now}" "$(date -u '+%Y-%m-%dT%H:%M:%SZ')"
_CI_MARK_PREV=${now}
}
mark_step_done "Configuration"
# ------------------------------------------------------------------------------
# Host / runner detection (CPU arch, Blackwell, USE_UV)
# ------------------------------------------------------------------------------
# Detect CPU architecture (x86_64 or aarch64)
ARCH=$(uname -m)
echo "Detected architecture: ${ARCH}"
# Detect GPU architecture (blackwell or not)
if [ "${IS_BLACKWELL+set}" = set ]; then
case "$IS_BLACKWELL" in 1 | true | yes) IS_BLACKWELL=1 ;; *) IS_BLACKWELL=0 ;; esac
echo "IS_BLACKWELL=${IS_BLACKWELL} (manually set via environment)"
else
IS_BLACKWELL=0
if command -v nvidia-smi >/dev/null 2>&1; then
while IFS= read -r cap; do
major="${cap%%.*}"
if [ "${major:-0}" -ge 10 ] 2>/dev/null; then
IS_BLACKWELL=1
break
fi
done <<< "$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null || true)"
fi
echo "IS_BLACKWELL=${IS_BLACKWELL} (auto-detected via nvidia-smi)"
fi
# Whether to use pip or uv to install dependencies
if [ "${USE_UV+set}" != set ]; then
if [ "$IS_BLACKWELL" = "1" ]; then
# Our current b200 runners have some issues with uv, so we default to pip
# It is a runner specific issue, not a GPU architecture issue.
USE_UV=false
else
USE_UV=true
fi
fi
case "$(printf '%s' "$USE_UV" | tr '[:upper:]' '[:lower:]')" in 1 | true | yes) USE_UV=1 ;; *) USE_UV=0 ;; esac
echo "USE_UV=${USE_UV}"
mark_step_done "Host / runner detection"
# ------------------------------------------------------------------------------
# Kill existing processes
# ------------------------------------------------------------------------------
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
python3 "${REPO_ROOT}/python/sglang/cli/killall.py"
KILLALL_EXIT=$?
echo "CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-}"
if [ $KILLALL_EXIT -ne 0 ]; then
echo "ERROR: killall.py detected uncleanable GPU memory. Aborting CI."
exit 1
fi
mark_step_done "Kill existing processes"
# ------------------------------------------------------------------------------
# Install apt packages
# ------------------------------------------------------------------------------
# Install apt packages (including python3/pip which may be missing on some runners)
# Use --no-install-recommends and ignore errors from unrelated broken packages on the runner
# The NVIDIA driver packages may have broken dependencies that are unrelated to these packages
# Run apt-get update first to refresh package index (stale index causes 404 on security.ubuntu.com)
apt-get update || true
CI_APT_PACKAGES=(
python3 python3-pip python3-venv python3-dev git libnuma-dev libssl-dev pkg-config
libibverbs-dev libibverbs1 ibverbs-providers ibverbs-utils
ffmpeg libavcodec-dev libavformat-dev libavutil-dev libswscale-dev
)
apt-get install -y --no-install-recommends "${CI_APT_PACKAGES[@]}" || {
echo "Warning: apt-get install failed, checking if required packages are available..."
for pkg in "${CI_APT_PACKAGES[@]}"; do
if ! dpkg -l "$pkg" 2>/dev/null | grep -q "^ii"; then
echo "ERROR: Required package $pkg is not installed and apt-get failed"
exit 1
fi
done
echo "All required packages are already installed, continuing..."
}
mark_step_done "Install apt packages"
# ------------------------------------------------------------------------------
# Python package site hygiene & install protoc
# ------------------------------------------------------------------------------
# Clear torch compilation cache
python3 -c 'import os, shutil, tempfile, getpass; cache_dir = os.environ.get("TORCHINDUCTOR_CACHE_DIR") or os.path.join(tempfile.gettempdir(), "torchinductor_" + getpass.getuser()); shutil.rmtree(cache_dir, ignore_errors=True)'
# Remove broken dist-info directories (missing METADATA per PEP 376)
SITE_PACKAGES=$(python3 -c "import site; print(site.getsitepackages()[0])")
if [ -d "$SITE_PACKAGES" ]; then
{ set +x; } 2>/dev/null
find "$SITE_PACKAGES" -maxdepth 1 -name "*.dist-info" -type d | while read -r d; do
if [ ! -f "$d/METADATA" ]; then
echo "Removing broken dist-info: $d"
rm -rf "$d"
fi
done
set -x
fi
# Install protoc
bash "${SCRIPT_DIR}/../utils/install_protoc.sh"
mark_step_done "Python package site hygiene & install protoc"
# ------------------------------------------------------------------------------
# Pip / uv toolchain & stale package cleanup
# ------------------------------------------------------------------------------
# Install pip and uv (use python3 -m pip for robustness since some runners only have pip3)
python3 -m pip install --upgrade pip
if [ "$USE_UV" = "0" ]; then
PIP_CMD="pip"
PIP_INSTALL_SUFFIX="--break-system-packages"
PIP_UNINSTALL_CMD="pip uninstall -y"
PIP_UNINSTALL_SUFFIX="--break-system-packages"
else
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"
PIP_UNINSTALL_SUFFIX=""
fi
# Clean up existing installations
$PIP_UNINSTALL_CMD sgl-kernel sglang-kernel sglang sgl-fa4 flash-attn-4 $PIP_UNINSTALL_SUFFIX || true
mark_step_done "Pip / uv toolchain & stale package cleanup"
# ------------------------------------------------------------------------------
# Uninstall Flashinfer
# ------------------------------------------------------------------------------
# Keep flashinfer packages installed if version matches to avoid re-downloading:
# - flashinfer-cubin: 150+ MB, plus extra cubins from ci_download_flashinfer_cubin.sh
# - flashinfer-jit-cache: 1.2+ GB, by far the largest download in CI
FLASHINFER_PYTHON_REQUIRED=$(grep -Po -m1 '(?<=flashinfer_python==)[0-9A-Za-z\.\-]+' python/pyproject.toml || echo "")
FLASHINFER_CUBIN_REQUIRED=$(grep -Po -m1 '(?<=flashinfer_cubin==)[0-9A-Za-z\.\-]+' python/pyproject.toml || echo "")
FLASHINFER_CUBIN_INSTALLED=$(pip show flashinfer-cubin 2>/dev/null | grep "^Version:" | awk '{print $2}' || echo "")
FLASHINFER_JIT_INSTALLED=$(pip show flashinfer-jit-cache 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed 's/+.*//' || echo "")
UNINSTALL_CUBIN=true
UNINSTALL_JIT_CACHE=true
if [ "$FLASHINFER_CUBIN_INSTALLED" = "$FLASHINFER_CUBIN_REQUIRED" ] && [ -n "$FLASHINFER_CUBIN_REQUIRED" ]; then
echo "flashinfer-cubin==${FLASHINFER_CUBIN_REQUIRED} already installed, keeping it"
UNINSTALL_CUBIN=false
else
echo "flashinfer-cubin version mismatch (installed: ${FLASHINFER_CUBIN_INSTALLED:-none}, required: ${FLASHINFER_CUBIN_REQUIRED}), reinstalling"
fi
if [ "$FLASHINFER_JIT_INSTALLED" = "$FLASHINFER_PYTHON_REQUIRED" ] && [ -n "$FLASHINFER_PYTHON_REQUIRED" ]; then
echo "flashinfer-jit-cache==${FLASHINFER_PYTHON_REQUIRED} already installed, keeping it"
UNINSTALL_JIT_CACHE=false
else
echo "flashinfer-jit-cache version mismatch (installed: ${FLASHINFER_JIT_INSTALLED:-none}, required: ${FLASHINFER_PYTHON_REQUIRED}), will reinstall"
fi
# Build uninstall list based on what needs updating
FLASHINFER_UNINSTALL="flashinfer-python"
[ "$UNINSTALL_CUBIN" = true ] && FLASHINFER_UNINSTALL="$FLASHINFER_UNINSTALL flashinfer-cubin"
[ "$UNINSTALL_JIT_CACHE" = true ] && FLASHINFER_UNINSTALL="$FLASHINFER_UNINSTALL flashinfer-jit-cache"
$PIP_UNINSTALL_CMD $FLASHINFER_UNINSTALL $PIP_UNINSTALL_SUFFIX || true
$PIP_UNINSTALL_CMD opencv-python opencv-python-headless $PIP_UNINSTALL_SUFFIX || true
mark_step_done "Uninstall Flashinfer"
# ------------------------------------------------------------------------------
# Install main package
# ------------------------------------------------------------------------------
# Install the main package
EXTRAS="dev,runai,tracing"
if [ -n "$OPTIONAL_DEPS" ]; then
EXTRAS="dev,runai,tracing,${OPTIONAL_DEPS}"
fi
echo "Installing python extras: [${EXTRAS}]"
source "$(dirname "$0")/cache_nvidia_wheels.sh"
$PIP_CMD install -e "python[${EXTRAS}]" --extra-index-url https://download.pytorch.org/whl/${CU_VERSION} $PIP_INSTALL_SUFFIX
mark_step_done "Install main package"
# ------------------------------------------------------------------------------
# Install sglang-kernel
# ------------------------------------------------------------------------------
# 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)
echo "SGL_KERNEL_VERSION_FROM_KERNEL=${SGL_KERNEL_VERSION_FROM_KERNEL} SGL_KERNEL_VERSION_FROM_SRT=${SGL_KERNEL_VERSION_FROM_SRT}"
if [ "${CUSTOM_BUILD_SGL_KERNEL:-}" = "true" ] && [ -d "sgl-kernel/dist" ]; then
ls -alh sgl-kernel/dist
# Determine wheel architecture
if [ "$ARCH" = "aarch64" ] || [ "$ARCH" = "arm64" ]; then
WHEEL_ARCH="aarch64"
else
WHEEL_ARCH="x86_64"
fi
$PIP_CMD install sgl-kernel/dist/sglang_kernel-${SGL_KERNEL_VERSION_FROM_KERNEL}-cp310-abi3-manylinux2014_${WHEEL_ARCH}.whl --force-reinstall $PIP_INSTALL_SUFFIX
elif [ "${CUSTOM_BUILD_SGL_KERNEL:-}" = "true" ] && [ ! -d "sgl-kernel/dist" ]; then
# CUSTOM_BUILD_SGL_KERNEL was set but artifacts not available (e.g., stage rerun without wheel build)
# Fail instead of falling back to PyPI - we need to test the built kernel, not PyPI version
echo "ERROR: CUSTOM_BUILD_SGL_KERNEL=true but sgl-kernel/dist not found."
echo "This usually happens when rerunning a stage without the sgl-kernel-build-wheels job."
echo "Please re-run the full workflow using /tag-and-rerun-ci to rebuild the kernel."
exit 1
else
# On Blackwell machines, skip reinstall if correct version already installed to avoid race conditions
if [ "$IS_BLACKWELL" = "1" ]; then
INSTALLED_SGL_KERNEL=$(pip show sglang-kernel 2>/dev/null | grep "^Version:" | awk '{print $2}' || echo "")
if [ "$INSTALLED_SGL_KERNEL" = "$SGL_KERNEL_VERSION_FROM_SRT" ]; then
echo "sglang-kernel==${SGL_KERNEL_VERSION_FROM_SRT} already installed, skipping reinstall"
else
echo "Installing sglang-kernel==${SGL_KERNEL_VERSION_FROM_SRT} (current: ${INSTALLED_SGL_KERNEL:-none})"
$PIP_CMD install sglang-kernel==${SGL_KERNEL_VERSION_FROM_SRT} $PIP_INSTALL_SUFFIX
fi
else
$PIP_CMD install sglang-kernel==${SGL_KERNEL_VERSION_FROM_SRT} --force-reinstall $PIP_INSTALL_SUFFIX
fi
fi
mark_step_done "Install sglang-kernel"
# ------------------------------------------------------------------------------
# Install sglang-router
# ------------------------------------------------------------------------------
# Install router for pd-disagg test
$PIP_CMD install sglang-router $PIP_INSTALL_SUFFIX
# Show current packages
$PIP_CMD list
mark_step_done "Install sglang-router"
# ------------------------------------------------------------------------------
# Download flashinfer artifacts
# ------------------------------------------------------------------------------
# Download flashinfer jit cache
UNINSTALL_JIT_CACHE="$UNINSTALL_JIT_CACHE" \
FLASHINFER_PYTHON_REQUIRED="$FLASHINFER_PYTHON_REQUIRED" \
CU_VERSION="$CU_VERSION" \
PIP_CMD="$PIP_CMD" \
PIP_INSTALL_SUFFIX="$PIP_INSTALL_SUFFIX" \
bash "${SCRIPT_DIR}/ci_download_flashinfer_jit_cache.sh"
# Download flashinfer cubins
bash "${SCRIPT_DIR}/ci_download_flashinfer_cubin.sh"
mark_step_done "Download flashinfer artifacts"
# ------------------------------------------------------------------------------
# Install extra dependency
# ------------------------------------------------------------------------------
# Install other python dependencies
if [ "$CU_VERSION" = "cu130" ]; then
NVRTC_SPEC="nvidia-cuda-nvrtc"
else
NVRTC_SPEC="nvidia-cuda-nvrtc-cu12"
fi
$PIP_CMD install mooncake-transfer-engine==0.3.10.post1 "${NVRTC_SPEC}" py-spy scipy huggingface_hub[hf_xet] pytest $PIP_INSTALL_SUFFIX
# Install other test dependencies
if [ "$IS_BLACKWELL" != "1" ]; then
# For lmms_evals evaluating MMMU
git clone --branch v0.5 --depth 1 https://github.com/EvolvingLMMs-Lab/lmms-eval.git
$PIP_CMD install -e lmms-eval/ $PIP_INSTALL_SUFFIX
fi
$PIP_CMD uninstall xformers || true
mark_step_done "Install extra dependency"
# ------------------------------------------------------------------------------
# Fix other dependencies
# ------------------------------------------------------------------------------
# Fix CUDA version mismatch between torch and torchaudio.
# PyPI's torch 2.9.1 bundles cu128 but torchaudio from pytorch.org/cu129 uses cu129.
# This mismatch causes torchaudio's C extension to fail loading, producing:
# "partially initialized module 'torchaudio' has no attribute 'lib'"
# We cannot replace torch with cu129 (breaks sgl_kernel ABI), so instead we reinstall
# torchaudio/torchvision from an index matching torch's CUDA version.
TORCH_CUDA_VER=$(python3 -c "import torch; v=torch.version.cuda; parts=v.split('.'); print(f'cu{parts[0]}{parts[1]}')")
echo "Detected torch CUDA version: ${TORCH_CUDA_VER}"
if [ "${TORCH_CUDA_VER}" != "${CU_VERSION}" ]; then
# Pin versions to match what was installed by pyproject.toml (strip +cuXYZ suffix)
TORCHAUDIO_VER=$(pip show torchaudio 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed 's/+.*//')
TORCHVISION_VER=$(pip show torchvision 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed 's/+.*//')
echo "Reinstalling torchaudio==${TORCHAUDIO_VER} torchvision==${TORCHVISION_VER} from ${TORCH_CUDA_VER} index to match torch..."
$PIP_CMD install "torchaudio==${TORCHAUDIO_VER}" "torchvision==${TORCHVISION_VER}" --index-url "https://download.pytorch.org/whl/${TORCH_CUDA_VER}" --force-reinstall --no-deps $PIP_INSTALL_SUFFIX
fi
# Fix dependencies: DeepEP depends on nvshmem 3.4.5 — skip reinstall when already correct (avoids pip races / wasted work)
INSTALLED_NVSHMEM=$(pip show nvidia-nvshmem-cu12 2>/dev/null | grep "^Version:" | awk '{print $2}' || echo "")
if [ "$INSTALLED_NVSHMEM" = "$NVIDIA_NVSHMEM_VERSION" ]; then
echo "nvidia-nvshmem-cu12==${NVIDIA_NVSHMEM_VERSION} already installed, skipping reinstall"
else
$PIP_CMD install nvidia-nvshmem-cu12==${NVIDIA_NVSHMEM_VERSION} $PIP_INSTALL_SUFFIX
fi
# Fix dependencies: Cudnn with version less than 9.16.0.29 will cause performance regression on Conv3D kernel
INSTALLED_CUDNN=$(pip show nvidia-cudnn-cu12 2>/dev/null | grep "^Version:" | awk '{print $2}' || echo "")
if [ "$INSTALLED_CUDNN" = "$NVIDIA_CUDNN_VERSION" ]; then
echo "nvidia-cudnn-cu12==${NVIDIA_CUDNN_VERSION} already installed, skipping reinstall"
else
$PIP_CMD install nvidia-cudnn-cu12==${NVIDIA_CUDNN_VERSION} $PIP_INSTALL_SUFFIX
fi
mark_step_done "Fix other dependencies"
# Force reinstall nvidia-cutlass-dsl to ensure the .pth file exists.
# The Docker image ships nvidia-cutlass-dsl-libs-base 4.3.5; upgrading to 4.4.2
# can delete the .pth file without reliably recreating it (pip race condition).
$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
# Install human-eval
pip install "setuptools==70.0.0"
git clone https://github.com/merrymercy/human-eval.git
cd human-eval
pip install -e . --no-build-isolation
# ------------------------------------------------------------------------------
# Prepare runner
# ------------------------------------------------------------------------------
# Prepare the CI runner (cleanup HuggingFace cache, etc.)
bash "${SCRIPT_DIR}/prepare_runner.sh"
mark_step_done "Prepare runner"
# ------------------------------------------------------------------------------
# Verify imports
# ------------------------------------------------------------------------------
# Show current packages
$PIP_CMD list
python3 -c "import torch; print(torch.version.cuda)"
python3 -c "import cutlass; import cutlass.cute;"
mark_step_done "Verify imports"

24
third_party/sglang/scripts/ci/cuda/ci_install_gateway_dependencies.sh vendored Executable file
View File

@@ -0,0 +1,24 @@
#!/bin/bash
set -euxo pipefail
# Check if sudo is available
if command -v sudo >/dev/null 2>&1; then
sudo apt-get update
sudo apt-get install -y libssl-dev pkg-config protobuf-compiler redis-server
else
apt-get update
apt-get install -y libssl-dev pkg-config protobuf-compiler redis-server
fi
# Install rustup (Rust installer and version manager)
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain 1.90
# Follow the installation prompts, then reload your shell
. "$HOME/.cargo/env"
source $HOME/.cargo/env
# Verify installation
rustc --version
cargo --version
protoc --version

106
third_party/sglang/scripts/ci/cuda/ci_start_disaggregation_servers.sh vendored Executable file
View File

@@ -0,0 +1,106 @@
#!/bin/bash
set -euo pipefail
# Optional: set DISAGG_READY_FILE to a filepath; when all servers are healthy, the script will
# create this file as a readiness signal (useful for CI to proceed to next steps).
DISAGG_READY_FILE="${DISAGG_READY_FILE:-}"
MODEL_PATH="/raid/models/meta-llama/Llama-3.1-8B-Instruct"
# Function to find the first available active IB device
find_active_ib_device() {
for device in mlx5_{0..11}; do
if ibv_devinfo $device >/dev/null 2>&1; then
state=$(ibv_devinfo $device | grep "state:" | head -1 | awk '{print $2}')
if [[ "$state" == "PORT_ACTIVE" ]]; then
echo "$device"
return 0
fi
fi
done
echo "No active IB device found" >&2
return 1
}
# Get the first available active IB device
DEVICE=$(find_active_ib_device)
echo "Using IB device: $DEVICE"
# Launch prefill servers on GPU 0–3
for i in {0..3}; do
PORT=$((30001 + i))
BOOTSTRAP_PORT=$((9001 + i))
HOST="127.0.0.$((i + 1))"
echo "Launching PREFILL server on GPU $i at $HOST:$PORT (bootstrap: $BOOTSTRAP_PORT)"
CUDA_VISIBLE_DEVICES=$i \
python3 -m sglang.launch_server \
--model-path "$MODEL_PATH" \
--disaggregation-mode prefill \
--host "$HOST" \
--port "$PORT" \
--disaggregation-ib-device "$DEVICE" \
--disaggregation-bootstrap-port "$BOOTSTRAP_PORT" &
done
# Launch decode servers on GPU 4–7
for i in {4..7}; do
PORT=$((30001 + i))
HOST="127.0.0.$((i + 1))"
echo "Launching DECODE server on GPU $i at $HOST:$PORT"
CUDA_VISIBLE_DEVICES=$i \
python3 -m sglang.launch_server \
--model-path "$MODEL_PATH" \
--disaggregation-mode decode \
--host "$HOST" \
--port "$PORT" \
--disaggregation-ib-device "$DEVICE" \
--base-gpu-id 0 &
done
# Wait for disaggregation servers to initialize
echo "Waiting for disaggregation servers to initialize..."
# Health check with 5-minute timeout
TIMEOUT=300
START_TIME=$(date +%s)
echo "Checking health of all 8 servers..."
while true; do
CURRENT_TIME=$(date +%s)
ELAPSED=$((CURRENT_TIME - START_TIME))
if [ $ELAPSED -ge $TIMEOUT ]; then
echo "❌ Timeout: Servers did not become healthy within 5 minutes"
exit 1
fi
HEALTHY_COUNT=0
# Check all 8 servers (127.0.0.1-8:30001-30008)
for i in {1..8}; do
if curl -s -f "http://127.0.0.$i:$((30000 + i))/health" >/dev/null 2>&1; then
HEALTHY_COUNT=$((HEALTHY_COUNT + 1))
fi
done
echo "Healthy servers: $HEALTHY_COUNT/8 (elapsed: ${ELAPSED}s)"
if [ $HEALTHY_COUNT -eq 8 ]; then
echo "✅ All 8 servers are healthy!"
# Emit readiness signal file if requested
if [ -n "$DISAGG_READY_FILE" ]; then
echo "Creating readiness flag: $DISAGG_READY_FILE"
# Ensure parent dir exists; ignore errors
mkdir -p "$(dirname "$DISAGG_READY_FILE")" 2>/dev/null || true
touch "$DISAGG_READY_FILE"
fi
break
else
sleep 10 # Wait 10 seconds before next check
fi
done
# Don't launch router here - just keep servers running
echo "✅ All disaggregation servers are ready and waiting for router connections"
# Keep the script running
wait

19
third_party/sglang/scripts/ci/cuda/prepare_runner.sh vendored Executable file
View File

@@ -0,0 +1,19 @@
#!/bin/bash
# Prepare the CI runner by cleaning up stale HuggingFace cache artifacts and validating models
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
echo "Preparing CI runner..."
echo ""
# Clean up stale HuggingFace cache artifacts from previous failed downloads
python3 "${SCRIPT_DIR}/../utils/cleanup_hf_cache.py"
echo ""
# Pre-validate cached models and write markers for offline mode
# This allows tests to run with HF_HUB_OFFLINE=1 for models that are fully cached
python3 "${SCRIPT_DIR}/../utils/prevalidate_cached_models.py"
echo ""
echo "CI runner preparation complete!"

399
third_party/sglang/scripts/ci/cuda/warmup_deep_gemm.py vendored Normal file
View File

@@ -0,0 +1,399 @@
"""
Lightweight DeepGEMM JIT compilation warmup without loading model weights.
Reads model config.json from HF cache to derive kernel shapes, then compiles
DeepGEMM kernels directly. This avoids the expensive model weight loading step
that the full `sglang.compile_deep_gemm` requires.
Supports DeepSeek V2/V3 family models. Falls back to `sglang.compile_deep_gemm`
for unsupported architectures.
Usage:
python3 scripts/ci/cuda/warmup_deep_gemm.py \
deepseek-ai/DeepSeek-V3-0324:8 \
deepseek-ai/DeepSeek-V3.2-Exp:8
"""
import json
import os
import subprocess
import sys
import time
from math import ceil
from pathlib import Path
# Configure DeepGEMM cache before importing deep_gemm
os.environ["DG_JIT_CACHE_DIR"] = os.getenv(
"SGLANG_DG_CACHE_DIR",
os.path.join(os.path.expanduser("~"), ".cache", "deep_gemm"),
)
os.environ["DG_JIT_USE_NVRTC"] = os.getenv("SGL_DG_USE_NVRTC", "0")
BLOCK_SIZE = 128
def get_config_json(model_name):
"""Load config.json for a cached model from HF cache."""
cache_dir = os.environ.get(
"HF_HOME", os.path.join(os.path.expanduser("~"), ".cache", "huggingface")
)
hub_dir = os.path.join(cache_dir, "hub")
safe_name = "models--" + model_name.replace("/", "--")
snapshots_dir = os.path.join(hub_dir, safe_name, "snapshots")
if not os.path.isdir(snapshots_dir):
return None
snapshots = sorted(
Path(snapshots_dir).iterdir(), key=lambda p: p.stat().st_mtime, reverse=True
)
for snapshot in snapshots:
config_path = snapshot / "config.json"
if config_path.exists():
with open(config_path) as f:
return json.load(f)
return None
def is_deepseek_v2v3(config):
"""Check if a model is from the DeepSeek V2/V3 family."""
architectures = config.get("architectures", [])
model_type = config.get("model_type", "")
return any(
"DeepseekV2" in a or "DeepseekV3" in a for a in architectures
) or model_type in ("deepseek_v2", "deepseek_v3")
def compute_deepseek_v2v3_shapes(config, tp):
"""Compute all DeepGEMM (kernel_type, N, K, num_groups) for DeepSeek V2/V3.
Shape derivation based on:
- MoE: python/sglang/srt/layers/moe/fused_moe_triton/layer.py
- MLA: python/sglang/srt/models/deepseek_v2.py
- FP8: python/sglang/srt/layers/quantization/fp8_kernel.py
"""
shapes = []
hidden_size = config["hidden_size"]
num_attention_heads = config.get("num_attention_heads", 128)
kv_lora_rank = config.get("kv_lora_rank", 512)
qk_nope_head_dim = config.get("qk_nope_head_dim", 128)
v_head_dim = config.get("v_head_dim", 128)
n_routed_experts = config.get("n_routed_experts", 0)
n_shared_experts = config.get("n_shared_experts", 0)
moe_intermediate_size = config.get("moe_intermediate_size", 0)
num_local_heads = num_attention_heads // tp
# Shared expert fusion is enabled by default (disable_shared_experts_fusion=False)
# so the FusedMoE weight tensor includes shared experts
num_local_experts = n_routed_experts + n_shared_experts
# --- MoE expert GEMM shapes ---
# FusedMoE shards intermediate_size across TP ranks (column parallel for gate/up,
# row parallel for down). All experts are replicated on each TP rank.
if n_routed_experts > 0 and moe_intermediate_size > 0:
moe_inter_per_tp = moe_intermediate_size // tp
# Gate-Up projection: (tokens, hidden_size) @ (experts, 2*inter_per_tp, hidden_size)^T
# Both masked and contiguous paths are used at runtime
shapes.append(("MASKED", moe_inter_per_tp * 2, hidden_size, num_local_experts))
shapes.append(("CONTIG", moe_inter_per_tp * 2, hidden_size, num_local_experts))
# Down projection: (tokens, inter_per_tp) @ (experts, hidden_size, inter_per_tp)^T
shapes.append(("MASKED", hidden_size, moe_inter_per_tp, num_local_experts))
shapes.append(("CONTIG", hidden_size, moe_inter_per_tp, num_local_experts))
# --- MLA attention GEMM shapes (masked grouped GEMM) ---
if kv_lora_rank > 0 and num_local_heads > 0:
# Q_nope -> compressed K: (heads, m, qk_nope_head_dim) @ (heads, kv_lora_rank, qk_nope_head_dim)^T
shapes.append(("MASKED", kv_lora_rank, qk_nope_head_dim, num_local_heads))
# Attention output -> V: (heads, m, kv_lora_rank) @ (heads, v_head_dim, kv_lora_rank)^T
shapes.append(("MASKED", v_head_dim, kv_lora_rank, num_local_heads))
# --- kv_b_proj (non-grouped GEMM via FP8 kernel) ---
# ColumnParallelLinear(kv_lora_rank, num_heads * (qk_nope + v_head_dim))
# Per TP rank: N = num_local_heads * (qk_nope_head_dim + v_head_dim)
if kv_lora_rank > 0 and num_local_heads > 0:
kv_b_proj_n = num_local_heads * (qk_nope_head_dim + v_head_dim)
shapes.append(("NORMAL", kv_b_proj_n, kv_lora_rank, 1))
return shapes
def get_architecture_key(config, tp):
"""Key for dedup: models with same key share DeepGEMM kernels."""
if config is None:
return None
fields = [
config.get("hidden_size", 0),
config.get("moe_intermediate_size", 0),
config.get("n_routed_experts", 0),
config.get("n_shared_experts", 0),
config.get("num_attention_heads", 0),
config.get("kv_lora_rank", 0),
config.get("qk_nope_head_dim", 0),
config.get("v_head_dim", 0),
tp,
]
return tuple(fields)
def compute_m_list(fast_warmup=False, chunked_prefill_size=8192):
"""Compute the list of M values to compile (matches compile_utils.py logic)."""
m_list = []
if fast_warmup:
m_list += list(range(1, 1025))
next_m, sample_step = 1024, 2
max_prefill_bs = min(chunked_prefill_size, 32 * 1024)
while next_m < max_prefill_bs:
m_list += list(range(next_m, 2 * next_m, sample_step))
next_m *= 2
sample_step *= 2
m_list.append(max_prefill_bs)
m_list = sorted(set(m_list))
else:
m_max = 16 * 1024
if chunked_prefill_size > 8192:
m_max = chunked_prefill_size * 2
m_max = min(128 * 1024, m_max)
m_list = list(range(1, m_max + 1))
return m_list
def _empty_token_fp8(size):
"""Create FP8 token tensor + per-block scale tensor."""
import torch
*dims, k = size
return (
torch.empty(size, device="cuda", dtype=torch.float8_e4m3fn),
torch.empty((*dims, ceil(k / BLOCK_SIZE)), device="cuda", dtype=torch.float32),
)
def _empty_block_fp8(size):
"""Create FP8 block tensor + per-block scale tensor."""
import torch
*dims, n, k = size
return (
torch.empty(size, device="cuda", dtype=torch.float8_e4m3fn),
torch.empty(
(*dims, ceil(n / BLOCK_SIZE), ceil(k / BLOCK_SIZE)),
device="cuda",
dtype=torch.float32,
),
)
def get_memory_requirement(kernel_type, max_m, n, k, num_groups):
"""Estimate GPU memory needed in GB for compilation buffers."""
_GB = 1 << 30
if kernel_type == "NORMAL":
return (max_m * k + n * k + max_m * n * 2) / _GB
elif kernel_type == "CONTIG":
return (max_m * k + num_groups * n * k + max_m * 4 + max_m * n * 2) / _GB
elif kernel_type == "MASKED":
return (
num_groups * max_m * k
+ num_groups * n * k
+ num_groups * 4
+ num_groups * max_m * n * 2
) / _GB
return 0
def compile_one_shape(kernel_type, n, k, num_groups, m_list):
"""Compile DeepGEMM kernels for one (kernel_type, N, K, num_groups) shape."""
import deep_gemm
import torch
from tqdm import tqdm
# Filter M list for contiguous layout alignment
if kernel_type == "CONTIG":
m_alignment = deep_gemm.get_mk_alignment_for_contiguous_layout()
m_list = sorted(set(m for m in m_list if m % m_alignment == 0))
if not m_list:
return
max_m = max(m_list)
# Reduce max_m if not enough GPU memory
mem_free = torch.cuda.mem_get_info()[0] / (1 << 30)
mem_required = get_memory_requirement(kernel_type, max_m, n, k, num_groups)
if mem_required > mem_free:
while (
get_memory_requirement(kernel_type, max_m, n, k, num_groups) > mem_free
and max_m > 4096
):
max_m //= 2
print(
f" Memory {mem_free:.1f}GB < required {mem_required:.1f}GB, "
f"reducing max_m to {max_m}"
)
m_list = [m for m in m_list if m <= max_m]
old_mode = deep_gemm.get_compile_mode()
deep_gemm.set_compile_mode(1)
try:
if kernel_type == "NORMAL":
lhs_q, lhs_s = _empty_token_fp8((max_m, k))
rhs_q, rhs_s = _empty_block_fp8((n, k))
out = torch.empty((max_m, n), device="cuda", dtype=torch.bfloat16)
for m in tqdm(m_list, desc=f" NORMAL N={n} K={k}"):
deep_gemm.fp8_gemm_nt((lhs_q[:m], lhs_s[:m]), (rhs_q, rhs_s), out[:m])
elif kernel_type == "CONTIG":
lhs_q, lhs_s = _empty_token_fp8((max_m, k))
rhs_q, rhs_s = _empty_block_fp8((num_groups, n, k))
m_indices = torch.zeros((max_m,), device="cuda", dtype=torch.int32)
out = torch.empty((max_m, n), device="cuda", dtype=torch.bfloat16)
for m in tqdm(m_list, desc=f" CONTIG N={n} K={k} G={num_groups}"):
deep_gemm.m_grouped_fp8_gemm_nt_contiguous(
(lhs_q[:m], lhs_s[:m]),
(rhs_q, rhs_s),
out[:m],
m_indices=m_indices[:m],
)
elif kernel_type == "MASKED":
lhs_q, lhs_s = _empty_token_fp8((num_groups, max_m, k))
rhs_q, rhs_s = _empty_block_fp8((num_groups, n, k))
masked_m = torch.zeros((num_groups,), device="cuda", dtype=torch.int32)
out = torch.empty(
(num_groups, max_m, n), device="cuda", dtype=torch.bfloat16
)
for m in tqdm(m_list, desc=f" MASKED N={n} K={k} G={num_groups}"):
deep_gemm.fp8_m_grouped_gemm_nt_masked(
(lhs_q, lhs_s),
(rhs_q, rhs_s),
out,
masked_m=masked_m,
expected_m=m,
)
finally:
deep_gemm.set_compile_mode(old_mode)
torch.cuda.current_stream().synchronize()
torch.cuda.empty_cache()
def compile_shapes_lightweight(shapes, m_list):
"""Compile all DeepGEMM shapes directly (no model loading)."""
for i, (kernel_type, n, k, num_groups) in enumerate(shapes, 1):
print(f"\n[{i}/{len(shapes)}] {kernel_type} N={n} K={k} G={num_groups}")
t0 = time.time()
compile_one_shape(kernel_type, n, k, num_groups, m_list)
elapsed = time.time() - t0
print(f" Done in {elapsed:.1f}s")
def fallback_compile_deep_gemm(model, tp):
"""Fall back to full sglang.compile_deep_gemm (loads model weights)."""
print(f"Falling back to full compile_deep_gemm for {model} (tp={tp})...")
cmd = [
sys.executable,
"-m",
"sglang.compile_deep_gemm",
"--model",
model,
"--tp",
str(tp),
"--trust-remote-code",
"--model-loader-extra-config",
'{"enable_multithread_load": true, "num_threads": 64}',
]
result = subprocess.run(cmd)
if result.returncode != 0:
print(f"Warning: fallback failed for {model} (exit code {result.returncode})")
return result.returncode == 0
def main():
if len(sys.argv) < 2 or sys.argv[1] in ("-h", "--help"):
print("Usage: warmup_deep_gemm.py model1:tp1 [model2:tp2 ...]")
print("\nDerives DeepGEMM kernel shapes from config.json without loading model")
print(
"weights. Falls back to full compile_deep_gemm for unknown architectures."
)
sys.exit(0)
# Parse model:tp pairs
model_tp_pairs = []
for arg in sys.argv[1:]:
if ":" not in arg:
print(f"Error: expected model:tp format, got '{arg}'")
sys.exit(1)
model, tp_str = arg.rsplit(":", 1)
model_tp_pairs.append((model, int(tp_str)))
fast_warmup = os.environ.get("SGLANG_JIT_DEEPGEMM_FAST_WARMUP", "0").lower() in (
"1",
"true",
)
print(f"=== DeepGEMM Lightweight Warmup ({len(model_tp_pairs)} model(s)) ===")
print(f" Fast warmup: {fast_warmup}")
print(
f" Cache dir: {os.environ.get('DG_JIT_CACHE_DIR', '~/.cache/deep_gemm')}\n"
)
# Load configs and deduplicate by architecture
seen_keys = {}
to_process = [] # (model, tp, config_or_None, shapes_or_None)
for model, tp in model_tp_pairs:
config = get_config_json(model)
if config is None:
print(f" SKIP {model} (tp={tp}): config.json not in HF cache")
continue
key = get_architecture_key(config, tp)
if key in seen_keys:
print(f" DEDUP {model} (tp={tp}): same shapes as {seen_keys[key]}")
continue
if is_deepseek_v2v3(config):
shapes = compute_deepseek_v2v3_shapes(config, tp)
seen_keys[key] = model
to_process.append((model, tp, config, shapes))
print(f" FOUND {model} (tp={tp}): {len(shapes)} DeepGEMM shape(s)")
else:
# Unknown architecture: will use fallback
seen_keys[key] = model
to_process.append((model, tp, config, None))
arch = config.get("architectures", ["unknown"])
print(f" FOUND {model} (tp={tp}): unknown arch {arch}, will use fallback")
if not to_process:
print("\nNo models to process. Done.")
return
m_list = compute_m_list(fast_warmup=fast_warmup)
print(f"\nM list: {len(m_list)} values (range {min(m_list)}-{max(m_list)})")
for model, tp, config, shapes in to_process:
print(f"\n{'=' * 60}")
print(f"Model: {model} (tp={tp})")
print(f"{'=' * 60}")
if shapes is None:
# Unknown architecture: fall back to full compile_deep_gemm
fallback_compile_deep_gemm(model, tp)
continue
# Print shape summary
for kernel_type, n, k, num_groups in shapes:
print(f" {kernel_type:8s} N={n:<6d} K={k:<6d} G={num_groups}")
t0 = time.time()
compile_shapes_lightweight(shapes, m_list)
elapsed = time.time() - t0
print(f"\nCompleted {model} in {elapsed:.1f}s")
print("\nDeepGEMM lightweight warmup complete.")
if __name__ == "__main__":
main()

313
third_party/sglang/scripts/ci/cuda/warmup_server.py vendored Normal file
View File

@@ -0,0 +1,313 @@
"""
Full server warmup to pre-warm Triton autotuning and CUDA graph capture.
On cold H200 nodes (new nodes or after container recreation), CUDA graph capture
triggers Triton autotuning which takes ~330s per server launch. This script
launches actual servers with CUDA graphs enabled to cache the autotuned kernels,
so subsequent test launches are fast (~30-60s).
Uses marker files to skip warmup on already-warm nodes. Marker files are
invalidated when Python, Triton, or PyTorch versions change.
Usage:
python3 scripts/ci/cuda/warmup_server.py \
deepseek-ai/DeepSeek-V3-0324:8 \
inclusionAI/Ring-2.5-1T:8
"""
import hashlib
import json
import os
import signal
import subprocess
import sys
import tempfile
import time
from pathlib import Path
# Reuse helpers from warmup_deep_gemm (same directory)
sys.path.insert(0, os.path.dirname(__file__))
from warmup_deep_gemm import get_architecture_key, get_config_json
MARKER_DIR = os.path.join(os.path.expanduser("~"), ".cache", "sglang", "warmup_markers")
HEALTH_POLL_INTERVAL = 10 # seconds between health checks
SERVER_STARTUP_TIMEOUT = 900 # 15 min max to wait for server ready
DEFAULT_PORT = 39876
def get_version_key():
"""Hash of Python + Triton + PyTorch versions to invalidate markers on upgrades."""
parts = [sys.version]
try:
import triton
parts.append(f"triton={triton.__version__}")
except ImportError:
parts.append("triton=none")
try:
import torch
parts.append(f"torch={torch.__version__}")
except ImportError:
parts.append("torch=none")
return hashlib.sha256("|".join(parts).encode()).hexdigest()[:12]
def get_marker_path(model, tp):
"""Get the marker file path for a model:tp pair."""
version_key = get_version_key()
safe_model = model.replace("/", "--")
return os.path.join(
MARKER_DIR, f"server_warmup_{safe_model}_tp{tp}_{version_key}.done"
)
def check_marker(model, tp):
"""Check if warmup marker exists (node already warm)."""
marker = get_marker_path(model, tp)
return os.path.exists(marker)
def write_marker(model, tp):
"""Write warmup marker after successful warmup."""
marker = get_marker_path(model, tp)
os.makedirs(os.path.dirname(marker), exist_ok=True)
Path(marker).write_text(
json.dumps(
{
"model": model,
"tp": tp,
"version_key": get_version_key(),
"timestamp": time.time(),
}
)
)
print(f" Wrote marker: {marker}")
def kill_server(proc):
"""Kill server process tree."""
if proc.poll() is not None:
return
try:
os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
except (ProcessLookupError, OSError):
pass
try:
proc.wait(timeout=15)
except subprocess.TimeoutExpired:
try:
os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
except (ProcessLookupError, OSError):
pass
try:
proc.wait(timeout=5)
except subprocess.TimeoutExpired:
pass
def wait_for_server(base_url, proc, timeout):
"""Poll /health_generate until server is ready or timeout."""
import requests
start = time.time()
while time.time() - start < timeout:
ret = proc.poll()
if ret is not None:
return False, f"Server exited with code {ret}"
try:
resp = requests.get(f"{base_url}/health_generate", timeout=5)
if resp.status_code == 200:
return True, None
except requests.RequestException:
pass
time.sleep(HEALTH_POLL_INTERVAL)
return False, "Timed out waiting for server"
def send_generate_request(base_url):
"""Send one /generate request to exercise the full inference path."""
import requests
payload = {
"input_ids": [0, 1, 2, 3],
"sampling_params": {
"max_new_tokens": 8,
"temperature": 0,
},
}
try:
resp = requests.post(f"{base_url}/generate", json=payload, timeout=120)
if resp.status_code == 200:
print(" Generate request succeeded")
else:
print(f" Warning: generate request returned {resp.status_code}")
except requests.RequestException as e:
print(f" Warning: generate request failed: {e}")
def warmup_one_model(model, tp, port):
"""Launch server, wait for ready, send one request, then kill."""
base_url = f"http://127.0.0.1:{port}"
cmd = [
sys.executable,
"-m",
"sglang.launch_server",
"--model-path",
model,
"--tp",
str(tp),
"--host",
"127.0.0.1",
"--port",
str(port),
"--trust-remote-code",
"--model-loader-extra-config",
'{"enable_multithread_load": true, "num_threads": 64}',
]
# Use a temp file for server output to avoid pipe buffer deadlock
# (server logs can exceed the 64KB pipe buffer during CUDA graph capture)
log_file = tempfile.NamedTemporaryFile(
mode="w", prefix="warmup_server_", suffix=".log", delete=False
)
log_path = log_file.name
print(f" Launching server: {' '.join(cmd)}")
print(f" Server log: {log_path}")
proc = subprocess.Popen(
cmd,
stdout=log_file,
stderr=subprocess.STDOUT,
preexec_fn=os.setsid,
)
try:
# Wait for server to be ready (includes CUDA graph capture)
print(
f" Waiting for server (timeout={SERVER_STARTUP_TIMEOUT}s, "
f"polling every {HEALTH_POLL_INTERVAL}s)..."
)
ok, err = wait_for_server(base_url, proc, SERVER_STARTUP_TIMEOUT)
if not ok:
print(f" Warning: server not ready: {err}")
# Dump last lines of server log for debugging
try:
log_file.flush()
with open(log_path) as f:
lines = f.readlines()
for line in lines[-20:]:
print(f" | {line.rstrip()}")
except Exception:
pass
return False
print(" Server ready, sending generate request...")
send_generate_request(base_url)
return True
finally:
print(" Killing server...")
kill_server(proc)
log_file.close()
try:
os.unlink(log_path)
except OSError:
pass
def main():
if len(sys.argv) < 2 or sys.argv[1] in ("-h", "--help"):
print("Usage: warmup_server.py model1:tp1 [model2:tp2 ...]")
print(
"\nLaunches full servers with CUDA graphs enabled to pre-warm"
" Triton autotuning."
)
print("Skips instantly on warm nodes (marker file exists).")
sys.exit(0)
# Parse model:tp pairs
model_tp_pairs = []
for arg in sys.argv[1:]:
if ":" not in arg:
print(f"Error: expected model:tp format, got '{arg}'")
sys.exit(1)
model, tp_str = arg.rsplit(":", 1)
model_tp_pairs.append((model, int(tp_str)))
print(f"=== Server CUDA Graph Warmup ({len(model_tp_pairs)} model(s)) ===")
print(f" Marker dir: {MARKER_DIR}")
print(f" Version key: {get_version_key()}\n")
# Deduplicate by architecture and check markers
seen_keys = {}
to_warmup = []
for model, tp in model_tp_pairs:
# Check marker first (fast path)
if check_marker(model, tp):
print(f" SKIP {model} (tp={tp}): already warm (marker exists)")
continue
# Architecture dedup
config = get_config_json(model)
if config is not None:
key = get_architecture_key(config, tp)
if key in seen_keys:
print(
f" DEDUP {model} (tp={tp}): same architecture as {seen_keys[key]}"
)
continue
seen_keys[key] = model
to_warmup.append((model, tp))
print(f" QUEUE {model} (tp={tp}): needs warmup")
if not to_warmup:
print("\nAll models already warm. Done.")
return
print(f"\n{len(to_warmup)} model(s) to warm up.\n")
port = DEFAULT_PORT
for i, (model, tp) in enumerate(to_warmup, 1):
print(f"\n{'=' * 60}")
print(f"[{i}/{len(to_warmup)}] {model} (tp={tp})")
print(f"{'=' * 60}")
t0 = time.time()
success = warmup_one_model(model, tp, port)
elapsed = time.time() - t0
if success:
print(f" Completed in {elapsed:.0f}s")
write_marker(model, tp)
# Also write markers for dedup'd models that share this architecture
config = get_config_json(model)
if config is not None:
key = get_architecture_key(config, tp)
for other_model, other_tp in model_tp_pairs:
if (other_model, other_tp) == (model, tp):
continue
other_config = get_config_json(other_model)
if other_config is not None:
other_key = get_architecture_key(other_config, other_tp)
if other_key == key and not check_marker(other_model, other_tp):
write_marker(other_model, other_tp)
print(
f" Also marked {other_model} (tp={other_tp}) as warm (same arch)"
)
else:
print(
f" Warning: warmup failed after {elapsed:.0f}s (non-fatal, tests will still work)"
)
# Use a different port for the next model to avoid bind conflicts
port += 100
print("\nServer CUDA graph warmup complete.")
if __name__ == "__main__":
main()