feat: support sgl-kernel PyPI (#2433)

Co-authored-by: Zhangyi <[email protected]>

.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 }}

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
     "

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",

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()