Bump FlashInfer to 0.6.17 and remove Kimi K3 workarounds (#33997)

2026-08-12 02:17:26 -07:00
parent 2d76d537e5
commit 00e57d74f0
19 changed files with 84 additions and 6496 deletions
@@ -13,10 +13,7 @@ ARG PIP_DEFAULT_INDEX
 ARG UBUNTU_MIRROR
 ARG GITHUB_ARTIFACTORY=github.com
 ARG INSTALL_FLASHINFER_JIT_CACHE=0
-ARG FLASHINFER_VERSION=0.6.15.post1
-ARG TRTLLM_GEN_MOE_CUBIN_URL="https://github.com/sgl-project/whl/releases/download/trtllm_gen_moe_cubin_20260617/trtllm_gen_moe_cubin_pool_20260617_v0613rc1.zip"
-ARG TRTLLM_GEN_MOE_CUBIN_SHA256="4900501cbe782a76b08a5858f9f07152287b97cb68114466dac286366b66c192"
-ARG TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT="trtllm_gen_moe_cubin_pool_20260617_v0613rc1"
+ARG FLASHINFER_VERSION=0.6.17
 ARG MOONCAKE_VERSION=0.3.12.post1
 ARG MSCCLPP_VERSION=sglang-v0.9.1
 #if need other arg please add in MOONCAKE_COMPILE_ARG
@@ -25,8 +22,7 @@ ARG MOONCAKE_COMPILE_ARG="-DUSE_HTTP=ON -DUSE_MNNVL=ON -DUSE_CUDA=ON -DWITH_EP=O
 ENV DEBIAN_FRONTEND=noninteractive \
    CUDA_HOME=/usr/local/cuda \
    GDRCOPY_HOME=/usr/src/gdrdrv-${GDRCOPY_VERSION}/ \
-    FLASHINFER_VERSION=${FLASHINFER_VERSION} \
-    SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL=/opt/trtllm_gen_moe_cubin_pool
+    FLASHINFER_VERSION=${FLASHINFER_VERSION}

 # Add GKE default lib and bin locations
 ENV PATH="${PATH}:/usr/local/nvidia/bin" \
@@ -73,7 +69,6 @@ RUN --mount=type=cache,target=/var/cache/apt,id=base-apt \
    build-essential \
    cmake \
    perl \
-    patch \
    patchelf \
    ccache \
    git-lfs \
@@ -164,7 +159,6 @@ ENV LANG=en_US.UTF-8 \
 #     |
 #     +-- devtools_builder (independent)
 #     +-- gateway_builder  (independent, only needs gateway source)
-#     +-- trtllm_cubin_builder (independent)
 #     |
 #     v
 #   framework (combines all artifacts)
@@ -400,27 +394,6 @@ RUN --mount=type=cache,target=/root/.cache/pip \
    && cp target/release/sgl-model-gateway /build/sgl-model-gateway-bin \
    && rm -rf /root/.cargo /root/.rustup /build/sgl-model-gateway/target /build/sgl-model-gateway/bindings/python/target

-########################################################
-# PARALLEL STAGE 6: TRT-LLM Generated-MoE Cubin Pool
-########################################################
-FROM base AS trtllm_cubin_builder
-
-RUN cubin_archive="/tmp/trtllm_gen_moe_cubin_pool.zip" && \
-    cubin_extract_dir="/tmp/trtllm_gen_moe_cubin_extract" && \
-    wget --no-verbose --output-document="${cubin_archive}" \
-      "${TRTLLM_GEN_MOE_CUBIN_URL}" && \
-    echo "${TRTLLM_GEN_MOE_CUBIN_SHA256}  ${cubin_archive}" | \
-      sha256sum --check --strict - && \
-    mkdir -p "${cubin_extract_dir}" && \
-    unzip -q "${cubin_archive}" -d "${cubin_extract_dir}" && \
-    test ! -e "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" && \
-    mv "${cubin_extract_dir}/${TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT}" \
-      "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" && \
-    test "$(find "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" \
-      -type f -name '*.cubin' | wc -l)" -eq 1696 && \
-    rm -f "${cubin_archive}" && \
-    rm -rf "${cubin_extract_dir}"
-
 ########################################################
 ########## Final Framework Image ######################
 ########################################################
@@ -455,21 +428,6 @@ RUN --mount=type=cache,target=/root/.cache/pip \
 # Copy flashinfer cubin (always) and jit-cache (if installed) packages
 COPY --from=flashinfer_cache /flashinfer_jit_output/ /usr/local/lib/python3.12/dist-packages/

-# Apply the FlashInfer CuTeDSL MLA decode-context-parallel runtime patch.
-# Exclude tests because they are not included in the installed wheel.
-COPY docker/kimi_k3/flashinfer-perkz-dcp-0.6.15.txt /tmp/flashinfer-perkz-dcp-0.6.15.txt
-RUN FLASHINFER_DCP_PATCH=/tmp/flashinfer-perkz-dcp-0.6.15.txt && \
-    FLASHINFER_SITE_PACKAGES="$(python3 -c 'from pathlib import Path; import flashinfer; print(Path(flashinfer.__file__).resolve().parent.parent)')" && \
-    sed '/^diff --git a\/tests\//,$d' "${FLASHINFER_DCP_PATCH}" | \
-      patch --dry-run --batch --forward --strip=1 --directory="${FLASHINFER_SITE_PACKAGES}" && \
-    sed '/^diff --git a\/tests\//,$d' "${FLASHINFER_DCP_PATCH}" | \
-      patch --batch --forward --strip=1 --directory="${FLASHINFER_SITE_PACKAGES}" && \
-    rm -f "${FLASHINFER_DCP_PATCH}" && \
-    rm -rf /root/.cache/flashinfer /root/.cache/pip
-
-# Copy the pinned FlashInfer MXFP4 MoE runner cubin pool
-COPY --from=trtllm_cubin_builder /opt/trtllm_gen_moe_cubin_pool /opt/trtllm_gen_moe_cubin_pool
-
 # Copy dev tools
 COPY --from=devtools_builder /tools/diff-so-fancy /usr/local/bin/
 COPY --from=devtools_builder /tools/clang-format /usr/local/bin/
@@ -734,8 +692,7 @@ ARG GDRCOPY_VERSION=2.5.1

 ENV DEBIAN_FRONTEND=noninteractive \
    CUDA_HOME=/usr/local/cuda \
-    GDRCOPY_HOME=/usr/src/gdrdrv-${GDRCOPY_VERSION}/ \
-    SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL=/opt/trtllm_gen_moe_cubin_pool
+    GDRCOPY_HOME=/usr/src/gdrdrv-${GDRCOPY_VERSION}/

 # Add GKE default lib and bin locations + CUDA compiler paths for FlashInfer JIT
 ENV PATH="${PATH}:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/cuda/nvvm/bin" \
@@ -821,9 +778,6 @@ RUN --mount=type=cache,target=/var/cache/apt,id=runtime-apt \
 # Copy Python site-packages from framework (already cleaned of __pycache__/tests/pyc files)
 COPY --from=framework_final /usr/local/lib/python3.12/dist-packages /usr/local/lib/python3.12/dist-packages

-# Copy the pinned FlashInfer MXFP4 MoE runner cubin pool
-COPY --from=framework_final /opt/trtllm_gen_moe_cubin_pool /opt/trtllm_gen_moe_cubin_pool
-
 # Copy SGLang workspace
 COPY --from=framework_final /sgl-workspace /sgl-workspace

@@ -4,15 +4,13 @@
 # (deepseek-ai@d28bd67 at /sgl-workspace/DeepEP), the deep_gemm pip package,
 # and the CUDA 12.9 toolchain.
 #
-# This image adds the four Kimi-K3-specific pieces that stock lacks:
+# This image adds the three Kimi-K3-specific pieces that stock lacks:
 #   1. the Kimi-K3 SGLang code (this repo), editable-installed
 #   2. DeepEP patch + rebuild:
 #        topk 11->16, SWITCH_HIDDEN += 3584, EP>8 SourceMeta alignment,
 #        and cross-node timeout headroom; rebuilt for sm_90 and sm_100a only
 #   3. DeepGEMM upgrade to 0.1.5.post2:
 #        official MegaMoE runtime-JIT header with Kimi-K3 SiTU support
-#   4. FlashInfer CuTeDSL MLA DCP patch:
-#        apply the seven runtime-file diffs; exclude tests absent from the wheel
 #
 # Build (on/for x86_64; nvcc cross-compiles the DeepEP cubin, no GPU needed):
 #   docker build -f docker/kimi_k3/kimi_k3_cu12.Dockerfile \
@@ -36,9 +34,7 @@ ENV RUSTUP_HOME="/usr/local/rustup" \
 RUN apt-get update && \
    apt-get install -y --no-install-recommends \
      ca-certificates \
-      curl \
-      unzip \
-      wget && \
+      curl && \
    curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
      sh -s -- -y --no-modify-path --profile minimal \
        --default-toolchain "${RUST_VERSION}" && \
@@ -62,12 +58,10 @@ RUN set -eu; \
 # --- 1. Kimi-K3 SGLang code (replaces the base's stock sglang, editable) ---
 # Keep the installed extension modules, but discard Rust and pip build
 # artifacts that are not used at runtime.
-ARG SGLANG_COMMIT="25035bff8d34f3fcce2c1a2a5b1fe610225e84ed"
 RUN rm -rf /sgl-workspace/sglang && \
-    git clone --no-checkout \
+    git clone --branch main \
      https://github.com/sgl-project/sglang.git /sgl-workspace/sglang && \
    cd /sgl-workspace/sglang && \
-    git checkout --detach "${SGLANG_COMMIT}" && \
    rm -rf .git && \
    test ! -e .git && \
    pip install -e python --no-deps && \
@@ -96,55 +90,24 @@ RUN python3 -m pip install \
      "nvidia-nvimgcodec-cu12[all]==${NVIMGCODEC_VERSION}" && \
    rm -rf /root/.cache/pip

-# Install the pinned FlashInfer MXFP4 MoE runner cubin pool.
-ARG TRTLLM_GEN_MOE_CUBIN_URL="https://github.com/sgl-project/whl/releases/download/trtllm_gen_moe_cubin_20260617/trtllm_gen_moe_cubin_pool_20260617_v0613rc1.zip"
-ARG TRTLLM_GEN_MOE_CUBIN_SHA256="4900501cbe782a76b08a5858f9f07152287b97cb68114466dac286366b66c192"
-ARG TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT="trtllm_gen_moe_cubin_pool_20260617_v0613rc1"
-ENV SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL="/opt/trtllm_gen_moe_cubin_pool"
-
-RUN cubin_archive="/tmp/trtllm_gen_moe_cubin_pool.zip" && \
-    cubin_extract_dir="/tmp/trtllm_gen_moe_cubin_extract" && \
-    wget --no-verbose --output-document="${cubin_archive}" \
-      "${TRTLLM_GEN_MOE_CUBIN_URL}" && \
-    echo "${TRTLLM_GEN_MOE_CUBIN_SHA256}  ${cubin_archive}" | \
-      sha256sum --check --strict - && \
-    mkdir -p "${cubin_extract_dir}" && \
-    unzip -q "${cubin_archive}" -d "${cubin_extract_dir}" && \
-    test ! -e "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" && \
-    mv "${cubin_extract_dir}/${TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT}" \
-      "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" && \
-    test "$(find "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" \
-      -type f -name '*.cubin' | wc -l)" -eq 1696 && \
-    rm -f "${cubin_archive}" && \
-    rm -rf "${cubin_extract_dir}"
-
-# Reinstall the matching FlashInfer package trio before patching its Python
-# sources. A mixed Python/cubin/JIT-cache installation fails at import time.
+# Install the matching official FlashInfer package trio. A mixed
+# Python/cubin/JIT-cache installation fails at import time.
 # flashinfer-python and flashinfer-cubin are CUDA-independent packages; the
 # JIT-cache wheel is selected from the official CUDA 12.9 index.
 RUN python3 -m pip uninstall -y \
      flashinfer-python flashinfer-cubin flashinfer-jit-cache && \
    rm -rf /root/.cache/flashinfer /root/.cache/pip && \
    python3 -m pip install --no-deps \
-      "flashinfer-python==0.6.15.post1" && \
+      "flashinfer-python==0.6.17" && \
    python3 -m pip install --no-deps \
-      "flashinfer-cubin==0.6.15.post1" \
+      "flashinfer-cubin==0.6.17" \
      --index-url https://flashinfer.ai/whl && \
    python3 -m pip install --no-deps \
-      "flashinfer-jit-cache==0.6.15.post1" \
+      "flashinfer-jit-cache==0.6.17" \
      --index-url https://flashinfer.ai/whl/cu129 && \
-    python3 -c 'from importlib.metadata import version; expected = "0.6.15.post1"; assert version("flashinfer-python").split("+", 1)[0] == expected; assert version("flashinfer-cubin").split("+", 1)[0] == expected; assert version("flashinfer-jit-cache").startswith(expected + "+cu129"), version("flashinfer-jit-cache")' && \
+    python3 -c 'from importlib.metadata import version; expected = "0.6.17"; assert version("flashinfer-python").split("+", 1)[0] == expected; assert version("flashinfer-cubin").split("+", 1)[0] == expected; assert version("flashinfer-jit-cache").startswith(expected + "+cu129"), version("flashinfer-jit-cache")' && \
    rm -rf /root/.cache/pip

-ENV FLASHINFER_VERSION="0.6.15.post1"
-
-# --- 4. FlashInfer: CuTeDSL MLA decode-context-parallel runtime patch ---
-RUN FLASHINFER_DCP_PATCH=/sgl-workspace/sglang/docker/kimi_k3/flashinfer-perkz-dcp-0.6.15.txt && \
-    FLASHINFER_SITE_PACKAGES="$(python3 -c 'from pathlib import Path; import flashinfer; print(Path(flashinfer.__file__).resolve().parent.parent)')" && \
-    sed '/^diff --git a\/tests\//,$d' "${FLASHINFER_DCP_PATCH}" | \
-      patch --dry-run --batch --forward --strip=1 --directory="${FLASHINFER_SITE_PACKAGES}" && \
-    sed '/^diff --git a\/tests\//,$d' "${FLASHINFER_DCP_PATCH}" | \
-      patch --batch --forward --strip=1 --directory="${FLASHINFER_SITE_PACKAGES}" && \
-    rm -rf /root/.cache/flashinfer /root/.cache/pip
+ENV FLASHINFER_VERSION="0.6.17"

 WORKDIR /sgl-workspace/sglang
@@ -4,7 +4,7 @@
 # (deepseek-ai@d28bd67 at /sgl-workspace/DeepEP), the deep_gemm pip package,
 # and the CUDA 13 toolchain (nvcc + /usr/local/cuda/include/cccl).
 #
-# This image adds the four Kimi-K3-specific pieces that stock lacks:
+# This image adds the three Kimi-K3-specific pieces that stock lacks:
 #   1. the Kimi-K3 SGLang code (this repo), editable-installed
 #   2. DeepEP patch + rebuild:
 #        topk 11->16, SWITCH_HIDDEN += 3584, EP>8 SourceMeta alignment,
@@ -12,8 +12,6 @@
 #        sm_90, sm_100a, and sm_103a
 #   3. DeepGEMM upgrade to 0.1.5.post2:
 #        official MegaMoE runtime-JIT header with Kimi-K3 SiTU support
-#   4. FlashInfer CuTeDSL MLA DCP patch:
-#        apply the seven runtime-file diffs; exclude tests absent from the wheel
 #
 # Build (on/for aarch64; nvcc cross-compiles the DeepEP cubin, no GPU needed):
 #   docker build -f docker/kimi_k3/kimi_k3_cu13.Dockerfile \
@@ -37,9 +35,7 @@ ENV RUSTUP_HOME="/usr/local/rustup" \
 RUN apt-get update && \
    apt-get install -y --no-install-recommends \
      ca-certificates \
-      curl \
-      unzip \
-      wget && \
+      curl && \
    curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \
      sh -s -- -y --no-modify-path --profile minimal \
        --default-toolchain "${RUST_VERSION}" && \
@@ -53,12 +49,10 @@ ARG TORCH_CUDA_ARCH_LIST="9.0;10.0a;10.3a"
 # --- 1. Kimi-K3 SGLang code (replaces the base's stock sglang, editable) ---
 # Keep the installed extension modules, but discard Rust and pip build
 # artifacts that are not used at runtime.
-ARG SGLANG_COMMIT="25035bff8d34f3fcce2c1a2a5b1fe610225e84ed"
 RUN rm -rf /sgl-workspace/sglang && \
-    git clone --no-checkout \
+    git clone --branch main \
      https://github.com/sgl-project/sglang.git /sgl-workspace/sglang && \
    cd /sgl-workspace/sglang && \
-    git checkout --detach "${SGLANG_COMMIT}" && \
    rm -rf .git && \
    test ! -e .git && \
    pip install -e python --no-deps && \
@@ -86,53 +80,22 @@ RUN python3 -m pip install \
      "nvidia-nvimgcodec-cu13[all]==${NVIMGCODEC_VERSION}" && \
    rm -rf /root/.cache/pip

-# Install the pinned FlashInfer MXFP4 MoE runner cubin pool.
-ARG TRTLLM_GEN_MOE_CUBIN_URL="https://github.com/sgl-project/whl/releases/download/trtllm_gen_moe_cubin_20260617/trtllm_gen_moe_cubin_pool_20260617_v0613rc1.zip"
-ARG TRTLLM_GEN_MOE_CUBIN_SHA256="4900501cbe782a76b08a5858f9f07152287b97cb68114466dac286366b66c192"
-ARG TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT="trtllm_gen_moe_cubin_pool_20260617_v0613rc1"
-ENV SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL="/opt/trtllm_gen_moe_cubin_pool"
-
-RUN cubin_archive="/tmp/trtllm_gen_moe_cubin_pool.zip" && \
-    cubin_extract_dir="/tmp/trtllm_gen_moe_cubin_extract" && \
-    wget --no-verbose --output-document="${cubin_archive}" \
-      "${TRTLLM_GEN_MOE_CUBIN_URL}" && \
-    echo "${TRTLLM_GEN_MOE_CUBIN_SHA256}  ${cubin_archive}" | \
-      sha256sum --check --strict - && \
-    mkdir -p "${cubin_extract_dir}" && \
-    unzip -q "${cubin_archive}" -d "${cubin_extract_dir}" && \
-    test ! -e "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" && \
-    mv "${cubin_extract_dir}/${TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT}" \
-      "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" && \
-    test "$(find "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" \
-      -type f -name '*.cubin' | wc -l)" -eq 1696 && \
-    rm -f "${cubin_archive}" && \
-    rm -rf "${cubin_extract_dir}"
-
-# Reinstall the matching FlashInfer package trio before patching its Python
-# sources. A mixed Python/cubin/JIT-cache installation fails at import time.
+# Install the matching official FlashInfer package trio. A mixed
+# Python/cubin/JIT-cache installation fails at import time.
 RUN python3 -m pip uninstall -y \
      flashinfer-python flashinfer-cubin flashinfer-jit-cache && \
    rm -rf /root/.cache/flashinfer /root/.cache/pip && \
    python3 -m pip install --no-deps \
-      "flashinfer-python==0.6.15.post1" && \
+      "flashinfer-python==0.6.17" && \
    python3 -m pip install --no-deps \
-      "flashinfer-cubin==0.6.15.post1" \
+      "flashinfer-cubin==0.6.17" \
      --index-url https://flashinfer.ai/whl && \
    python3 -m pip install --no-deps \
-      "flashinfer-jit-cache==0.6.15.post1" \
+      "flashinfer-jit-cache==0.6.17" \
      --index-url https://flashinfer.ai/whl/cu130 && \
-    python3 -c 'from importlib.metadata import version; expected = "0.6.15.post1"; packages = ("flashinfer-python", "flashinfer-cubin", "flashinfer-jit-cache"); actual = {package: version(package).split("+", 1)[0] for package in packages}; assert all(value == expected for value in actual.values()), actual' && \
+    python3 -c 'from importlib.metadata import version; expected = "0.6.17"; packages = ("flashinfer-python", "flashinfer-cubin", "flashinfer-jit-cache"); actual = {package: version(package).split("+", 1)[0] for package in packages}; assert all(value == expected for value in actual.values()), actual' && \
    rm -rf /root/.cache/pip

-ENV FLASHINFER_VERSION="0.6.15.post1"
-
-# --- 4. FlashInfer: CuTeDSL MLA decode-context-parallel runtime patch ---
-RUN FLASHINFER_DCP_PATCH=/sgl-workspace/sglang/docker/kimi_k3/flashinfer-perkz-dcp-0.6.15.txt && \
-    FLASHINFER_SITE_PACKAGES="$(python3 -c 'from pathlib import Path; import flashinfer; print(Path(flashinfer.__file__).resolve().parent.parent)')" && \
-    sed '/^diff --git a\/tests\//,$d' "${FLASHINFER_DCP_PATCH}" | \
-      patch --dry-run --batch --forward --strip=1 --directory="${FLASHINFER_SITE_PACKAGES}" && \
-    sed '/^diff --git a\/tests\//,$d' "${FLASHINFER_DCP_PATCH}" | \
-      patch --batch --forward --strip=1 --directory="${FLASHINFER_SITE_PACKAGES}" && \
-    rm -rf /root/.cache/flashinfer /root/.cache/pip
+ENV FLASHINFER_VERSION="0.6.17"

 WORKDIR /sgl-workspace/sglang