From ab7d5829cde4109ce68154c5fd3af3e857cad928 Mon Sep 17 00:00:00 2001 From: Alan Kao Date: Thu, 8 Jan 2026 12:18:29 +0800 Subject: [PATCH] [AMD] Add pip install / wheel build support for ROCm sgl-kernel (#15627) --- .github/workflows/release-whl-kernel.yml | 79 +++++++++ 3rdparty/amd/sgl-kernel/CMakeLists_rocm.txt | 159 ++++++++++++++++++ 3rdparty/amd/sgl-kernel/build_rocm.sh | 123 ++++++++++++++ 3rdparty/amd/sgl-kernel/rename_wheels_rocm.sh | 30 ++++ 3rdparty/amd/sgl-kernel/rocm_hipify.py | 40 +++++ scripts/update_kernel_whl_index.py | 35 +++- 6 files changed, 462 insertions(+), 4 deletions(-) create mode 100644 3rdparty/amd/sgl-kernel/CMakeLists_rocm.txt create mode 100755 3rdparty/amd/sgl-kernel/build_rocm.sh create mode 100755 3rdparty/amd/sgl-kernel/rename_wheels_rocm.sh create mode 100644 3rdparty/amd/sgl-kernel/rocm_hipify.py diff --git a/.github/workflows/release-whl-kernel.yml b/.github/workflows/release-whl-kernel.yml index fd105f16d..2fe1e8aef 100644 --- a/.github/workflows/release-whl-kernel.yml +++ b/.github/workflows/release-whl-kernel.yml @@ -196,3 +196,82 @@ jobs: git add -A git commit -m "update whl index" git push + + build-rocm700: + if: github.repository == 'sgl-project/sglang' + runs-on: amd-docker-scale + strategy: + matrix: + python-version: ["3.10"] + rocm-version: ["700"] + 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 + run: | + cp 3rdparty/amd/sgl-kernel/* sgl-kernel/ + cd sgl-kernel + chmod +x ./build_rocm.sh + ./build_rocm.sh "${{ matrix.rocm-version }}" + + - name: Upload artifacts + uses: actions/upload-artifact@v4 + with: + name: wheel-python${{ matrix.python-version }}-rocm${{ matrix.rocm-version }} + path: sgl-kernel/dist/* + + release-rocm700: + needs: build-rocm700 + 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.GH_PAT_FOR_WHL_RELEASE }} + 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.GH_PAT_FOR_WHL_RELEASE }} + + - name: Update wheel index + run: python3 scripts/update_kernel_whl_index.py --rocm 700 + + - name: Push wheel index + run: | + cd sgl-whl + git config --local user.name "sglang-bot" + git config --local user.email "sglangbot@gmail.com" + git add -A + git commit -m "update whl index" + git push diff --git a/3rdparty/amd/sgl-kernel/CMakeLists_rocm.txt b/3rdparty/amd/sgl-kernel/CMakeLists_rocm.txt new file mode 100644 index 000000000..5d225d872 --- /dev/null +++ b/3rdparty/amd/sgl-kernel/CMakeLists_rocm.txt @@ -0,0 +1,159 @@ +cmake_minimum_required(VERSION 3.24 FATAL_ERROR) +project(sgl_kernel LANGUAGES CXX) + +# Cmake +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_POSITION_INDEPENDENT_CODE ON) +set(CMAKE_SHARED_LIBRARY_PREFIX "") + +set(CMAKE_COLOR_DIAGNOSTICS ON) +set(CMAKE_VERBOSE_MAKEFILE ON CACHE BOOL "ON") + +# Python / Torch +find_package(Python COMPONENTS Interpreter Development.Module ${SKBUILD_SABI_COMPONENT} REQUIRED) + +execute_process( + COMMAND ${Python_EXECUTABLE} -c "import torch; print(torch.utils.cmake_prefix_path)" + OUTPUT_VARIABLE TORCH_PY_PREFIX + OUTPUT_STRIP_TRAILING_WHITESPACE +) + +set(Torch_DIR "${TORCH_PY_PREFIX}/Torch") +list(APPEND CMAKE_PREFIX_PATH "${TORCH_PY_PREFIX}/Torch") +find_package(Torch REQUIRED) + +execute_process( + COMMAND ${Python_EXECUTABLE} -c "import torch; print(int(torch._C._GLIBCXX_USE_CXX11_ABI))" + OUTPUT_VARIABLE TORCH_CXX11_ABI + OUTPUT_STRIP_TRAILING_WHITESPACE +) +if(TORCH_CXX11_ABI STREQUAL "0") + add_compile_definitions(_GLIBCXX_USE_CXX11_ABI=0) +else() + add_compile_definitions(_GLIBCXX_USE_CXX11_ABI=1) +endif() + +# ROCm/HIP +enable_language(HIP) +find_package(hip REQUIRED CONFIG) + +# Determine AMDGPU target from environment variable or default to gfx942 +set(AMDGPU_TARGET_ENV "$ENV{AMDGPU_TARGET}") + +if(AMDGPU_TARGET_ENV) + # Use environment variable if specified + set(AMDGPU_TARGETS "${AMDGPU_TARGET_ENV}") + message(STATUS "Using AMDGPU_TARGET from environment: ${AMDGPU_TARGETS}") +else() + # Default to gfx942 only + set(AMDGPU_TARGETS "gfx942") + message(STATUS "AMDGPU_TARGET not set, defaulting to gfx942") +endif() + +# Set HIP architectures +set(CMAKE_HIP_ARCHITECTURES ${AMDGPU_TARGETS}) + +# FP8 macro selection +# Always define HIP_FP8_TYPE_FNUZ=1 (for gfx942 and host compilation) +# Additionally define HIP_FP8_TYPE_E4M3=1 when building for gfx950 +# The existing utils.h logic will pick the right one based on architecture +set(SGL_FP8_MACROS "-DHIP_FP8_TYPE_FNUZ=1") + +if(AMDGPU_TARGETS MATCHES "gfx950") + list(APPEND SGL_FP8_MACROS "-DHIP_FP8_TYPE_E4M3=1") + message(STATUS "Multi-arch build: Enabling both HIP_FP8_TYPE_FNUZ (gfx942) and HIP_FP8_TYPE_E4M3 (gfx950)") +elseif(AMDGPU_TARGETS MATCHES "gfx942") + message(STATUS "Single-arch build: Enabling HIP_FP8_TYPE_FNUZ for gfx942") +else() + message(FATAL_ERROR "Unsupported AMDGPU_TARGET '${AMDGPU_TARGETS}'. Expected 'gfx942' or 'gfx950' or both.") +endif() + +# TopK dynamic smem bytes +# Dynamic shared-memory budget for the TopK kernels. +# - gfx942 (MI300/MI325): LDS is typically 64KB per workgroup -> keep dynamic smem <= ~48KB +# (leaves room for static shared allocations in the kernel). +# - gfx95x (MI350): LDS is larger (e.g. 160KB per CU) -> allow the original 128KB dynamic smem. +if(AMDGPU_TARGET_ONE STREQUAL "gfx942") + math(EXPR SGL_TOPK_DYNAMIC_SMEM_BYTES "48 * 1024") +else() + math(EXPR SGL_TOPK_DYNAMIC_SMEM_BYTES "32 * 1024 * 4") +endif() + +set(SGL_TOPK_MACROS "-DSGL_TOPK_DYNAMIC_SMEM_BYTES=${SGL_TOPK_DYNAMIC_SMEM_BYTES}") + +# Paths / includes +set(PROJ_ROOT ${CMAKE_CURRENT_LIST_DIR}) +set(SGL_INCLUDE_DIRS + ${PROJ_ROOT}/include + ${PROJ_ROOT}/include/impl + ${PROJ_ROOT}/csrc + ${TORCH_INCLUDE_DIRS} +) + +# Platform-specific library directory +set(PLAT_LIB_DIR "/usr/lib/x86_64-linux-gnu") +link_directories(${PLAT_LIB_DIR}) + +# Sources +set(SOURCES +${PROJ_ROOT}/csrc/allreduce/custom_all_reduce.hip +${PROJ_ROOT}/csrc/allreduce/deterministic_all_reduce.hip +${PROJ_ROOT}/csrc/allreduce/quick_all_reduce.hip +${PROJ_ROOT}/csrc/common_extension_rocm.cc +${PROJ_ROOT}/csrc/elementwise/activation.hip +${PROJ_ROOT}/csrc/elementwise/topk.hip +${PROJ_ROOT}/csrc/grammar/apply_token_bitmask_inplace_hip.hip +${PROJ_ROOT}/csrc/moe/moe_align_kernel.hip +${PROJ_ROOT}/csrc/moe/moe_topk_softmax_kernels.hip +${PROJ_ROOT}/csrc/moe/moe_topk_sigmoid_kernels.hip +${PROJ_ROOT}/csrc/speculative/eagle_utils.hip +${PROJ_ROOT}/csrc/kvcacheio/transfer.hip +${PROJ_ROOT}/csrc/elementwise/pos_enc.hip +) +set_source_files_properties( + ${SOURCES} + PROPERTIES + LANGUAGE HIP +) + +# Compile / Link flags +add_compile_options($<$:-O3>) + +set(SGL_HIP_FLAGS + -DNDEBUG + -DOPERATOR_NAMESPACE=sgl_kernel + -O3 + -std=c++17 + -DENABLE_BF16 + -DENABLE_FP8 + ${SGL_FP8_MACROS} + -Wno-pass-failed + -Wundefined-internal + ${SGL_TOPK_MACROS} +) + +# Python extension +Python_add_library(common_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI ${SOURCES}) +target_include_directories(common_ops PRIVATE ${SGL_INCLUDE_DIRS}) + +# Apply per-language flags +target_compile_options(common_ops PRIVATE + $<$:${SGL_HIP_FLAGS}> +) + +target_link_libraries(common_ops PRIVATE + ${TORCH_LIBRARIES} + hip::device + hip::host + hiprtc + amdhip64 +) + +target_link_options(common_ops PRIVATE + "SHELL:-Wl,-rpath,'\$ORIGIN/../../torch/lib'" +) + +install(TARGETS common_ops + LIBRARY DESTINATION sgl_kernel +) diff --git a/3rdparty/amd/sgl-kernel/build_rocm.sh b/3rdparty/amd/sgl-kernel/build_rocm.sh new file mode 100755 index 000000000..1022d8bb5 --- /dev/null +++ b/3rdparty/amd/sgl-kernel/build_rocm.sh @@ -0,0 +1,123 @@ +#!/bin/bash +set -euo pipefail +ROCM_VERSION=$1 + +PYTHON_ROOT_PATH="/opt/venv/bin" +AMDGPU_TARGET="gfx942;gfx950" + +echo "Python root path is: $PYTHON_ROOT_PATH" + +# Get version from git tags +SGLANG_VERSION="v0.5.6" # Default version, will be overridden if git tags are found + +# Fetch tags from origin to ensure we have the latest +if git fetch --tags origin; then + # Get the latest version tag sorted by version number (e.g., v0.5.7) + VERSION_FROM_TAG=$(git tag -l 'v[0-9]*' --sort=-v:refname | head -1) + if [ -n "$VERSION_FROM_TAG" ]; then + SGLANG_VERSION="$VERSION_FROM_TAG" + echo "Using SGLang version from git tags: $SGLANG_VERSION" + else + echo "Warning: No version tags found; using default $SGLANG_VERSION" >&2 + fi +else + echo "Warning: Failed to fetch tags from origin; using default $SGLANG_VERSION" >&2 +fi + +# Default base tags (can be overridden by command line arguments) +DEFAULT_MI30X_BASE_TAG="${SGLANG_VERSION}-rocm700-mi30x" +DEFAULT_MI35X_BASE_TAG="${SGLANG_VERSION}-rocm700-mi35x" + +# Parse command line arguments +MI30X_BASE_TAG="${DEFAULT_MI30X_BASE_TAG}" +MI35X_BASE_TAG="${DEFAULT_MI35X_BASE_TAG}" + +# Detect GPU architecture from the Kubernetes runner hostname +HOSTNAME_VALUE=$(hostname) +GPU_ARCH="mi30x" # default + +# Host names look like: linux-mi35x-gpu-1-xxxxx-runner-zzzzz +if [[ "${HOSTNAME_VALUE}" =~ ^linux-(mi[0-9]+[a-z]*)-gpu-[0-9]+ ]]; then + GPU_ARCH="${BASH_REMATCH[1]}" + echo "Detected GPU architecture from hostname: ${GPU_ARCH}" +else + echo "Warning: could not parse GPU architecture from '${HOSTNAME_VALUE}', defaulting to ${GPU_ARCH}" +fi + +case "${GPU_ARCH}" in + mi35x) + echo "Runner uses ${GPU_ARCH}; will fetch mi35x image." + ;; + mi30x|mi300|mi325) + echo "Runner uses ${GPU_ARCH}; will fetch mi30x image." + GPU_ARCH="mi30x" + ;; + *) + echo "Runner architecture '${GPU_ARCH}' unrecognised; defaulting to mi30x image." >&2 + GPU_ARCH="mi30x" + ;; +esac + +if [[ -f /etc/podinfo/gha-render-devices ]]; then + DEVICE_FLAG=$(cat /etc/podinfo/gha-render-devices) +else + DEVICE_FLAG="--device /dev/dri" +fi + +# Find the latest image +find_latest_image() { + local gpu_arch=$1 + local base_tag days_back image_tag + + case "${gpu_arch}" in + mi30x) base_tag="${MI30X_BASE_TAG}" ;; + mi35x) base_tag="${MI35X_BASE_TAG}" ;; + *) echo "Error: unsupported GPU architecture '${gpu_arch}'" >&2; return 1 ;; + esac + + for days_back in {0..6}; do + image_tag="${base_tag}-$(date -d "${days_back} days ago" +%Y%m%d)" + echo "Checking for image: rocm/sgl-dev:${image_tag}" >&2 + if docker manifest inspect "rocm/sgl-dev:${image_tag}" >/dev/null 2>&1; then + echo "Found available image: rocm/sgl-dev:${image_tag}" >&2 + echo "rocm/sgl-dev:${image_tag}" + return 0 + fi + done + + echo "Error: no ${gpu_arch} image found in the last 7 days for base ${base_tag}" >&2 + echo "Using hard-coded fallback…" >&2 + if [[ "${gpu_arch}" == "mi35x" ]]; then + echo "rocm/sgl-dev:v0.5.3-rocm700-mi35x-20251009" + else + echo "rocm/sgl-dev:v0.5.3-rocm700-mi30x-20251009" + fi +} + +# Pull and run the latest image +IMAGE=$(find_latest_image "${GPU_ARCH}") +echo "Pulling Docker image: ${IMAGE}" +docker pull "${IMAGE}" + +docker run --rm \ + -v $(pwd):/sgl-kernel \ + -e AMDGPU_TARGET="${AMDGPU_TARGET}" \ + ${IMAGE} \ + bash -c " + # Install CMake (version >= 3.26) - Robust Installation + export CMAKE_VERSION_MAJOR=3.31 + export CMAKE_VERSION_MINOR=1 + echo \"Downloading CMake from: https://cmake.org/files/v\${CMAKE_VERSION_MAJOR}/cmake-\${CMAKE_VERSION_MAJOR}.\${CMAKE_VERSION_MINOR}-linux-x86_64.tar.gz\" + wget https://cmake.org/files/v\${CMAKE_VERSION_MAJOR}/cmake-\${CMAKE_VERSION_MAJOR}.\${CMAKE_VERSION_MINOR}-linux-x86_64.tar.gz + tar -xzf cmake-\${CMAKE_VERSION_MAJOR}.\${CMAKE_VERSION_MINOR}-linux-x86_64.tar.gz + mv cmake-\${CMAKE_VERSION_MAJOR}.\${CMAKE_VERSION_MINOR}-linux-x86_64 /opt/cmake + export PATH=/opt/cmake/bin:\$PATH + + ${PYTHON_ROOT_PATH}/pip install --no-cache-dir ninja setuptools wheel numpy uv scikit-build-core && \ + + cd /sgl-kernel && \ + rm -rf CMakeLists.txt && mv CMakeLists_rocm.txt CMakeLists.txt && \ + ${PYTHON_ROOT_PATH}/python rocm_hipify.py && \ + ${PYTHON_ROOT_PATH}/python -m uv build --wheel -Cbuild-dir=build . --color=always --no-build-isolation && \ + ./rename_wheels_rocm.sh +" diff --git a/3rdparty/amd/sgl-kernel/rename_wheels_rocm.sh b/3rdparty/amd/sgl-kernel/rename_wheels_rocm.sh new file mode 100755 index 000000000..691407a3e --- /dev/null +++ b/3rdparty/amd/sgl-kernel/rename_wheels_rocm.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +set -ex + +WHEEL_DIR="dist" + +wheel_files=($WHEEL_DIR/*.whl) +for wheel in "${wheel_files[@]}"; do + intermediate_wheel="${wheel/linux/manylinux2014}" + + # Extract the current python version from the wheel name + if [[ $intermediate_wheel =~ -cp([0-9]+)- ]]; then + cp_version="${BASH_REMATCH[1]}" + else + echo "Could not extract Python version from wheel name: $intermediate_wheel" + continue + fi + + # Detect ROCm version and add appropriate suffix + if ls /opt | grep -q "7.0"; then + new_wheel="${intermediate_wheel/-cp${cp_version}/+rocm700-cp${cp_version}}" + 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." diff --git a/3rdparty/amd/sgl-kernel/rocm_hipify.py b/3rdparty/amd/sgl-kernel/rocm_hipify.py new file mode 100644 index 000000000..be83dd4b0 --- /dev/null +++ b/3rdparty/amd/sgl-kernel/rocm_hipify.py @@ -0,0 +1,40 @@ +from pathlib import Path + +import torch +from torch.utils.cpp_extension import CUDAExtension + +root = Path(__file__).parent.resolve() + +include_dirs = [ + root / "include", + root / "include" / "impl", + root / "csrc", +] + +sources = [ + "csrc/allreduce/custom_all_reduce.hip", + "csrc/allreduce/deterministic_all_reduce.hip", + "csrc/allreduce/quick_all_reduce.cu", + "csrc/common_extension_rocm.cc", + "csrc/elementwise/activation.cu", + "csrc/elementwise/topk.cu", + "csrc/grammar/apply_token_bitmask_inplace_cuda.cu", + "csrc/moe/moe_align_kernel.cu", + "csrc/moe/moe_topk_softmax_kernels.cu", + "csrc/moe/moe_topk_sigmoid_kernels.cu", + "csrc/speculative/eagle_utils.cu", + "csrc/kvcacheio/transfer.cu", + "csrc/elementwise/pos_enc.cu", +] + +libraries = ["hiprtc", "amdhip64", "c10", "torch", "torch_python"] + +ext_modules = [ + CUDAExtension( + name="sgl_kernel.common_ops", + sources=sources, + include_dirs=include_dirs, + libraries=libraries, + py_limited_api=False, + ), +] diff --git a/scripts/update_kernel_whl_index.py b/scripts/update_kernel_whl_index.py index 1b812a6a3..fbe7eb940 100644 --- a/scripts/update_kernel_whl_index.py +++ b/scripts/update_kernel_whl_index.py @@ -11,6 +11,10 @@ DEFAULT_CUDA_VERSION = "129" def check_wheel_cuda_version(path_name, target_cuda_version): + # Pass ROCm wheel + if re.search(f"rocm", path_name): + return False + # For other CUDA versions, the wheel path name will contain the cuda version suffix, e.g. sgl_kernel-0.3.16.post5+cu130-cp310-abi3-manylinux2014_x86_64.whl if target_cuda_version != DEFAULT_CUDA_VERSION: return target_cuda_version in path_name @@ -23,9 +27,9 @@ def check_wheel_cuda_version(path_name, target_cuda_version): return True -def update_wheel_index(cuda_version=DEFAULT_CUDA_VERSION): +def update_wheel_index(cuda_version=DEFAULT_CUDA_VERSION, rocm_version=None): index_dir = pathlib.Path(f"sgl-whl/cu{cuda_version}/sgl-kernel") - index_dir.mkdir(exist_ok=True) + index_dir.mkdir(exist_ok=True, parents=True) base_url = "https://github.com/sgl-project/whl/releases/download" for path in sorted(pathlib.Path("sgl-kernel/dist").glob("*.whl")): @@ -42,11 +46,34 @@ def update_wheel_index(cuda_version=DEFAULT_CUDA_VERSION): f.write(f'{path.name}
\n') +def update_wheel_index_rocm(rocm_version): + index_dir = pathlib.Path(f"sgl-whl/rocm{rocm_version}/sgl-kernel") + index_dir.mkdir(exist_ok=True, parents=True) + base_url = "https://github.com/sgl-project/whl/releases/download" + + for path in sorted(pathlib.Path("sgl-kernel/dist").glob("*.whl")): + # Skip the wheel if not rocm + if re.search(f"rocm", path.name) is None: + continue + with open(path, "rb") as f: + sha256 = hashlib.sha256(f.read()).hexdigest() + ver = re.findall( + r"sgl_kernel-([0-9.]+(?:\.post[0-9]+)?)(?:\+rocm[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'{path.name}
\n') + + def main(): parser = argparse.ArgumentParser() - parser.add_argument("--cuda", type=str, default="118") + parser.add_argument("--cuda", type=str, default=DEFAULT_CUDA_VERSION) + parser.add_argument("--rocm", type=str, default=None) args = parser.parse_args() - update_wheel_index(args.cuda) + if args.rocm is None: + update_wheel_index(args.cuda) + else: + update_wheel_index_rocm(args.rocm) if __name__ == "__main__":