minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Dep] Upgrade flashinfer to 0.6.14 (#29910)

Browse Source 
Co-authored-by: Brayden Zhong <b8zhong@uwaterloo.ca>
Co-authored-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
This commit is contained in:
Baizhou Zhang
2026-07-09 17:52:33 -07:00
committed by GitHub
co-authored by Brayden Zhong Mohammad Miadh Angkad
parent b2f9a95867
commit 2c6cd1ef41
6 changed files with 48 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docker/Dockerfile
+7 -3
View File
@@ -19,7 +19,7 @@ ARG PIP_DEFAULT_INDEX
ARG UBUNTU_MIRROR
ARG GITHUB_ARTIFACTORY=github.com
ARG INSTALL_FLASHINFER_JIT_CACHE=0
ARG FLASHINFER_VERSION=0.6.12
ARG FLASHINFER_VERSION=0.6.14
ARG MOONCAKE_VERSION=0.3.11.post1
ARG MSCCLPP_VERSION=sglang-v0.9.1
#if need other arg please add in MOONCAKE_COMPILE_ARG
@@ -332,7 +332,7 @@ ARG CUDA_VERSION
ARG INSTALL_FLASHINFER_JIT_CACHE
ARG FLASHINFER_VERSION
# Stage jit-cache artifacts into /flashinfer_jit_output for clean COPY later
# Stage jit-cache/cubin artifacts into /flashinfer_jit_output for clean COPY later
RUN --mount=type=cache,target=/root/.cache/pip \
case "$CUDA_VERSION" in \
12.6.1) CUINDEX=126 ;; \
@@ -342,6 +342,10 @@ RUN --mount=type=cache,target=/root/.cache/pip \
*) echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1 ;; \
esac \
&& mkdir -p /flashinfer_jit_output \
# flashinfer-cubin is CUDA-version-agnostic, unlike jit-cache, so its index-url has no cu${CUINDEX} suffix
&& python3 -m pip install flashinfer-cubin==${FLASHINFER_VERSION} --index-url https://flashinfer.ai/whl \
&& cp -r /usr/local/lib/python3.12/dist-packages/flashinfer_cubin /flashinfer_jit_output/ \
&& cp -r /usr/local/lib/python3.12/dist-packages/flashinfer_cubin-*.dist-info /flashinfer_jit_output/ \
&& if [ "$INSTALL_FLASHINFER_JIT_CACHE" = "1" ]; then \
python3 -m pip install flashinfer-jit-cache==${FLASHINFER_VERSION} --index-url https://flashinfer.ai/whl/cu${CUINDEX} \
&& cp -r /usr/local/lib/python3.12/dist-packages/flashinfer_jit_cache /flashinfer_jit_output/ \
@@ -458,7 +462,7 @@ COPY --from=deepep_builder /build/DeepEP /sgl-workspace/DeepEP
RUN --mount=type=cache,target=/root/.cache/pip \
pip install /tmp/wheels/deepep/*.whl && rm -rf /tmp/wheels/deepep
# Copy flashinfer jit-cache package (if installed)
# Copy flashinfer cubin (always) and jit-cache (if installed) packages
COPY --from=flashinfer_cache /flashinfer_jit_output/ /usr/local/lib/python3.12/dist-packages/
# Copy dev tools
python/pyproject.toml
+1 -2
View File
@@ -31,8 +31,7 @@ dependencies = [
"einops",
"fastapi",
"flash-attn-4==4.0.0b15",
"flashinfer_cubin==0.6.12",
"flashinfer_python[cu13]==0.6.12", # keep it aligned with jit-cache version in Dockerfile
"flashinfer_python[cu13]==0.6.14", # keep it aligned with jit-cache version in Dockerfile
"gguf",
"interegular",
"IPython",
python/sglang/srt/entrypoints/engine.py
+1 -1
View File
@@ -1303,7 +1303,7 @@ def _set_envs_and_config(server_args: ServerArgs):
if server_args.attention_backend == "flashinfer":
assert_pkg_version(
"flashinfer_python",
"0.6.12",
"0.6.14",
"Please uninstall the old version and "
"reinstall the latest version by following the instructions "
"at https://docs.flashinfer.ai/installation.html.",
python/sglang/srt/utils/common.py
+1 -1
View File
@@ -1891,7 +1891,7 @@ def check_pkg_version_at_least(pkg: str, min_version: str) -> bool:
Args:
pkg: Package name (distribution name, e.g., "flashinfer-python")
min_version: Minimum version required (e.g., "0.6.12")
min_version: Minimum version required (e.g., "0.6.14")
Returns:
True if package is installed and version >= min_version, False otherwise
scripts/ci/cuda/ci_install_dependency.sh
+14 -1
View File
@@ -222,7 +222,8 @@ uninstall_stale_flashinfer() {
# - flashinfer-cubin: 150+ MB
# - flashinfer-jit-cache: 1.2+ GB
FLASHINFER_PYTHON_REQUIRED=$(grep -Po -m1 'flashinfer_python(\[[^]]+\])?==\K[0-9A-Za-z\.\-]+' python/pyproject.toml || echo "")
FLASHINFER_CUBIN_REQUIRED=$(grep -Po -m1 'flashinfer_cubin(\[[^]]+\])?==\K[0-9A-Za-z\.\-]+' python/pyproject.toml || echo "")
# flashinfer-cubin is no longer a pyproject dependency (installed explicitly below), tracks the same version as flashinfer_python
FLASHINFER_CUBIN_REQUIRED="$FLASHINFER_PYTHON_REQUIRED"
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 "")
FLASHINFER_JIT_CU_VERSION=$(pip show flashinfer-jit-cache 2>/dev/null | grep "^Version:" | awk '{print $2}' | sed -n 's/.*+//p' || echo "")
@@ -371,6 +372,17 @@ install_sglang_router() {
mark_step_done "${FUNCNAME[0]}"
}
install_flashinfer_cubin() {
if [ "$UNINSTALL_CUBIN" = false ]; then
echo "flashinfer-cubin==${FLASHINFER_CUBIN_REQUIRED} already installed, skipping install"
else
# flashinfer-cubin is CUDA-version-agnostic, unlike jit-cache, so its index-url has no cu${CU_VERSION} suffix
$PIP_CMD install "flashinfer-cubin==${FLASHINFER_CUBIN_REQUIRED}" --index-url https://flashinfer.ai/whl $PIP_INSTALL_SUFFIX
fi
mark_step_done "${FUNCNAME[0]}"
}
download_flashinfer_cache() {
UNINSTALL_JIT_CACHE="$UNINSTALL_JIT_CACHE" \
FLASHINFER_PYTHON_REQUIRED="$FLASHINFER_PYTHON_REQUIRED" \
@@ -583,6 +595,7 @@ main() {
fi
install_sglang_kernel
install_sglang_router
install_flashinfer_cubin
download_flashinfer_cache
force_reinstall_cutlass_dsl_libs_cu13
stabilize_flashinfer_jit_paths
test/registered/moe/test_cutedsl_moe.py
+24 -17
View File
@@ -577,14 +577,7 @@ class TestCuteDslV2(unittest.TestCase):
"CuteDslMoEWrapper / convert_sf_to_mma_layout not available",
)
def test_v2_cuda_graph_parity(self):
"""Verify non-graph and cuda_graph v2 wrappers produce identical results.
Also checks both match the pure-PyTorch reference, and that a second
cuda_graph pass reuses buffers deterministically (subsumes the former
cuda_graph check).
"""
test_cases = [
# (num_tokens, hidden_size, intermediate_size, num_experts, top_k)
(128, 256, 512, 256, 2),
(256, 256, 512, 256, 4),
]
@@ -627,22 +620,36 @@ class TestCuteDslV2(unittest.TestCase):
with torch.no_grad():
out_no_graph = _run_wrapper(wrapper_no_graph, tensors)
out_graph = _run_wrapper(wrapper_graph, tensors)
out_graph2 = _run_wrapper(wrapper_graph, tensors)
for _ in range(3):
_run_wrapper(wrapper_graph, tensors)
torch.cuda.synchronize()
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
graph_output = _run_wrapper(wrapper_graph, tensors)
torch.cuda.synchronize()
graph.replay()
torch.cuda.synchronize()
out_graph1 = graph_output.clone()
graph.replay()
torch.cuda.synchronize()
out_graph2 = graph_output.clone()
torch.testing.assert_close(
out_no_graph,
out_graph,
out_graph1,
atol=1e-2,
rtol=1e-2,
msg="non-graph vs cuda_graph wrapper outputs diverge",
)
torch.testing.assert_close(
out_graph,
out_graph2,
atol=1e-5,
rtol=1e-5,
msg="second cuda_graph pass should reuse buffers identically",
max_diff = (out_graph1 - out_graph2).abs().max().item()
self.assertLess(
max_diff,
0.5,
f"cuda_graph replay diverged too much: max_diff={max_diff}",
)
ref_output = _compute_reference_moe_fp4(
@@ -658,7 +665,7 @@ class TestCuteDslV2(unittest.TestCase):
fc2_input_scale=tensors["fc2_input_scale"],
)
out_f32 = out_graph.float()
out_f32 = out_graph1.float()
ref_f32 = ref_output.float()
output_scale = max(ref_f32.std().item(), 0.01)
atol = max(0.1, 3.0 * output_scale)