chore: vendor sglang v0.5.10 snapshot

This commit is contained in:
2026-04-24 12:29:36 +00:00
parent 78f0d15221
commit bded08301f
4308 changed files with 1200894 additions and 2 deletions

View File

@@ -0,0 +1,161 @@
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)
list(APPEND CMAKE_PREFIX_PATH "/opt/rocm/lib/cmake/hip-lang")
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/pos_enc.hip
${PROJ_ROOT}/csrc/elementwise/topk.hip
${PROJ_ROOT}/csrc/grammar/apply_token_bitmask_inplace_hip.hip
${PROJ_ROOT}/csrc/kvcacheio/transfer.hip
${PROJ_ROOT}/csrc/memory/weak_ref_tensor.cpp
${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
)
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
)

View File

@@ -0,0 +1,50 @@
#!/bin/bash
set -euo pipefail
ROCM_VERSION=${1:-}
if [[ "${ROCM_VERSION}" == "700" ]]; then
IMAGE="lmsysorg/sglang:v0.5.8.post1-rocm700-mi35x"
elif [[ "${ROCM_VERSION}" == "720" ]]; then
IMAGE="rocm/pytorch:rocm7.2_ubuntu22.04_py3.10_pytorch_release_2.9.1"
else
echo "ERROR: Unsupported ROCM_VERSION='${ROCM_VERSION}'. Only '700' and '720' are supported." >&2
exit 1
fi
PYTHON_ROOT_PATH="/opt/venv/bin"
AMDGPU_TARGET="gfx942;gfx950"
# Pull and run the latest image
echo "Pulling Docker image: ${IMAGE}"
docker pull "${IMAGE}"
docker run --rm \
-v $(pwd):/sgl-kernel \
-e AMDGPU_TARGET="${AMDGPU_TARGET}" \
-e PYTORCH_ROCM_ARCH="${AMDGPU_TARGET}" \
${IMAGE} \
bash -c "
# Install torch, triton, and friends, depending on the ROCm version
if [[ "${ROCM_VERSION}" == "700" ]]; then
${PYTHON_ROOT_PATH}/pip install https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/torch-2.9.1.dev20251204%2Brocm7.0.2.lw.git351ff442-cp310-cp310-linux_x86_64.whl https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/triton-3.5.1%2Brocm7.0.2.gita272dfa8-cp310-cp310-linux_x86_64.whl https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/torchaudio-2.9.0%2Brocm7.0.2.gite3c6ee2b-cp310-cp310-linux_x86_64.whl https://repo.radeon.com/rocm/manylinux/rocm-rel-7.0.2/torchvision-0.24.0%2Brocm7.0.2.gitb919bd0c-cp310-cp310-linux_x86_64.whl
elif [[ "${ROCM_VERSION}" == "720" ]]; then
${PYTHON_ROOT_PATH}/pip install https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torch-2.9.1%2Brocm7.2.0.lw.git7e1940d4-cp310-cp310-linux_x86_64.whl https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/triton-3.5.1%2Brocm7.2.0.gita272dfa8-cp310-cp310-linux_x86_64.whl https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchaudio-2.9.0%2Brocm7.2.0.gite3c6ee2b-cp310-cp310-linux_x86_64.whl https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.24.0%2Brocm7.2.0.gitb919bd0c-cp310-cp310-linux_x86_64.whl
fi
# 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
"

View File

@@ -0,0 +1,28 @@
#!/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}"
[[ "$intermediate_wheel" == *"+rocm"* ]] && continue
# 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
ver_abrv=$(realpath /opt/rocm-* | sed -e 's/.*-//' -e 's/\.//g')
new_wheel=${intermediate_wheel/-cp${cp_version}/+rocm${ver_abrv}-cp${cp_version}}
if [[ "$wheel" != "$new_wheel" ]]; then
echo "Renaming $wheel to $new_wheel"
mv -- "$wheel" "$new_wheel"
fi
done
echo "Wheel renaming completed."

View File

@@ -0,0 +1,41 @@
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/pos_enc.cu",
"csrc/elementwise/topk.cu",
"csrc/grammar/apply_token_bitmask_inplace_cuda.cu",
"csrc/kvcacheio/transfer.cu",
"csrc/memory/weak_ref_tensor.cpp",
"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",
]
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,
),
]