speed up sgl-kernel build (#18586)

.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

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 \

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 \
+  --target deps \
-  --output "type=local,dest=${DIST_DIR}" \
+  --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

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