diff --git a/.github/workflows/release-whl-kernel.yml b/.github/workflows/release-whl-kernel.yml index 1c4b29397..ae2393b18 100644 --- a/.github/workflows/release-whl-kernel.yml +++ b/.github/workflows/release-whl-kernel.yml @@ -46,7 +46,8 @@ jobs: chmod +x ./build.sh ./build.sh "${{ matrix.python-version }}" "${{ matrix.cuda-version }}" ${{ matrix.arch == 'aarch64' && 'aarch64' || '' }} env: - USE_CCACHE: 0 + BUILD_JOBS: 64 + NVCC_THREADS: 8 - name: Upload to PyPI working-directory: sgl-kernel @@ -139,7 +140,8 @@ jobs: chmod +x ./build.sh ./build.sh "${{ matrix.python-version }}" "${{ matrix.cuda-version }}" ${{ matrix.arch == 'aarch64' && 'aarch64' || '' }} env: - USE_CCACHE: 0 + BUILD_JOBS: 64 + NVCC_THREADS: 8 - name: Upload artifacts uses: actions/upload-artifact@v4 diff --git a/sgl-kernel/Dockerfile b/sgl-kernel/Dockerfile index e5bc1c339..50cefe427 100644 --- a/sgl-kernel/Dockerfile +++ b/sgl-kernel/Dockerfile @@ -98,6 +98,11 @@ ARG ENABLE_CMAKE_PROFILE ARG ENABLE_BUILD_PROFILE ARG ARCH=x86_64 ARG USE_CCACHE=1 +# Parallelism knobs (override via --build-arg) +# BUILD_JOBS: number of parallel compilation units (ninja -j) +# NVCC_THREADS: per-compilation-unit NVCC --threads (multi-arch PTXAS) +ARG BUILD_JOBS=0 +ARG NVCC_THREADS=32 RUN --mount=type=cache,id=sgl-kernel-ccache,target=/ccache \ --mount=type=cache,id=sgl-kernel-pip,target=/root/.cache/pip \ @@ -121,13 +126,18 @@ RUN --mount=type=cache,id=sgl-kernel-ccache,target=/ccache \ export CMAKE_BUILD_PARALLEL_LEVEL=2; \ export NINJAFLAGS="-j2"; \ echo "ARM detected: Using extra conservative settings (2 parallel jobs)"; \ + elif [ "${BUILD_JOBS}" -gt 0 ] 2>/dev/null; then \ + export CMAKE_BUILD_PARALLEL_LEVEL=${BUILD_JOBS}; \ else \ - export CMAKE_BUILD_PARALLEL_LEVEL=$(echo "$(( $(nproc) / 3 )) 48" | awk '{print ($1 < $2) ? $1 : $2}'); \ + 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}"; \ if [ -n "${ENABLE_CMAKE_PROFILE:-}" ]; then \ echo "CMake profiling enabled - will save to /sgl-kernel/cmake-profile.json"; \ - export CMAKE_ARGS="--profiling-output=/sgl-kernel/cmake-profile.json --profiling-format=google-trace"; \ + export CMAKE_ARGS="${CMAKE_ARGS} --profiling-output=/sgl-kernel/cmake-profile.json --profiling-format=google-trace"; \ fi; \ + echo "Build parallelism: CMAKE_BUILD_PARALLEL_LEVEL=${CMAKE_BUILD_PARALLEL_LEVEL}, NVCC_THREADS=${NVCC_THREADS}"; \ + echo "CMAKE_ARGS=${CMAKE_ARGS}"; \ ${PYTHON_ROOT_PATH}/bin/python -m uv build --wheel -Cbuild-dir=build . --color=always --no-build-isolation; \ ./rename_wheels.sh; \ if [ -n "${ENABLE_BUILD_PROFILE:-}" ] && [ -f /sgl-kernel/build/.ninja_log ]; then \ diff --git a/sgl-kernel/build.sh b/sgl-kernel/build.sh index b29d4d64a..a5ea4905e 100755 --- a/sgl-kernel/build.sh +++ b/sgl-kernel/build.sh @@ -20,7 +20,8 @@ fi # Using home directory to persist across workspace cleanups/checkouts CACHE_DIR="${HOME}/.cache/sgl-kernel" BUILDX_CACHE_DIR="${CACHE_DIR}/buildx" -mkdir -p "${BUILDX_CACHE_DIR}" +CCACHE_HOST_DIR="${CACHE_DIR}/ccache" +mkdir -p "${BUILDX_CACHE_DIR}" "${CCACHE_HOST_DIR}" # Ensure a buildx builder with docker-container driver (required for cache export) BUILDER_NAME="sgl-kernel-builder" @@ -45,13 +46,23 @@ echo "BASE_IMG: ${BASE_IMG}" echo "PYTHON_TAG: ${PY_TAG}" echo "Output: ${DIST_DIR}/" echo "Buildx cache: ${BUILDX_CACHE_DIR}" +echo "ccache dir: ${CCACHE_HOST_DIR}" echo "Builder: ${BUILDER_NAME}" +echo "BUILD_JOBS: ${BUILD_JOBS:-auto}" +echo "NVCC_THREADS: ${NVCC_THREADS:-32}" +echo "USE_CCACHE: ${USE_CCACHE:-1}" echo "----------------------------------------" -# Optional profiling build-args (empty string disables) +# Optional build-args (empty string disables) BUILD_ARGS=() [ -n "${ENABLE_CMAKE_PROFILE:-}" ] && BUILD_ARGS+=(--build-arg ENABLE_CMAKE_PROFILE="${ENABLE_CMAKE_PROFILE}") [ -n "${ENABLE_BUILD_PROFILE:-}" ] && BUILD_ARGS+=(--build-arg ENABLE_BUILD_PROFILE="${ENABLE_BUILD_PROFILE}") +[ -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}") + +# ---- Step 1: Build deps image (layer cached, fast on repeat) ---- +DEPS_TAG="sgl-kernel-deps:cuda${CUDA_VERSION}-${PY_TAG}-${ARCH}" docker buildx build \ --builder "${BUILDER_NAME}" \ @@ -64,8 +75,70 @@ docker buildx build \ "${BUILD_ARGS[@]}" \ --cache-from type=local,src=${BUILDX_CACHE_DIR} \ --cache-to type=local,dest=${BUILDX_CACHE_DIR},mode=max \ - --target artifact \ - --output "type=local,dest=${DIST_DIR}" \ + --target deps \ + --load \ + -t "${DEPS_TAG}" \ --network=host +echo "Deps image ready: ${DEPS_TAG}" + +# ---- Step 2: Build wheel with host-mounted ccache ---- +# This allows ccache to persist on the host filesystem across builds. +CCACHE_FLAG="${USE_CCACHE:-1}" +BUILD_JOBS_FLAG="${BUILD_JOBS:-0}" +NVCC_THREADS_FLAG="${NVCC_THREADS:-32}" + +docker run --rm \ + --network=host \ + -v "$(pwd):/sgl-kernel" \ + -v "${CCACHE_HOST_DIR}:/ccache" \ + -w /sgl-kernel \ + -e ARCH="${ARCH}" \ + "${DEPS_TAG}" \ + bash -c ' +set -eux + +USE_CCACHE='"${CCACHE_FLAG}"' +BUILD_JOBS='"${BUILD_JOBS_FLAG}"' +NVCC_THREADS='"${NVCC_THREADS_FLAG}"' + +if [ "${USE_CCACHE}" = "1" ]; then + export CCACHE_DIR=/ccache + export CCACHE_BASEDIR=/sgl-kernel + export CCACHE_MAXSIZE=10G + export CCACHE_COMPILERCHECK=content + export CCACHE_COMPRESS=true + export CCACHE_SLOPPINESS=file_macro,time_macros,include_file_mtime,include_file_ctime + export CMAKE_C_COMPILER_LAUNCHER=ccache + export CMAKE_CXX_COMPILER_LAUNCHER=ccache + export CMAKE_CUDA_COMPILER_LAUNCHER=ccache + echo "=== ccache stats (before) ===" + ccache -sV +fi + +if [ "'"${ARCH}"'" = "aarch64" ]; then + export CUDA_NVCC_FLAGS="-Xcudafe --threads=8" + export MAKEFLAGS="-j8" + export CMAKE_BUILD_PARALLEL_LEVEL=2 + export NINJAFLAGS="-j4" + echo "ARM detected: Using extra conservative settings (2 parallel jobs)" +elif [ "${BUILD_JOBS}" -gt 0 ] 2>/dev/null; then + export CMAKE_BUILD_PARALLEL_LEVEL=${BUILD_JOBS} +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}" +echo "Build parallelism: CMAKE_BUILD_PARALLEL_LEVEL=${CMAKE_BUILD_PARALLEL_LEVEL}, NVCC_THREADS=${NVCC_THREADS}" + +${PYTHON_ROOT_PATH}/bin/python -m uv build --wheel -Cbuild-dir=build . --color=always --no-build-isolation +./rename_wheels.sh + +if [ "${USE_CCACHE}" = "1" ]; then + echo "=== ccache stats (after) ===" + ccache -s +fi +' + echo "Done. Wheels are in ${DIST_DIR}/" +ls -lh "${DIST_DIR}"/*.whl 2>/dev/null || true diff --git a/sgl-kernel/csrc/elementwise/concat_mla.cu b/sgl-kernel/csrc/elementwise/concat_mla.cu index 7d5b8595c..5c157c9f8 100644 --- a/sgl-kernel/csrc/elementwise/concat_mla.cu +++ b/sgl-kernel/csrc/elementwise/concat_mla.cu @@ -215,3 +215,4 @@ void concat_mla_absorb_q(at::Tensor a, at::Tensor b, at::Tensor out) { cudaError_t err = cudaGetLastError(); TORCH_CHECK(err == cudaSuccess, "CUDA kernel launch failed: ", cudaGetErrorString(err)); } +// test-1