From aa103eab8df407e87102cd290314a9f37dca5038 Mon Sep 17 00:00:00 2001
From: ishandhanani <82981111+ishandhanani@users.noreply.github.com>
Date: Thu, 9 Apr 2026 17:34:57 -0500
Subject: [PATCH] [Docker] Optimize Dockerfile for BuildKit layer caching
 (#22160)

---
 docker/Dockerfile | 490 +++++++++++++++++++++++++++++++---------------
 1 file changed, 328 insertions(+), 162 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 29aea71aa..2287f9e15 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -152,48 +152,38 @@ ENV LANG=en_US.UTF-8 \
     LC_ALL=en_US.UTF-8
 
 ########################################################
-########## Framework Development Image ################
+########## PARALLEL BUILDER STAGES ####################
 ########################################################
+#
+# These stages run IN PARALLEL via BuildKit:
+#
+#   base
+#     |
+#     +-- torch_deps ------> deepep_builder (needs torch)
+#     |                  \-> flashinfer_cache (needs flashinfer)
+#     |
+#     +-- devtools_builder (independent)
+#     +-- gateway_builder  (independent, only needs gateway source)
+#     |
+#     v
+#   framework (combines all artifacts)
+#
 
-# Copy local source if building from local
-FROM scratch AS local_src
-COPY . /src
+########################################################
+# PARALLEL STAGE 1: Torch/Deps Builder (starts from base)
+########################################################
+FROM base AS torch_deps
 
-FROM base AS framework
-
-ARG BRANCH_TYPE
-ARG BUILD_TYPE
 ARG CUDA_VERSION
-ARG BUILD_AND_DOWNLOAD_PARALLEL
+ARG BUILD_TYPE
 ARG SGL_KERNEL_VERSION
-ARG SGL_VERSION
-ARG USE_LATEST_SGLANG
-ARG INSTALL_FLASHINFER_JIT_CACHE
-ARG FLASHINFER_VERSION
-ARG GRACE_BLACKWELL
-ARG GRACE_BLACKWELL_DEEPEP_BRANCH
-ARG DEEPEP_COMMIT
-ARG TRITON_LANG_COMMIT
 ARG GITHUB_ARTIFACTORY
 
 WORKDIR /sgl-workspace
 
-# Install SGLang
-COPY --from=local_src /src /tmp/local_src
-RUN if [ "$BRANCH_TYPE" = "local" ]; then \
-        cp -r /tmp/local_src /sgl-workspace/sglang; \
-    elif [ "$USE_LATEST_SGLANG" = "1" ]; then \
-        git clone --depth=1 https://github.com/sgl-project/sglang.git /sgl-workspace/sglang; \
-    elif [ -z "$SGL_VERSION" ]; then \
-        echo "ERROR: SGL_VERSION must be set when USE_LATEST_SGLANG=0 and BRANCH_TYPE!=local" && exit 1; \
-    else \
-        git clone --depth=1 --branch v${SGL_VERSION} https://github.com/sgl-project/sglang.git /sgl-workspace/sglang; \
-    fi \
-    && rm -rf /tmp/local_src
-
+# Install sgl-kernel (from pre-built wheel)
 RUN --mount=type=cache,target=/root/.cache/pip \
     python3 -m pip install --upgrade pip setuptools wheel html5lib six \
-    && cd sglang \
     && case "$CUDA_VERSION" in \
         12.6.1) CUINDEX=126 ;; \
         12.8.1) CUINDEX=128 ;; \
@@ -213,25 +203,50 @@ RUN --mount=type=cache,target=/root/.cache/pip \
     else \
         echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1 \
     ; \
-    fi \
-    && python3 -m pip install -e "python[${BUILD_TYPE}]" --extra-index-url https://download.pytorch.org/whl/cu${CUINDEX} \
-    && if [ "$INSTALL_FLASHINFER_JIT_CACHE" = "1" ]; then \
-        python3 -m pip install flashinfer-jit-cache==${FLASHINFER_VERSION} --index-url https://flashinfer.ai/whl/cu${CUINDEX} ; \
-    fi \
-    && kernels lock python \
-    && ( success=0; for i in 1 2 3; do \
-        echo "Attempt $i/3: downloading sgl-kernel cubins..." && \
-        kernels download python && \
-        success=1 && break; \
-        echo "sgl-kernel cubin download failed, retrying in 30s..." && sleep 30; \
-    done; [ "$success" = "1" ] ) \
-    && mv python/kernels.lock /root/.cache/sglang
+    fi
 
-# DeepEP
-# We use Tom's DeepEP fork for GB200 for now; the 1fd57b0276311d035d16176bb0076426166e52f3 commit is https://github.com/fzyzcjy/DeepEP/tree/gb200_blog_part_2
-# TODO: move from Tom's branch to DeepEP hybrid-ep branch
-# We use the nvshmem version that ships with torch 2.9.1
-# CU12 uses 3.3.20 and CU13 uses 3.3.24
+# Copy ONLY dependency specification file (for better caching)
+COPY python/pyproject.toml /tmp/sglang_deps/python/pyproject.toml
+
+# Install sglang dependencies (torch, transformers, etc.)
+# This layer is cached unless pyproject.toml changes
+# Generate constraints.txt to prevent reinstalling these deps in later stages
+RUN --mount=type=cache,target=/root/.cache/pip \
+    case "$CUDA_VERSION" in \
+        12.6.1) CUINDEX=126 ;; \
+        12.8.1) CUINDEX=128 ;; \
+        12.9.1) CUINDEX=129 ;; \
+        13.0.1) CUINDEX=130 ;; \
+        *) echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1 ;; \
+    esac \
+    && cd /tmp/sglang_deps/python \
+    && mkdir -p sglang \
+    && touch sglang/__init__.py \
+    && echo '__version__ = "0.0.0"' > sglang/version.py \
+    && touch README.md \
+    && touch LICENSE \
+    && python3 -m pip install --extra-index-url https://download.pytorch.org/whl/cu${CUINDEX} ".[${BUILD_TYPE}]" \
+    && cd /sgl-workspace \
+    && rm -rf /tmp/sglang_deps \
+    && pip freeze | grep -v "^sglang==" > /sgl-workspace/constraints.txt
+
+########################################################
+# PARALLEL STAGE 2: DeepEP Builder (needs torch_deps)
+########################################################
+FROM torch_deps AS deepep_builder
+
+ARG CUDA_VERSION
+ARG BUILD_AND_DOWNLOAD_PARALLEL
+ARG GRACE_BLACKWELL
+ARG GRACE_BLACKWELL_DEEPEP_BRANCH
+ARG HOPPER_SBO
+ARG HOPPER_SBO_DEEPEP_COMMIT
+ARG DEEPEP_COMMIT
+ARG GITHUB_ARTIFACTORY
+
+WORKDIR /build
+
+# Clone DeepEP
 RUN set -eux; \
     if [ "$GRACE_BLACKWELL" = "1" ]; then \
       git clone https://github.com/fzyzcjy/DeepEP.git && \
@@ -256,19 +271,17 @@ RUN set -eux; \
         cd .. ; \
     fi
 
-# Install DeepEP
+# Build DeepEP wheel
 RUN --mount=type=cache,target=/root/.cache/pip \
-    cd /sgl-workspace/DeepEP && \
+    cd /build/DeepEP && \
     case "$CUDA_VERSION" in \
         12.6.1) \
             CHOSEN_TORCH_CUDA_ARCH_LIST='9.0' \
             ;; \
         12.8.1) \
-            # FIXED: 12.8.1 does NOT support Blackwell 10.3 \
             CHOSEN_TORCH_CUDA_ARCH_LIST='9.0;10.0' \
             ;; \
         12.9.1|13.0.1) \
-            # 12.9.1+ properly supports Blackwell 10.3 \
             CHOSEN_TORCH_CUDA_ARCH_LIST='9.0;10.0;10.3' \
             ;; \
         *) \
@@ -278,76 +291,163 @@ RUN --mount=type=cache,target=/root/.cache/pip \
     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}" MAX_JOBS=${BUILD_AND_DOWNLOAD_PARALLEL} pip install --no-build-isolation .
+    TORCH_CUDA_ARCH_LIST="${CHOSEN_TORCH_CUDA_ARCH_LIST}" MAX_JOBS=${BUILD_AND_DOWNLOAD_PARALLEL} \
+        python3 setup.py bdist_wheel -d /wheels
 
-# Install Mooncake
-RUN --mount=type=cache,target=/root/.cache/pip \
-    CUDA_MAJOR="${CUDA_VERSION%%.*}" && \
-    if [ "$CUDA_MAJOR" -ge 13 ]; then \
-        echo "CUDA >= 13, installing mooncake-transfer-engine from source code"; \
-        git clone --branch v${MOONCAKE_VERSION} --depth 1 https://github.com/kvcache-ai/Mooncake.git && \
-        cd Mooncake && \
-        bash dependencies.sh && \
-        mkdir -p build && \
-        cd build && \
-        cmake .. ${MOONCAKE_COMPILE_ARG} && \
-        make -j$(nproc) && \
-        make install; \
-    else \
-        echo "CUDA < 13, installing mooncake-transfer-engine from pip"; \
-        python3 -m pip install mooncake-transfer-engine==${MOONCAKE_VERSION}; \
-    fi
-# Install essential Python packages
-RUN --mount=type=cache,target=/root/.cache/pip \
-    python3 -m pip install \
-    datamodel_code_generator \
-    pre-commit \
-    pytest \
-    black \
-    isort \
-    icdiff \
-    uv \
-    wheel \
-    scikit-build-core \
-    nixl \
-    py-spy \
-    cubloaty \
-    google-cloud-storage
+########################################################
+# PARALLEL STAGE 3: FlashInfer Cache (needs torch_deps)
+########################################################
+FROM torch_deps AS flashinfer_cache
 
-# Build and install sgl-model-gateway (install Rust, build, then remove Rust toolchain)
-# Cleanup runs unconditionally via trap to ensure Rust artifacts don't bloat the layer
+ARG CUDA_VERSION
+ARG INSTALL_FLASHINFER_JIT_CACHE
+ARG FLASHINFER_VERSION
+ARG BUILD_AND_DOWNLOAD_PARALLEL
+
+# Download flashinfer cubins
+# Stage jit-cache artifacts into /flashinfer_jit_output for clean COPY later
 RUN --mount=type=cache,target=/root/.cache/pip \
-    cleanup() { rm -rf /root/.cargo /root/.rustup /sgl-workspace/sglang/sgl-model-gateway/target /sgl-workspace/sglang/sgl-model-gateway/bindings/python/target /sgl-workspace/sglang/sgl-model-gateway/bindings/python/dist; sed -i '/\.cargo\/env/d' /root/.profile /root/.bashrc 2>/dev/null; } \
-    && trap cleanup EXIT \
-    && curl --proto '=https' --tlsv1.2 --retry 3 --retry-delay 2 -sSf https://sh.rustup.rs | sh -s -- -y \
+    case "$CUDA_VERSION" in \
+        12.6.1) CUINDEX=126 ;; \
+        12.8.1) CUINDEX=128 ;; \
+        12.9.1) CUINDEX=129 ;; \
+        13.0.1) CUINDEX=130 ;; \
+        *) echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1 ;; \
+    esac \
+    && mkdir -p /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/ \
+        && cp -r /usr/local/lib/python3.12/dist-packages/flashinfer_jit_cache-*.dist-info /flashinfer_jit_output/ ; \
+    fi \
+    && FLASHINFER_CUBIN_DOWNLOAD_THREADS=${BUILD_AND_DOWNLOAD_PARALLEL} FLASHINFER_LOGGING_LEVEL=warning python3 -m flashinfer --download-cubin
+
+########################################################
+# PARALLEL STAGE 4: Dev Tools Builder (starts from base)
+########################################################
+FROM base AS devtools_builder
+
+ARG GITHUB_ARTIFACTORY
+
+WORKDIR /tools
+
+# Minimal apt deps needed for oh-my-zsh install in this stage
+# Full dev apt packages (gdb, vim, tmux, nsight, etc.) are installed in the framework stage
+RUN --mount=type=cache,target=/var/cache/apt,id=devtools-apt \
+    apt-get update && apt-get install -y --no-install-recommends zsh git \
+    && rm -rf /var/lib/apt/lists/*
+
+# Download CLI tools (each in its own layer for parallel downloads)
+RUN curl --retry 3 --retry-delay 2 -LSso /tools/diff-so-fancy \
+        https://${GITHUB_ARTIFACTORY}/so-fancy/diff-so-fancy/releases/download/v1.4.4/diff-so-fancy \
+    && chmod +x /tools/diff-so-fancy
+
+RUN curl --retry 3 --retry-delay 2 -LSso /tools/clang-format \
+        https://${GITHUB_ARTIFACTORY}/muttleyxd/clang-tools-static-binaries/releases/download/master-32d3ac78/clang-format-16_linux-amd64 \
+    && chmod +x /tools/clang-format
+
+RUN curl --retry 3 --retry-delay 2 -fsSL -o /tmp/clangd.zip \
+        https://${GITHUB_ARTIFACTORY}/clangd/clangd/releases/download/18.1.3/clangd-linux-18.1.3.zip \
+    && unzip -q /tmp/clangd.zip -d /tmp \
+    && cp /tmp/clangd_18.1.3/bin/* /tools/ \
+    && mkdir -p /tools/lib && cp -r /tmp/clangd_18.1.3/lib/* /tools/lib/ \
+    && rm -rf /tmp/clangd.zip /tmp/clangd_18.1.3
+
+RUN CMAKE_VERSION=3.31.1 \
+    && ARCH=$(uname -m) \
+    && CMAKE_INSTALLER="cmake-${CMAKE_VERSION}-linux-${ARCH}" \
+    && curl --retry 3 --retry-delay 2 -fsSL -o "/tmp/${CMAKE_INSTALLER}.tar.gz" \
+        "https://${GITHUB_ARTIFACTORY}/Kitware/CMake/releases/download/v${CMAKE_VERSION}/${CMAKE_INSTALLER}.tar.gz" \
+    && tar -xzf "/tmp/${CMAKE_INSTALLER}.tar.gz" -C /tmp \
+    && cp -r "/tmp/${CMAKE_INSTALLER}/bin/"* /tools/ \
+    && mkdir -p /tools/share && cp -r "/tmp/${CMAKE_INSTALLER}/share/"* /tools/share/ \
+    && rm -rf "/tmp/${CMAKE_INSTALLER}" "/tmp/${CMAKE_INSTALLER}.tar.gz"
+
+RUN curl --proto '=https' --tlsv1.2 --retry 3 --retry-delay 2 -sSf https://just.systems/install.sh | \
+    sed "s|https://github.com|https://${GITHUB_ARTIFACTORY}|g" | \
+    bash -s -- --tag 1.42.4 --to /tools
+
+# Install oh-my-zsh and plugins
+RUN sh -c "$(curl --retry 3 --retry-delay 2 -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)" "" --unattended \
+    && git clone --depth 1 https://github.com/zsh-users/zsh-autosuggestions ${ZSH_CUSTOM:-/root/.oh-my-zsh/custom}/plugins/zsh-autosuggestions \
+    && git clone --depth 1 https://github.com/zsh-users/zsh-syntax-highlighting.git ${ZSH_CUSTOM:-/root/.oh-my-zsh/custom}/plugins/zsh-syntax-highlighting
+
+########################################################
+# PARALLEL STAGE 5: Gateway Builder (starts from base)
+########################################################
+# Builds sgl-model-gateway in isolation so Python-only changes
+# don't trigger a full Rust recompilation.
+FROM base AS gateway_builder
+
+ARG GITHUB_ARTIFACTORY
+ARG BRANCH_TYPE
+ARG SGL_VERSION
+ARG USE_LATEST_SGLANG
+
+WORKDIR /build
+
+# Copy ONLY the gateway source (not the full repo)
+COPY sgl-model-gateway /build/sgl-model-gateway
+
+# Install Rust, build gateway binary and Python bindings, then clean up Rust toolchain
+RUN --mount=type=cache,target=/root/.cache/pip \
+    curl --proto '=https' --tlsv1.2 --retry 3 --retry-delay 2 -sSf https://sh.rustup.rs | sh -s -- -y \
     && export PATH="/root/.cargo/bin:${PATH}" \
-    && rustc --version && cargo --version \
     && python3 -m pip install maturin \
-    && cd /sgl-workspace/sglang/sgl-model-gateway/bindings/python \
-    && ulimit -n 65536 && maturin build --release --features vendored-openssl --out dist \
-    && python3 -m pip install --force-reinstall dist/*.whl \
-    && cd /sgl-workspace/sglang/sgl-model-gateway \
+    && cd /build/sgl-model-gateway/bindings/python \
+    && ulimit -n 65536 && maturin build --release --features vendored-openssl --out /build/gateway_wheels \
+    && cd /build/sgl-model-gateway \
     && cargo build --release --bin sgl-model-gateway --features vendored-openssl \
-    && cp target/release/sgl-model-gateway /usr/local/bin/sgl-model-gateway
+    && 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
 
+########################################################
+########## Final Framework Image ######################
+########################################################
+#
+# Combines all artifacts from parallel builder stages
+#
+FROM torch_deps AS framework
+
+ARG BRANCH_TYPE
+ARG BUILD_TYPE
+ARG CUDA_VERSION
+ARG BUILD_AND_DOWNLOAD_PARALLEL
+ARG SGL_VERSION
+ARG USE_LATEST_SGLANG
+ARG GITHUB_ARTIFACTORY
+ARG MOONCAKE_VERSION
+ARG MOONCAKE_COMPILE_ARG
+
+WORKDIR /sgl-workspace
+
+# =============================================================================
+# Copy artifacts from parallel builders
+# =============================================================================
+
+# Copy DeepEP wheel and install
+COPY --from=deepep_builder /wheels /tmp/wheels/deepep
+COPY --from=deepep_builder /build/DeepEP /sgl-workspace/DeepEP
 RUN --mount=type=cache,target=/root/.cache/pip \
-   python3 -m pip install "nvidia-cutlass-dsl>=4.4.1" "nvidia-cutlass-dsl-libs-base>=4.4.1" --force-reinstall --no-deps;
+    pip install /tmp/wheels/deepep/*.whl && rm -rf /tmp/wheels/deepep
 
-# Patching packages for CUDA 12/13 compatibility
-# TODO: Remove when torch version covers these packages
-RUN --mount=type=cache,target=/root/.cache/pip if [ "${CUDA_VERSION%%.*}" = "12" ]; then \
-    python3 -m pip install nvidia-nccl-cu12==2.28.3 --force-reinstall --no-deps ; \
-    python3 -m pip install nvidia-cudnn-cu12==9.16.0.29 --force-reinstall --no-deps ; \
-    python3 -m pip install cuda-python==12.9 ; \
-elif [ "${CUDA_VERSION%%.*}" = "13" ]; then \
-    python3 -m pip install nvidia-nccl-cu13==2.28.3 --force-reinstall --no-deps ; \
-    python3 -m pip install nvidia-cudnn-cu13==9.16.0.29 --force-reinstall --no-deps ; \
-    python3 -m pip install nvidia-cublas==13.1.0.3 --force-reinstall --no-deps ; \
-    python3 -m pip install nixl-cu13 --no-deps ; \
-    python3 -m pip install cuda-python==13.2.0 ; \
-fi
+# Copy flashinfer cache and jit-cache package (if installed)
+COPY --from=flashinfer_cache /root/.cache/flashinfer /root/.cache/flashinfer
+COPY --from=flashinfer_cache /flashinfer_jit_output/ /usr/local/lib/python3.12/dist-packages/
 
-# Install development tools
+# Copy dev tools
+COPY --from=devtools_builder /tools/diff-so-fancy /usr/local/bin/
+COPY --from=devtools_builder /tools/clang-format /usr/local/bin/
+COPY --from=devtools_builder /tools/clangd /usr/local/bin/
+COPY --from=devtools_builder /tools/lib /usr/local/lib/
+COPY --from=devtools_builder /tools/cmake /usr/local/bin/
+COPY --from=devtools_builder /tools/ctest /usr/local/bin/
+COPY --from=devtools_builder /tools/cpack /usr/local/bin/
+COPY --from=devtools_builder /tools/share/cmake-3.31 /usr/local/share/cmake-3.31
+COPY --from=devtools_builder /tools/just /usr/local/bin/
+COPY --from=devtools_builder /root/.oh-my-zsh /root/.oh-my-zsh
+
+# Install dev apt packages (need to re-run since we're in a different stage)
 RUN --mount=type=cache,target=/var/cache/apt,id=framework-apt \
     apt-get update && apt-get install -y --no-install-recommends \
     gdb \
@@ -386,63 +486,69 @@ RUN --mount=type=cache,target=/var/cache/apt,id=framework-apt \
     && apt install -y --no-install-recommends nsight-systems-cli \
     && rm -rf /var/lib/apt/lists/*
 
-# Install minimal Python dev packages
+# =============================================================================
+# Python packages and tools (before source copy for better caching)
+# =============================================================================
+
+# Install Mooncake
 RUN --mount=type=cache,target=/root/.cache/pip \
-    python3 -m pip install --break-system-packages \
+    CUDA_MAJOR="${CUDA_VERSION%%.*}" && \
+    if [ "$CUDA_MAJOR" -ge 13 ]; then \
+        echo "CUDA >= 13, installing mooncake-transfer-engine from source code"; \
+        git clone --branch v${MOONCAKE_VERSION} --depth 1 https://github.com/kvcache-ai/Mooncake.git && \
+        cd Mooncake && \
+        bash dependencies.sh && \
+        mkdir -p build && \
+        cd build && \
+        cmake .. ${MOONCAKE_COMPILE_ARG} && \
+        make -j$(nproc) && \
+        make install; \
+    else \
+        echo "CUDA < 13, installing mooncake-transfer-engine from pip"; \
+        python3 -m pip install mooncake-transfer-engine==${MOONCAKE_VERSION}; \
+    fi
+
+# Install essential Python packages (use constraints to prevent conflicts)
+RUN --mount=type=cache,target=/root/.cache/pip \
+    python3 -m pip install -c /sgl-workspace/constraints.txt \
+    datamodel_code_generator \
+    pre-commit \
     pytest \
     black \
     isort \
     icdiff \
-    scikit-build-core \
     uv \
-    pre-commit \
+    wheel \
+    scikit-build-core \
+    nixl \
+    py-spy \
+    cubloaty \
+    google-cloud-storage \
     pandas \
     matplotlib \
     tabulate \
     termplotlib
 
-# diff-so-fancy
-RUN curl --retry 3 --retry-delay 2 -LSso /usr/local/bin/diff-so-fancy \
-        https://${GITHUB_ARTIFACTORY}/so-fancy/diff-so-fancy/releases/download/v1.4.4/diff-so-fancy \
-    && chmod +x /usr/local/bin/diff-so-fancy
+RUN --mount=type=cache,target=/root/.cache/pip \
+   python3 -m pip install "nvidia-cutlass-dsl>=4.4.1" "nvidia-cutlass-dsl-libs-base>=4.4.1" --force-reinstall --no-deps;
 
-# clang-format
-RUN curl --retry 3 --retry-delay 2 -LSso /usr/local/bin/clang-format \
-        https://${GITHUB_ARTIFACTORY}/muttleyxd/clang-tools-static-binaries/releases/download/master-32d3ac78/clang-format-16_linux-amd64 \
-    && chmod +x /usr/local/bin/clang-format
-
-# clangd
-RUN curl --retry 3 --retry-delay 2 -fsSL -o clangd.zip \
-        https://${GITHUB_ARTIFACTORY}/clangd/clangd/releases/download/18.1.3/clangd-linux-18.1.3.zip \
-    && unzip -q clangd.zip \
-    && cp -r clangd_18.1.3/bin/* /usr/local/bin/ \
-    && cp -r clangd_18.1.3/lib/* /usr/local/lib/ \
-    && rm -rf clangd_18.1.3 clangd.zip
-
-# CMake
-RUN CMAKE_VERSION=3.31.1 \
-    && ARCH=$(uname -m) \
-    && CMAKE_INSTALLER="cmake-${CMAKE_VERSION}-linux-${ARCH}" \
-    && curl --retry 3 --retry-delay 2 -fsSL -o "${CMAKE_INSTALLER}.tar.gz" \
-        "https://${GITHUB_ARTIFACTORY}/Kitware/CMake/releases/download/v${CMAKE_VERSION}/${CMAKE_INSTALLER}.tar.gz" \
-    && tar -xzf "${CMAKE_INSTALLER}.tar.gz" \
-    && cp -r "${CMAKE_INSTALLER}/bin/"* /usr/local/bin/ \
-    && cp -r "${CMAKE_INSTALLER}/share/"* /usr/local/share/ \
-    && rm -rf "${CMAKE_INSTALLER}" "${CMAKE_INSTALLER}.tar.gz"
-
-# Install just
-RUN curl --proto '=https' --tlsv1.2 --retry 3 --retry-delay 2 -sSf https://just.systems/install.sh | \
-    sed "s|https://github.com|https://${GITHUB_ARTIFACTORY}|g" | \
-    bash -s -- --tag 1.42.4 --to /usr/local/bin
+# Patching packages for CUDA 12/13 compatibility
+# TODO: Remove when torch version covers these packages
+RUN --mount=type=cache,target=/root/.cache/pip if [ "${CUDA_VERSION%%.*}" = "12" ]; then \
+    python3 -m pip install nvidia-nccl-cu12==2.28.3 --force-reinstall --no-deps ; \
+    python3 -m pip install nvidia-cudnn-cu12==9.16.0.29 --force-reinstall --no-deps ; \
+    python3 -m pip install cuda-python==12.9 ; \
+elif [ "${CUDA_VERSION%%.*}" = "13" ]; then \
+    python3 -m pip install nvidia-nccl-cu13==2.28.3 --force-reinstall --no-deps ; \
+    python3 -m pip install nvidia-cudnn-cu13==9.16.0.29 --force-reinstall --no-deps ; \
+    python3 -m pip install nvidia-cublas==13.1.0.3 --force-reinstall --no-deps ; \
+    python3 -m pip install nixl-cu13 --no-deps ; \
+    python3 -m pip install cuda-python==13.2.0 ; \
+fi
 
 # Add yank script
 COPY --chown=root:root --chmod=755 docker/configs/yank /usr/local/bin/yank
 
-# Install oh-my-zsh and plugins
-RUN sh -c "$(curl --retry 3 --retry-delay 2 -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)" "" --unattended \
-    && git clone --depth 1 https://github.com/zsh-users/zsh-autosuggestions ${ZSH_CUSTOM:-~/.oh-my-zsh/custom}/plugins/zsh-autosuggestions \
-    && git clone --depth 1 https://github.com/zsh-users/zsh-syntax-highlighting.git ${ZSH_CUSTOM:-~/.oh-my-zsh/custom}/plugins/zsh-syntax-highlighting
-
 # These configs are optional; users can override them by mounting their own files
 COPY docker/configs/opt/.vimrc /opt/sglang/.vimrc
 COPY docker/configs/opt/.tmux.conf /opt/sglang/.tmux.conf
@@ -476,6 +582,60 @@ RUN --mount=type=cache,target=/var/cache/apt,id=framework-apt \
     dpkg dpkg-dev libdpkg-perl \
     && rm -rf /var/lib/apt/lists/*
 
+# =============================================================================
+# Copy sglang source and do editable install (LAST for better caching)
+# =============================================================================
+
+# Copy local source if building from local
+FROM scratch AS local_src
+COPY . /src
+
+FROM framework AS framework_final
+
+ARG BRANCH_TYPE
+ARG BUILD_TYPE
+ARG CUDA_VERSION
+ARG SGL_VERSION
+ARG USE_LATEST_SGLANG
+
+WORKDIR /sgl-workspace
+
+COPY --from=local_src /src /tmp/local_src
+RUN if [ "$BRANCH_TYPE" = "local" ]; then \
+        cp -r /tmp/local_src /sgl-workspace/sglang; \
+    elif [ "$USE_LATEST_SGLANG" = "1" ]; then \
+        git clone --depth=1 https://github.com/sgl-project/sglang.git /sgl-workspace/sglang; \
+    elif [ -z "$SGL_VERSION" ]; then \
+        echo "ERROR: SGL_VERSION must be set when USE_LATEST_SGLANG=0 and BRANCH_TYPE!=local" && exit 1; \
+    else \
+        git clone --depth=1 --branch v${SGL_VERSION} https://github.com/sgl-project/sglang.git /sgl-workspace/sglang; \
+    fi \
+    && rm -rf /tmp/local_src
+
+# Editable install (fast - dependencies already installed via constraints)
+# Clean up __pycache__/tests/pyc in same RUN to avoid writing ~28k files to layer
+RUN --mount=type=cache,target=/root/.cache/pip \
+    cd /sgl-workspace/sglang \
+    && python3 -m pip install --no-deps -e "python[${BUILD_TYPE}]" \
+    && kernels lock python \
+    && ( success=0; for i in 1 2 3; do \
+        echo "Attempt $i/3: downloading sgl-kernel cubins..." && \
+        kernels download python && \
+        success=1 && break; \
+        echo "sgl-kernel cubin download failed, retrying in 30s..." && sleep 30; \
+    done; [ "$success" = "1" ] ) \
+    && mv python/kernels.lock /root/.cache/sglang \
+    && find /usr/local/lib/python3.12/dist-packages -type d -name "__pycache__" -exec rm -rf {} + 2>/dev/null || true \
+    && find /usr/local/lib/python3.12/dist-packages -type d -name "tests" -exec rm -rf {} + 2>/dev/null || true \
+    && find /usr/local/lib/python3.12/dist-packages -name "*.pyc" -delete 2>/dev/null || true
+
+# Install pre-built gateway artifacts from parallel builder
+COPY --from=gateway_builder /build/sgl-model-gateway-bin /usr/local/bin/sgl-model-gateway
+COPY --from=gateway_builder /build/gateway_wheels /tmp/gateway_wheels
+RUN --mount=type=cache,target=/root/.cache/pip \
+    python3 -m pip install --force-reinstall /tmp/gateway_wheels/*.whl \
+    && rm -rf /tmp/gateway_wheels
+
 # Set workspace directory
 WORKDIR /sgl-workspace/sglang
 
@@ -590,15 +750,21 @@ RUN --mount=type=cache,target=/var/cache/apt,id=runtime-apt \
     dpkg dpkg-dev libdpkg-perl \
     && rm -rf /var/lib/apt/lists/*
 
-# Copy Python site-packages from framework (contains all built packages)
-COPY --from=framework /usr/local/lib/python3.12/dist-packages /usr/local/lib/python3.12/dist-packages
+# 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 SGLang workspace
-COPY --from=framework /sgl-workspace /sgl-workspace
+COPY --from=framework_final /sgl-workspace /sgl-workspace
+
+# Copy sgl-model-gateway binary
+COPY --from=framework_final /usr/local/bin/sgl-model-gateway /usr/local/bin/sgl-model-gateway
+
+# Copy py-spy binary
+COPY --from=framework_final /usr/local/bin/py-spy /usr/local/bin/py-spy
 
 # Copy cache for kernels from kernels community
-COPY --from=framework /root/.cache/huggingface /root/.cache/huggingface
-COPY --from=framework /root/.cache/sglang /root/.cache/sglang
+COPY --from=framework_final /root/.cache/huggingface /root/.cache/huggingface
+COPY --from=framework_final /root/.cache/sglang /root/.cache/sglang
 
 # Fix Triton to use system ptxas for Blackwell (sm_103a) support (CUDA 13+ only)
 RUN if [ "${CUDA_VERSION%%.*}" = "13" ] && [ -d /usr/local/lib/python3.12/dist-packages/triton/backends/nvidia/bin ]; then \
@@ -607,9 +773,9 @@ RUN if [ "${CUDA_VERSION%%.*}" = "13" ] && [ -d /usr/local/lib/python3.12/dist-p
     fi
 
 # Copy GDRCopy runtime libraries (but not the build artifacts)
-COPY --from=framework /usr/lib/libgdrapi.so* /usr/lib/
-COPY --from=framework /usr/bin/gdrcopy_* /usr/bin/
-COPY --from=framework /usr/src/gdrdrv-2.5.1 /usr/src/gdrdrv-2.5.1
+COPY --from=framework_final /usr/lib/libgdrapi.so* /usr/lib/
+COPY --from=framework_final /usr/bin/gdrcopy_* /usr/bin/
+COPY --from=framework_final /usr/src/gdrdrv-2.5.1 /usr/src/gdrdrv-2.5.1
 
 # Fix DeepEP IBGDA symlink in runtime
 RUN ln -sf /usr/lib/$(uname -m)-linux-gnu/libmlx5.so.1 /usr/lib/$(uname -m)-linux-gnu/libmlx5.so