From 94e1251131ca27260cb0e8938aeb7b4a4e630b19 Mon Sep 17 00:00:00 2001
From: gongwei-130 <56567052+gongwei-130@users.noreply.github.com>
Date: Thu, 11 Dec 2025 23:20:44 -0800
Subject: [PATCH] enable flashinfer-jit-cache in image build and ci install to
 speed up model launch (#14959)

---
 .github/workflows/release-docker.yml | 5 +++++
 scripts/ci/ci_install_dependency.sh  | 3 +++
 2 files changed, 8 insertions(+)

diff --git a/.github/workflows/release-docker.yml b/.github/workflows/release-docker.yml
index a95e5ddf5..5360a2c14 100644
--- a/.github/workflows/release-docker.yml
+++ b/.github/workflows/release-docker.yml
@@ -69,6 +69,7 @@ jobs:
             --build-arg CUDA_VERSION=${{ matrix.variant.cuda_version }} \
             --build-arg BUILD_TYPE=${{ matrix.variant.build_type }} \
             --build-arg GRACE_BLACKWELL=${{ matrix.variant.grace_blackwell }} \
+            --build-arg INSTALL_FLASHINFER_JIT_CACHE=1 \
             -t lmsysorg/sglang:${tag} \
             --no-cache \
             .
@@ -86,6 +87,7 @@ jobs:
             --build-arg CUDA_VERSION=${{ matrix.variant.cuda_version }} \
             --build-arg BUILD_TYPE=${{ matrix.variant.build_type }} \
             --build-arg GRACE_BLACKWELL=${{ matrix.variant.grace_blackwell }} \
+            --build-arg INSTALL_FLASHINFER_JIT_CACHE=1 \
             -t lmsysorg/sglang:${tag} \
             --no-cache \
             .
@@ -102,6 +104,7 @@ jobs:
             -f docker/Dockerfile \
             --build-arg CUDA_VERSION=13.0.1 \
             --build-arg BUILD_TYPE=${{ matrix.variant.build_type }} \
+            --build-arg INSTALL_FLASHINFER_JIT_CACHE=1 \
             --build-arg GRACE_BLACKWELL=0 \
             -t lmsysorg/sglang:${tag} \
             --no-cache \
@@ -146,6 +149,7 @@ jobs:
             --build-arg CUDA_VERSION=${{ matrix.variant.cuda_version }} \
             --build-arg BUILD_TYPE=${{ matrix.variant.build_type }} \
             --build-arg GRACE_BLACKWELL=${{ matrix.variant.grace_blackwell }} \
+            --build-arg INSTALL_FLASHINFER_JIT_CACHE=1 \
             -t lmsysorg/sglang:${tag} \
             --no-cache \
             .
@@ -163,6 +167,7 @@ jobs:
             --build-arg CUDA_VERSION=${{ matrix.variant.cuda_version }} \
             --build-arg BUILD_TYPE=${{ matrix.variant.build_type }} \
             --build-arg GRACE_BLACKWELL=${{ matrix.variant.grace_blackwell }} \
+            --build-arg INSTALL_FLASHINFER_JIT_CACHE=1 \
             -t lmsysorg/sglang:${tag} \
             --no-cache \
             .
diff --git a/scripts/ci/ci_install_dependency.sh b/scripts/ci/ci_install_dependency.sh
index df9654b12..2a56a0460 100755
--- a/scripts/ci/ci_install_dependency.sh
+++ b/scripts/ci/ci_install_dependency.sh
@@ -5,6 +5,7 @@ set -euxo pipefail
 # Set up environment variables
 IS_BLACKWELL=${IS_BLACKWELL:-0}
 CU_VERSION="cu129"
+FLASHINFER_VERSION=0.5.3
 OPTIONAL_DEPS="${1:-}"
 
 # Detect system architecture
@@ -154,6 +155,8 @@ $PIP_CMD install nvidia-nvshmem-cu12==3.4.5 --force-reinstall $PIP_INSTALL_SUFFI
 $PIP_CMD install nvidia-cudnn-cu12==9.16.0.29 --force-reinstall $PIP_INSTALL_SUFFIX
 $PIP_CMD uninstall xformers || true
 
+$PIP_CMD install flashinfer-jit-cache==${FLASHINFER_VERSION} --index-url https://flashinfer.ai/whl/${CU_VERSION}
+
 # Show current packages
 $PIP_CMD list
 python3 -c "import torch; print(torch.version.cuda)"