From b75275b6f2186032bec665fb4719bfbb8d321dbe Mon Sep 17 00:00:00 2001
From: Yineng Zhang <me@zhyncs.com>
Date: Fri, 11 Apr 2025 01:58:46 -0700
Subject: [PATCH] feat: add cu128 identifier for sgl-kernel (#5287)

---
 .github/workflows/release-whl-kernel-cu128.yml |  2 +-
 sgl-kernel/build.sh                            |  1 -
 sgl-kernel/rename_wheels.sh                    | 10 +++++++---
 3 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/release-whl-kernel-cu128.yml b/.github/workflows/release-whl-kernel-cu128.yml
index 9d070afd6..c6f59998b 100644
--- a/.github/workflows/release-whl-kernel-cu128.yml
+++ b/.github/workflows/release-whl-kernel-cu128.yml
@@ -14,7 +14,7 @@ on:
 jobs:
   build-wheels:
     if: github.repository == 'sgl-project/sglang'
-    runs-on: ubuntu-latest
+    runs-on: sgl-kernel-build-node
     strategy:
       matrix:
         python-version: ['3.9']
diff --git a/sgl-kernel/build.sh b/sgl-kernel/build.sh
index 608cbb647..945d60ca5 100755
--- a/sgl-kernel/build.sh
+++ b/sgl-kernel/build.sh
@@ -4,7 +4,6 @@ set -ex
 PYTHON_VERSION=$1
 CUDA_VERSION=$2
 PYTHON_ROOT_PATH=/opt/python/cp${PYTHON_VERSION//.}-cp${PYTHON_VERSION//.}
-ENABLE_SM90A=$(( ${CUDA_VERSION%.*} >= 12 ? ON : OFF ))
 
 if [ ${CUDA_VERSION} = "12.8" ]; then
    DOCKER_IMAGE="pytorch/manylinux2_28-builder:cuda${CUDA_VERSION}"
diff --git a/sgl-kernel/rename_wheels.sh b/sgl-kernel/rename_wheels.sh
index 95c241893..b52a94a92 100755
--- a/sgl-kernel/rename_wheels.sh
+++ b/sgl-kernel/rename_wheels.sh
@@ -5,13 +5,17 @@ WHEEL_DIR="dist"
 
 wheel_files=($WHEEL_DIR/*.whl)
 for wheel in "${wheel_files[@]}"; do
-    new_wheel="${wheel/linux/manylinux2014}"
+    intermediate_wheel="${wheel/linux/manylinux2014}"
+
+    if ls /usr/local/ | grep -q "12.8"; then
+        new_wheel="${intermediate_wheel/-cp39/+cu128-cp39}"
+    else
+        new_wheel="$intermediate_wheel"
+    fi
 
     if [[ "$wheel" != "$new_wheel" ]]; then
         echo "Renaming $wheel to $new_wheel"
         mv -- "$wheel" "$new_wheel"
-
     fi
 done
-
 echo "Wheel renaming completed."