[AMD] Add pip install / wheel build support for ROCm sgl-kernel (#15627)

2026-01-07 20:18:29 -08:00
parent 261860e17b
commit ab7d5829cd
6 changed files with 462 additions and 4 deletions
@@ -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($<$<COMPILE_LANGUAGE:CXX>:-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
+  $<$<COMPILE_LANGUAGE:HIP>:${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
+)
@@ -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
+"
@@ -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."
@@ -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,
+    ),
+]