Add release workflow for sgl-deep-gemm wheels (#24348)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-04 00:58:44 -07:00
parent 52b4609789
commit de08c804c5
5 changed files with 408 additions and 0 deletions
@@ -0,0 +1,66 @@
+#!/bin/bash
+# Build sgl-deep-gemm wheel inside a CUDA-versioned container.
+#
+# Usage: build_sgl_deep_gemm.sh <PYTHON_VERSION> <CUDA_VERSION> <DEEPGEMM_SRC> [ARCH]
+#   PYTHON_VERSION: e.g. 3.10
+#   CUDA_VERSION:   e.g. 12.9 or 13.0
+#   DEEPGEMM_SRC:   path to a checkout of sgl-project/DeepGEMM
+#   ARCH:           x86_64 (default) or aarch64
+#
+# The wheel is written into <DEEPGEMM_SRC>/dist by build_sgl_deep_gemm.sh.
+set -ex
+
+if [ $# -lt 3 ]; then
+  echo "Usage: $0 <PYTHON_VERSION> <CUDA_VERSION> <DEEPGEMM_SRC> [ARCH]"
+  exit 1
+fi
+
+PYTHON_VERSION="$1"
+CUDA_VERSION="$2"
+DEEPGEMM_SRC="$(cd "$3" && pwd)"
+ARCH="${4:-$(uname -i)}"
+
+if [ "${ARCH}" = "aarch64" ]; then
+  BASE_IMG="pytorch/manylinuxaarch64-builder"
+else
+  BASE_IMG="pytorch/manylinux2_28-builder"
+fi
+
+PY_TAG="cp${PYTHON_VERSION//.}-cp${PYTHON_VERSION//.}"
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
+DOCKERFILE="${REPO_ROOT}/docker/sgl-deep-gemm.Dockerfile"
+
+DEPS_TAG="sgl-deep-gemm-deps:cuda${CUDA_VERSION}-${PY_TAG}-${ARCH}"
+
+echo "----------------------------------------"
+echo "PYTHON_VERSION: ${PYTHON_VERSION}"
+echo "CUDA_VERSION:   ${CUDA_VERSION}"
+echo "ARCH:           ${ARCH}"
+echo "BASE_IMG:       ${BASE_IMG}"
+echo "DEEPGEMM_SRC:   ${DEEPGEMM_SRC}"
+echo "DEPS_TAG:       ${DEPS_TAG}"
+echo "----------------------------------------"
+
+docker build \
+  -f "${DOCKERFILE}" "$(dirname "${DOCKERFILE}")" \
+  --build-arg BASE_IMG="${BASE_IMG}" \
+  --build-arg CUDA_VERSION="${CUDA_VERSION}" \
+  --build-arg ARCH="${ARCH}" \
+  --build-arg PYTHON_VERSION="${PYTHON_VERSION}" \
+  --build-arg PYTHON_TAG="${PY_TAG}" \
+  -t "${DEPS_TAG}" \
+  --network=host
+
+mkdir -p "${DEEPGEMM_SRC}/dist"
+
+docker run --rm \
+  --network=host \
+  -v "${DEEPGEMM_SRC}:/deepgemm" \
+  -w /deepgemm \
+  "${DEPS_TAG}" \
+  bash build_sgl_deep_gemm.sh
+
+echo "Wheels written to ${DEEPGEMM_SRC}/dist:"
+ls -lh "${DEEPGEMM_SRC}/dist"/*.whl 2>/dev/null || true
@@ -0,0 +1,66 @@
+#!/usr/bin/env bash
+# Rename a freshly-built sgl-deep-gemm wheel so its filename, METADATA Version,
+# and WHEEL platform tag carry the +cuXXX local version and manylinux2014_<arch>
+# tag expected by the sgl-whl index and PyPI upload step.
+#
+# Input:  dist/sgl_deep_gemm-<VERSION>-py3-none-any.whl
+# Output: dist/sgl_deep_gemm-<VERSION>+<CU_TAG>-py3-none-manylinux2014_<ARCH>.whl
+#
+# Usage: rename_wheels.sh <WHEEL_DIR> <CU_TAG> <ARCH>
+#   WHEEL_DIR: directory containing the *.whl file (e.g. DeepGEMM/dist)
+#   CU_TAG:    cu129 | cu130
+#   ARCH:      x86_64 | aarch64
+set -ex
+
+if [ $# -lt 3 ]; then
+  echo "Usage: $0 <WHEEL_DIR> <CU_TAG> <ARCH>"
+  exit 1
+fi
+
+WHEEL_DIR="$1"
+CU_TAG="$2"
+ARCH="$3"
+PLAT_TAG="manylinux2014_${ARCH}"
+
+PYTHON="${PYTHON:-python3}"
+"${PYTHON}" -m pip install --quiet wheel
+
+shopt -s nullglob
+wheel_files=("${WHEEL_DIR}"/sgl_deep_gemm-*.whl)
+if [ ${#wheel_files[@]} -eq 0 ]; then
+  echo "No sgl_deep_gemm wheel found under ${WHEEL_DIR}" >&2
+  exit 1
+fi
+
+for wheel in "${wheel_files[@]}"; do
+  TMPDIR=$(mktemp -d)
+  trap 'rm -rf -- "$TMPDIR"' ERR
+
+  "${PYTHON}" -m wheel unpack "$wheel" --dest "$TMPDIR"
+  UNPACKED=$(find "$TMPDIR" -mindepth 1 -maxdepth 1 -type d | head -1)
+  DIST_INFO=$(find "$UNPACKED" -maxdepth 1 -type d -name "*.dist-info" | head -1)
+  WHEEL_META="${DIST_INFO}/WHEEL"
+  METADATA_FILE="${DIST_INFO}/METADATA"
+
+  # Replace the py3-none-any tag with a platform-specific one.
+  sed -i "s/^Tag: py3-none-any$/Tag: py3-none-${PLAT_TAG}/" "$WHEEL_META"
+
+  ORIG_VERSION=$(grep '^Version:' "$METADATA_FILE" | head -1 | sed 's/^Version:[[:space:]]*//')
+  if [[ "$ORIG_VERSION" == *"+${CU_TAG}"* ]]; then
+    NEW_VERSION="$ORIG_VERSION"
+  else
+    NEW_VERSION="${ORIG_VERSION}+${CU_TAG}"
+    sed -i "s/^Version:.*/Version: ${NEW_VERSION}/" "$METADATA_FILE"
+    OLD_BASE=$(basename "$DIST_INFO")
+    NEW_BASE="${OLD_BASE/${ORIG_VERSION}/${NEW_VERSION}}"
+    mv "$DIST_INFO" "${UNPACKED}/${NEW_BASE}"
+  fi
+
+  rm -f "$wheel"
+  "${PYTHON}" -m wheel pack "$UNPACKED" --dest-dir "$WHEEL_DIR"
+  rm -rf "$TMPDIR"
+  trap - ERR
+done
+
+echo "Renamed wheels in ${WHEEL_DIR}:"
+ls -lh "${WHEEL_DIR}"/*.whl
@@ -0,0 +1,45 @@
+# Generates a PEP 503 simple index for sgl-deep-gemm wheels under
+# sgl-whl/cu<version>/sgl-deep-gemm/index.html. Mirrors the layout used by
+# update_kernel_whl_index.py so consumers can `pip install
+# sgl-deep-gemm --extra-index-url https://...whl/cu129`.
+
+import argparse
+import hashlib
+import pathlib
+import re
+
+SUPPORTED_CUDA_VERSIONS = ["129", "130"]
+
+
+def update_wheel_index(cuda_version, wheel_dir):
+    index_dir = pathlib.Path(f"sgl-whl/cu{cuda_version}/sgl-deep-gemm")
+    index_dir.mkdir(exist_ok=True, parents=True)
+    base_url = "https://github.com/sgl-project/whl/releases/download"
+
+    suffix = f"+cu{cuda_version}"
+    for path in sorted(pathlib.Path(wheel_dir).glob("*.whl")):
+        if suffix not in path.name:
+            continue
+        with open(path, "rb") as f:
+            sha256 = hashlib.sha256(f.read()).hexdigest()
+        match = re.match(r"sgl_deep_gemm-([0-9][^-+]*)(?:\+cu[0-9]+)?-", path.name)
+        if not match:
+            continue
+        ver = match.group(1)
+        full_url = f"{base_url}/v{ver}/{path.name}#sha256={sha256}"
+        with (index_dir / "index.html").open("a") as f:
+            f.write(f'<a href="{full_url}">{path.name}</a><br>\n')
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--cuda", type=str, required=True, choices=SUPPORTED_CUDA_VERSIONS
+    )
+    parser.add_argument("--wheel-dir", type=str, default="dist")
+    args = parser.parse_args()
+    update_wheel_index(args.cuda, args.wheel_dir)
+
+
+if __name__ == "__main__":
+    main()