diff --git a/.github/workflows/release-pypi-kernel.yml b/.github/workflows/release-pypi-kernel.yml new file mode 100644 index 00000000000..4b09f0854b0 --- /dev/null +++ b/.github/workflows/release-pypi-kernel.yml @@ -0,0 +1,37 @@ +name: Release SGLang Kernel to PyPI + +on: + push: + branches: + - main + paths: + - sgl-kernel/pyproject.toml + workflow_dispatch: + +jobs: + build-wheels: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ['3.9', '3.10', '3.11', '3.12'] + cuda-version: ['11.8', '12.1', '12.4'] + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Build wheels for Python ${{ matrix.python-version }} and CUDA ${{ matrix.cuda-version }} + run: | + cd sgl-kernel + chmod +x ./build.sh + ./build.sh "${{ matrix.python-version }}" "${{ matrix.cuda-version }}" + + - name: Upload to pypi + working-directory: sgl-kernel + run: | + pip install twine + python3 -m twine upload dist/* -u __token__ -p ${{ secrets.PYPI_TOKEN }} diff --git a/sgl-kernel/build.sh b/sgl-kernel/build.sh index b276f0141c2..799b724dfe6 100755 --- a/sgl-kernel/build.sh +++ b/sgl-kernel/build.sh @@ -1,13 +1,16 @@ #!/bin/bash - set -ex +PYTHON_VERSION=$1 +CUDA_VERSION=$2 +PYTHON_ROOT_PATH=/opt/python/cp${PYTHON_VERSION//.}-cp${PYTHON_VERSION//.} -docker run --rm -it \ +docker run --rm \ -v "$(pwd)":/sgl-kernel \ - pytorch/manylinux-builder:cuda12.1 \ + pytorch/manylinux-builder:cuda${CUDA_VERSION} \ bash -c " - pip install --no-cache-dir torch==2.4.0 --index-url https://download.pytorch.org/whl/cu121 && \ + ${PYTHON_ROOT_PATH}/bin/pip install --no-cache-dir torch==2.4.0 --index-url https://download.pytorch.org/whl/cu${CUDA_VERSION//.} && \ export TORCH_CUDA_ARCH_LIST='7.5 8.0 8.9 9.0+PTX' && \ + export CUDA_VERSION=${CUDA_VERSION} && \ cd /sgl-kernel && \ - python setup.py bdist_wheel + ${PYTHON_ROOT_PATH}/bin/python setup.py bdist_wheel " diff --git a/sgl-kernel/pyproject.toml b/sgl-kernel/pyproject.toml index 4330d2e19be..ce434bc34ab 100644 --- a/sgl-kernel/pyproject.toml +++ b/sgl-kernel/pyproject.toml @@ -12,8 +12,7 @@ license = { file = "LICENSE" } classifiers = [ "Programming Language :: Python :: 3", "License :: OSI Approved :: Apache Software License", - "Programming Language :: C++", - "Programming Language :: CUDA", + "Environment :: GPU :: NVIDIA CUDA" ] dependencies = [ "torch", diff --git a/sgl-kernel/setup.py b/sgl-kernel/setup.py index f2af83643e5..aaaae62bea2 100644 --- a/sgl-kernel/setup.py +++ b/sgl-kernel/setup.py @@ -1,10 +1,56 @@ -from setuptools import find_packages, setup +import os +import shutil +import zipfile +from pathlib import Path + +from setuptools import setup from torch.utils.cpp_extension import BuildExtension, CUDAExtension +root = Path(__file__).parent.resolve() + + +def get_version(): + with open(root / "pyproject.toml") as f: + for line in f: + if line.startswith("version"): + return line.split("=")[1].strip().strip('"') + + +def rename_wheel(): + if not os.environ.get("CUDA_VERSION"): + return + cuda_version = os.environ["CUDA_VERSION"].replace(".", "") + base_version = get_version() + + wheel_dir = Path("dist") + old_wheel = next(wheel_dir.glob("*.whl")) + tmp_dir = wheel_dir / "tmp" + tmp_dir.mkdir(exist_ok=True) + + with zipfile.ZipFile(old_wheel, "r") as zip_ref: + zip_ref.extractall(tmp_dir) + + old_info = tmp_dir / f"sgl_kernel-{base_version}.dist-info" + new_info = tmp_dir / f"sgl_kernel-{base_version}+cu{cuda_version}.dist-info" + old_info.rename(new_info) + + new_wheel = ( + wheel_dir + / f"sgl_kernel-{base_version}+cu{cuda_version}-{old_wheel.name.split('-', 2)[-1]}" + ) + with zipfile.ZipFile(new_wheel, "w", zipfile.ZIP_DEFLATED) as new_zip: + for file_path in tmp_dir.rglob("*"): + if file_path.is_file(): + new_zip.write(file_path, file_path.relative_to(tmp_dir)) + + old_wheel.unlink() + shutil.rmtree(tmp_dir) + + setup( name="sgl-kernel", - version="0.0.2", - packages=find_packages(where="src"), + version=get_version(), + packages=["sgl_kernel"], package_dir={"": "src"}, ext_modules=[ CUDAExtension( @@ -30,3 +76,5 @@ cmdclass={"build_ext": BuildExtension}, install_requires=["torch"], ) + +rename_wheel()