From 9b7cf9ee6c299e85e3273842ee2b007312f9276d Mon Sep 17 00:00:00 2001
From: Yineng Zhang <me@zhyncs.com>
Date: Mon, 24 Mar 2025 20:53:23 -0700
Subject: [PATCH] support cu128 sgl-kernel (#4744)

---
 .../workflows/release-whl-kernel-cu128.yml    | 92 +++++++++++++++++++
 scripts/update_kernel_whl_index.py            | 31 +++++--
 2 files changed, 115 insertions(+), 8 deletions(-)
 create mode 100644 .github/workflows/release-whl-kernel-cu128.yml

diff --git a/.github/workflows/release-whl-kernel-cu128.yml b/.github/workflows/release-whl-kernel-cu128.yml
new file mode 100644
index 000000000..9d070afd6
--- /dev/null
+++ b/.github/workflows/release-whl-kernel-cu128.yml
@@ -0,0 +1,92 @@
+name: Release SGLang Kernel Wheel (cu128)
+
+on:
+  workflow_dispatch:
+    inputs:
+      tag_name:
+        type: string
+  push:
+    branches:
+      - main
+    paths:
+      - sgl-kernel/python/sgl_kernel/version.py
+
+jobs:
+  build-wheels:
+    if: github.repository == 'sgl-project/sglang'
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ['3.9']
+        cuda-version: ['12.8']
+
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          submodules: 'recursive'
+
+      - 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 artifacts
+        uses: actions/upload-artifact@v4
+        with:
+          name: wheel-python${{ matrix.python-version }}-cuda${{ matrix.cuda-version }}
+          path: sgl-kernel/dist/*
+
+  release:
+    needs: build-wheels
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Download artifacts
+        uses: actions/download-artifact@v4
+        with:
+          path: sgl-kernel/dist/
+          merge-multiple: true
+          pattern: wheel-*
+
+      - name: Set tag name
+        id: set_tag_name
+        run: |
+          if [ -z "${{ inputs.tag_name }}" ]; then
+            TAG_NAME="v$(cat sgl-kernel/python/sgl_kernel/version.py | cut -d'"' -f2)"
+            echo "tag_name=$TAG_NAME" >> $GITHUB_OUTPUT
+          else
+            echo "tag_name=${{ inputs.tag_name }}" >> $GITHUB_OUTPUT
+          fi
+
+      - name: Release
+        uses: softprops/action-gh-release@v2
+        with:
+          tag_name: ${{ steps.set_tag_name.outputs.tag_name }}
+          repository: sgl-project/whl
+          token: ${{ secrets.WHL_TOKEN }}
+          files: |
+            sgl-kernel/dist/*
+
+      - name: Clone wheel index
+        run: git clone https://oauth2:${WHL_TOKEN}@github.com/sgl-project/whl.git sgl-whl
+        env:
+          WHL_TOKEN: ${{ secrets.WHL_TOKEN }}
+
+      - name: Update wheel index
+        run: python3 scripts/update_kernel_whl_index.py --cuda 128
+
+      - name: Push wheel index
+        run: |
+          cd sgl-whl
+          git config --local user.name "github-actions[bot]"
+          git config --local user.email "41898282+github-actions[bot]@users.noreply.github.com"
+          git add -A
+          git commit -m "update whl index"
+          git push
diff --git a/scripts/update_kernel_whl_index.py b/scripts/update_kernel_whl_index.py
index a42969641..8e0a2b5f3 100644
--- a/scripts/update_kernel_whl_index.py
+++ b/scripts/update_kernel_whl_index.py
@@ -1,16 +1,31 @@
 # Reference: https://github.com/flashinfer-ai/flashinfer/blob/v0.2.0/scripts/update_whl_index.py
 
+import argparse
 import hashlib
 import pathlib
 import re
 
-for path in sorted(pathlib.Path("sgl-kernel/dist").glob("*.whl")):
-    with open(path, "rb") as f:
-        sha256 = hashlib.sha256(f.read()).hexdigest()
-    ver = re.findall(r"sgl_kernel-([0-9.]+(?:\.post[0-9]+)?)-", path.name)[0]
-    index_dir = pathlib.Path(f"sgl-whl/cu118/sgl-kernel")
+
+def update_wheel_index(cuda_version="118"):
+    index_dir = pathlib.Path(f"sgl-whl/cu{cuda_version}/sgl-kernel")
     index_dir.mkdir(exist_ok=True)
     base_url = "https://github.com/sgl-project/whl/releases/download"
-    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')
+
+    for path in sorted(pathlib.Path("sgl-kernel/dist").glob("*.whl")):
+        with open(path, "rb") as f:
+            sha256 = hashlib.sha256(f.read()).hexdigest()
+        ver = re.findall(r"sgl_kernel-([0-9.]+(?:\.post[0-9]+)?)-", path.name)[0]
+        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, default="118")
+    args = parser.parse_args()
+    update_wheel_index(args.cuda)
+
+
+if __name__ == "__main__":
+    main()