build(sgl-kernel): support configurable mirrors for restricted networks (#27387)

2026-06-08 08:52:26 -07:00
parent 3607cbd65a
commit b047bb3e92
5 changed files with 82 additions and 27 deletions
@@ -13,6 +13,9 @@ set(CMAKE_VERBOSE_MAKEFILE ON CACHE BOOL "ON")
 set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 set(CMAKE_SHARED_LIBRARY_PREFIX "")

+# GitHub Artifactory
+set(GITHUB_ARTIFACTORY "github.com" CACHE STRING "GitHub mirror URL")
+
 # Python
 find_package(Python COMPONENTS Interpreter Development.Module ${SKBUILD_SABI_COMPONENT} REQUIRED)

@@ -46,57 +49,61 @@ clear_cuda_arches(CMAKE_FLAG)
 # cutlass
 FetchContent_Declare(
    repo-cutlass
-    GIT_REPOSITORY https://github.com/NVIDIA/cutlass
-    GIT_TAG        57e3cfb47a2d9e0d46eb6335c3dc411498efa198
-    GIT_SHALLOW    OFF
+    URL      https://${GITHUB_ARTIFACTORY}/NVIDIA/cutlass/archive/57e3cfb47a2d9e0d46eb6335c3dc411498efa198.tar.gz
+    URL_HASH SHA256=09237099a70f80bff1dc8bb80c843a674bb4fdcb46e43cc6993e711c5ca89bb5
 )
 FetchContent_Populate(repo-cutlass)

 # fmt
 FetchContent_Declare(
    repo-fmt
-    GIT_REPOSITORY https://github.com/fmtlib/fmt
-    GIT_TAG        553ec11ec06fbe0beebfbb45f9dc3c9eabd83d28
-    GIT_SHALLOW    OFF
+    URL      https://${GITHUB_ARTIFACTORY}/fmtlib/fmt/archive/553ec11ec06fbe0beebfbb45f9dc3c9eabd83d28.tar.gz
+    URL_HASH SHA256=c314292789d28c3c3b420e75a7b2d1706f685f7fb63289128d46aeaea2c6be71
 )
 FetchContent_Populate(repo-fmt)

 # Triton kernel
 FetchContent_Declare(
    repo-triton
-    GIT_REPOSITORY "https://github.com/triton-lang/triton"
-    GIT_TAG        v3.6.0
-    GIT_SHALLOW    OFF
+    URL      https://${GITHUB_ARTIFACTORY}/triton-lang/triton/archive/v3.6.0.tar.gz
+    URL_HASH SHA256=be270ed11ca5a8fbd9d7941c5bbe9a23a9f6e2ffd372c8398346928bee464774
 )
 FetchContent_Populate(repo-triton)

 # flashinfer
 FetchContent_Declare(
    repo-flashinfer
-    GIT_REPOSITORY https://github.com/flashinfer-ai/flashinfer.git
-    GIT_TAG        bc29697ba20b7e6bdb728ded98f04788e16ee021
-    GIT_SHALLOW    OFF
+    URL      https://${GITHUB_ARTIFACTORY}/flashinfer-ai/flashinfer/archive/bc29697ba20b7e6bdb728ded98f04788e16ee021.tar.gz
+    URL_HASH SHA256=931dfd118f4b6de8c7d98702153c7c03840139170af21a07607693bd9749744d
 )
 FetchContent_Populate(repo-flashinfer)

 # flash-attention
 FetchContent_Declare(
    repo-flash-attention
-    GIT_REPOSITORY https://github.com/sgl-project/sgl-attn
-    GIT_TAG        bcf72ccc6816b36a5fae2c5a3c027604629785e0
-    GIT_SHALLOW    OFF
+    URL      https://${GITHUB_ARTIFACTORY}/sgl-project/sgl-attn/archive/bcf72ccc6816b36a5fae2c5a3c027604629785e0.tar.gz
+    URL_HASH SHA256=2110d8ca1ed9b330b9f99c1d4088be12a37e44152c39fdf90dfb6526f532ba96
 )
 FetchContent_Populate(repo-flash-attention)

 # mscclpp
 FetchContent_Declare(
    repo-mscclpp
-    GIT_REPOSITORY https://github.com/microsoft/mscclpp.git
-    GIT_TAG        51eca89d20f0cfb3764ccd764338d7b22cd486a6
-    GIT_SHALLOW    OFF
+    URL      https://${GITHUB_ARTIFACTORY}/microsoft/mscclpp/archive/51eca89d20f0cfb3764ccd764338d7b22cd486a6.tar.gz
+    URL_HASH SHA256=b064de701da5253e32f4031d16b01a61d78c369327d585ff7bfb521bbe742677
 )
 FetchContent_Populate(repo-mscclpp)

+# mscclpp's own CMakeLists.txt hardcodes a github.com FetchContent for nlohmann/json.
+# Patch it to route through GITHUB_ARTIFACTORY so it follows the same mirror as our deps.
+file(READ "${repo-mscclpp_SOURCE_DIR}/CMakeLists.txt" _mscclpp_cmakelists)
+string(REPLACE
+    "https://github.com/nlohmann/json"
+    "https://${GITHUB_ARTIFACTORY}/nlohmann/json"
+    _mscclpp_cmakelists "${_mscclpp_cmakelists}")
+file(WRITE "${repo-mscclpp_SOURCE_DIR}/CMakeLists.txt" "${_mscclpp_cmakelists}")
+unset(_mscclpp_cmakelists)
+
 # ccache option
 option(ENABLE_CCACHE "Whether to use ccache" ON)
 find_program(CCACHE_FOUND ccache)
@@ -376,6 +383,8 @@ endif()
 set(MSCCLPP_USE_CUDA ON)
 set(MSCCLPP_BYPASS_GPU_CHECK ON)
 set(MSCCLPP_BUILD_TESTS OFF)
+set(MSCCLPP_BUILD_PYTHON_BINDINGS OFF)
+set(MSCCLPP_BUILD_APPS_NCCL OFF)
 add_subdirectory(
    ${repo-mscclpp_SOURCE_DIR}
    ${CMAKE_CURRENT_BINARY_DIR}/mscclpp-build
@@ -16,8 +16,13 @@ ARG CMAKE_VERSION_MINOR=1
 ARG USE_CCACHE=1
 ARG CCACHE_VERSION=4.12.1
 ARG GITHUB_ARTIFACTORY=github.com
-ARG PYTORCH_MIRROR=download.pytorch.org
+ARG PYTORCH_INDEX_BASE=https://download.pytorch.org/whl
 ARG PIP_DEFAULT_INDEX=https://pypi.python.org/simple
+# Optional mirror for the manylinux base image's yum repos (AlmaLinux 8).
+# Pass scheme + host (and optional path prefix), e.g.
+#   --build-arg YUM_MIRROR=https://mirrors.aliyun.com
+# Empty (default) keeps upstream repo.almalinux.org.
+ARG YUM_MIRROR=

 ENV PYTHON_ROOT_PATH=/opt/python/${PYTHON_TAG}
 ENV PATH=/opt/cmake/bin:${PATH}
@@ -28,8 +33,21 @@ ENV FLASHINFER_CUDA_ARCH_LIST="8.0 8.9 9.0a 10.0a 12.0a"
 ENV CPLUS_INCLUDE_PATH=/usr/local/cuda/include/cccl${CPLUS_INCLUDE_PATH:+:${CPLUS_INCLUDE_PATH}}
 ENV C_INCLUDE_PATH=/usr/local/cuda/include/cccl${C_INCLUDE_PATH:+:${C_INCLUDE_PATH}}

-# Install build dependencies
-RUN yum install gcc gcc-c++ make wget tar numactl-devel libibverbs -y --nogpgcheck \
+RUN if [ -n "${YUM_MIRROR}" ]; then \
+      set -eux; \
+      sed -i \
+        -e 's|^mirrorlist=|#mirrorlist=|g' \
+        -e 's|^# *baseurl=https://repo.almalinux.org|baseurl='"${YUM_MIRROR}"'|g' \
+        /etc/yum.repos.d/almalinux*.repo; \
+      sed -i 's|^enabled=1|enabled=0|g' /etc/yum.repos.d/epel*.repo; \
+    fi
+
+# Install build dependencies. libzstd-devel + xxhash-devel let ccache's
+# FindZstd.cmake / FindXxhash.cmake skip their hardcoded github.com FetchContent
+# fallbacks — critical when github.com is flaky/blocked. xxhash-devel is in
+# PowerTools (disabled by default on AlmaLinux 8); libzstd-devel is in BaseOS.
+RUN yum install gcc gcc-c++ make wget tar numactl-devel libibverbs libzstd-devel -y --nogpgcheck \
+ && yum --enablerepo=powertools install xxhash-devel -y --nogpgcheck \
 && ln -sv /usr/lib64/libibverbs.so.1 /usr/lib64/libibverbs.so \
 && yum clean all && rm -rf /var/cache/yum

@@ -84,7 +102,7 @@ RUN --mount=type=cache,id=sgl-kernel-pip,target=/root/.cache/pip \
      12.8) TORCH_VER=2.11.0; CU_TAG=cu128 ;; \
      *)    TORCH_VER=2.11.0; CU_TAG=cu126 ;; \
    esac; \
-    ${PYTHON_ROOT_PATH}/bin/pip install torch==${TORCH_VER} --index-url https://${PYTORCH_MIRROR}/whl/${CU_TAG}; \
+    ${PYTHON_ROOT_PATH}/bin/pip install torch==${TORCH_VER} --index-url ${PYTORCH_INDEX_BASE}/${CU_TAG}; \
    ${PYTHON_ROOT_PATH}/bin/pip install ninja setuptools==75.0.0 wheel==0.41.0 numpy uv scikit-build-core --index-url ${PIP_DEFAULT_INDEX}

 # Build stage: copy source and build wheel
@@ -103,6 +121,8 @@ ARG USE_CCACHE=1
 #   NVCC_THREADS: per-compilation-unit NVCC --threads (multi-arch PTXAS)
 ARG BUILD_JOBS=0
 ARG NVCC_THREADS=32
+# Redeclare so CMake third-party FetchContent uses the same mirror as deps stage
+ARG GITHUB_ARTIFACTORY=github.com

 RUN --mount=type=cache,id=sgl-kernel-ccache,target=/ccache \
    --mount=type=cache,id=sgl-kernel-pip,target=/root/.cache/pip \
@@ -131,7 +151,7 @@ RUN --mount=type=cache,id=sgl-kernel-ccache,target=/ccache \
    else \
      export CMAKE_BUILD_PARALLEL_LEVEL=$(echo "$(( $(nproc) * 2 / 3 )) 64" | awk '{print ($1 < $2) ? $1 : $2}'); \
    fi; \
-    export CMAKE_ARGS="${CMAKE_ARGS:-} -DSGL_KERNEL_COMPILE_THREADS=${NVCC_THREADS}"; \
+    export CMAKE_ARGS="${CMAKE_ARGS:-} -DSGL_KERNEL_COMPILE_THREADS=${NVCC_THREADS} -DGITHUB_ARTIFACTORY=${GITHUB_ARTIFACTORY}"; \
    if [ -n "${ENABLE_CMAKE_PROFILE:-}" ]; then \
      echo "CMake profiling enabled - will save to /sgl-kernel/cmake-profile.json"; \
      export CMAKE_ARGS="${CMAKE_ARGS} --profiling-output=/sgl-kernel/cmake-profile.json --profiling-format=google-trace"; \
@@ -11,6 +11,13 @@ CMAKE_BUILD_PARALLEL_LEVEL ?= $(MAX_JOBS)
 UV_BUILD_DIR ?= build
 CMAKE_POLICY_VERSION_MINIMUM ?= 3.5

+# Optional GitHub mirror for FetchContent (e.g. internal artifactory).
+# Empty -> CMakeLists.txt default (github.com).
+GITHUB_ARTIFACTORY ?=
+ifneq ($(strip $(GITHUB_ARTIFACTORY)),)
+CMAKE_ARGS += -DGITHUB_ARTIFACTORY=$(GITHUB_ARTIFACTORY)
+endif
+
 # Show help for each target
 help: ## Show this help message
 	@echo "Available targets:"
@@ -60,6 +60,10 @@ echo "BUILD_JOBS:     ${BUILD_JOBS:-auto}"
 echo "NVCC_THREADS:   ${NVCC_THREADS:-32}"
 echo "USE_CCACHE:     ${USE_CCACHE:-1}"
 echo "RESET_BUILDER:  ${RESET_BUILDER:-0}"
+echo "GITHUB_ARTIFACTORY: ${GITHUB_ARTIFACTORY:-github.com}"
+echo "PYTORCH_INDEX_BASE: ${PYTORCH_INDEX_BASE:-https://download.pytorch.org/whl}"
+echo "PIP_DEFAULT_INDEX:  ${PIP_DEFAULT_INDEX:-https://pypi.python.org/simple}"
+echo "YUM_MIRROR:         ${YUM_MIRROR:-(upstream)}"
 echo "----------------------------------------"

 # Optional build-args (empty string disables)
@@ -69,6 +73,10 @@ BUILD_ARGS=()
 [ -n "${USE_CCACHE:-}" ]           && BUILD_ARGS+=(--build-arg USE_CCACHE="${USE_CCACHE}")
 [ -n "${BUILD_JOBS:-}" ]           && BUILD_ARGS+=(--build-arg BUILD_JOBS="${BUILD_JOBS}")
 [ -n "${NVCC_THREADS:-}" ]         && BUILD_ARGS+=(--build-arg NVCC_THREADS="${NVCC_THREADS}")
+[ -n "${GITHUB_ARTIFACTORY:-}" ]   && BUILD_ARGS+=(--build-arg GITHUB_ARTIFACTORY="${GITHUB_ARTIFACTORY}")
+[ -n "${PYTORCH_INDEX_BASE:-}" ]   && BUILD_ARGS+=(--build-arg PYTORCH_INDEX_BASE="${PYTORCH_INDEX_BASE}")
+[ -n "${PIP_DEFAULT_INDEX:-}" ]    && BUILD_ARGS+=(--build-arg PIP_DEFAULT_INDEX="${PIP_DEFAULT_INDEX}")
+[ -n "${YUM_MIRROR:-}" ]           && BUILD_ARGS+=(--build-arg YUM_MIRROR="${YUM_MIRROR}")

 # ---- Step 1: Build deps image (layer cached, fast on repeat) ----
 DEPS_TAG="sgl-kernel-deps:cuda${CUDA_VERSION}-${PY_TAG}-${ARCH}"
@@ -96,6 +104,7 @@ echo "Deps image ready: ${DEPS_TAG}"
 CCACHE_FLAG="${USE_CCACHE:-1}"
 BUILD_JOBS_FLAG="${BUILD_JOBS:-0}"
 NVCC_THREADS_FLAG="${NVCC_THREADS:-32}"
+GITHUB_ARTIFACTORY_FLAG="${GITHUB_ARTIFACTORY:-github.com}"

 docker run --rm \
  --network=host \
@@ -103,6 +112,7 @@ docker run --rm \
  -v "${CCACHE_HOST_DIR}:/ccache" \
  -w /sgl-kernel \
  -e ARCH="${ARCH}" \
+  -e GITHUB_ARTIFACTORY="${GITHUB_ARTIFACTORY_FLAG}" \
  "${DEPS_TAG}" \
  bash -c '
 set -eux
@@ -137,8 +147,9 @@ else
  export CMAKE_BUILD_PARALLEL_LEVEL=$(echo "$(( $(nproc) * 2 / 3 )) 64" | awk "{print (\$1 < \$2) ? \$1 : \$2}")
 fi

-export CMAKE_ARGS="${CMAKE_ARGS:-} -DSGL_KERNEL_COMPILE_THREADS=${NVCC_THREADS}"
+export CMAKE_ARGS="${CMAKE_ARGS:-} -DSGL_KERNEL_COMPILE_THREADS=${NVCC_THREADS} -DGITHUB_ARTIFACTORY=${GITHUB_ARTIFACTORY}"
 echo "Build parallelism: CMAKE_BUILD_PARALLEL_LEVEL=${CMAKE_BUILD_PARALLEL_LEVEL}, NVCC_THREADS=${NVCC_THREADS}"
+echo "GitHub mirror: GITHUB_ARTIFACTORY=${GITHUB_ARTIFACTORY}"

 ${PYTHON_ROOT_PATH}/bin/python -m uv build --wheel -Cbuild-dir=build . --color=always --no-build-isolation
 PYTHON=${PYTHON_ROOT_PATH}/bin/python ./rename_wheels.sh
@@ -1,12 +1,20 @@
 # flash_mla
 FetchContent_Declare(
    repo-flashmla
-    GIT_REPOSITORY https://github.com/sgl-project/FlashMLA
-    GIT_TAG df022ebafb88578eab9f0300606ee765608d8b5c
-    GIT_SHALLOW OFF
+    URL      https://${GITHUB_ARTIFACTORY}/sgl-project/FlashMLA/archive/df022ebafb88578eab9f0300606ee765608d8b5c.tar.gz
+    URL_HASH SHA256=45992d7de7d051dc897aff33156a2b7515d745fc489008adc7924ab884578d52
 )
 FetchContent_Populate(repo-flashmla)

+# flashmla submodule pin: NVIDIA/cutlass @ 147f5673d0c1c3dcf66f78d677fd647e4a020219
+FetchContent_Declare(
+    repo-flashmla-cutlass
+    URL      https://${GITHUB_ARTIFACTORY}/NVIDIA/cutlass/archive/147f5673d0c1c3dcf66f78d677fd647e4a020219.tar.gz
+    URL_HASH SHA256=9f6c53320a85b4a570975e557918cde65168cd311f081920446c238437347dc6
+    SOURCE_DIR ${repo-flashmla_SOURCE_DIR}/csrc/cutlass
+)
+FetchContent_Populate(repo-flashmla-cutlass)
+
 set(FLASHMLA_CUDA_FLAGS
    "--expt-relaxed-constexpr"
    "--expt-extended-lambda"