Motivation

This PR aims to build and host sgl-kernel wheel for ROCm, which is a pre-requisite to build sglang wheel for ROCm.

Modifications

1. Added sgl-kernel/CMakeLists_rocm.txt that would be used by CMake for building wheel for ROCm, similar to NVIDIA's CMakeLists.txt. We use CMake over setup_rocm.py for building the wheel since the wheel built with CMake is Python version agnostic, whereas a wheel built with setup_rocm.py will only work with the Python version used to build it.
2. Added sgl-kernel/build_rocm.sh similar to NVIDIA's sgl-kernel/build.sh that builds the ROCm wheel inside a docker container (used by Github Workflows).
3. Added sgl-kernel/rename_wheels_rocm.sh similar to existing NVIDIA's sgl-kernel/rename_wheels.sh to rename wheels to the standard format. This script detects the rocm version used to build the wheel and renames it to the following template: sgl_kernel-<sglang_version>+rocm<rocm_version>-cp310-abi3-manylinux2014_x86_64.whl.
4. Added sgl-kernel/rocm_hipify.py that hipifies the sources using PyTorch's built in hipify module. This is required by CMake for build, as CMake expects HIP files to be compiled.
5. Updated .github/workflows/release-whl-kernel.yml to build and push ROCm 7.0 wheels to SGLang's kernel index.
6. Added ROCm support to scripts/ci/update_kernel_whl_index.py to update sgl-kernel wheel index.

Dependencies

This PR builds and releases sgl-kernel for ROCm 7.0, ensuring the package is available in the Releases section of github.com/sgl-project/whl.

However, for it to show up on the SGLang kernel wheel index completely, this small PR that adds the index file needs to be merged. Though this should not disrupt any functionality, as we only consume the URL for the package from the Releases section of github.com/sgl-project/whl

Usage Instructions

This PR adds support for building and releasing sgl-kernel to the SGLang Kernel Index.

Users are not expected to install sgl-kernel themselves, since it would be installed as a dependency by the SGLang wheel built with #14802. Users are only expected to pip install SGLang using the wheel built with #14802 and that should automatically install sgl-kernel as part of it.

However, if the user chooses to, they can install sgl-kernel on ROCm by choosing the specific ROCm version from SGLang kernel index and running the install command:

pip install sgl-kernel --index-url https://docs.sglang.io/whl/rocm700

Note: The sgl-kenel package is built with Torch 2.10.0 Nightly for ROCm 7.0 and the user must install the same torch version if they choose to install sgl-kernel using the pre-built wheel. Refer to the specific torch version in #14802 .

Maintenance

• New kernel: When a new kernel is added, similar to how we currently update sources in sgl-kernel/setup_rocm.py, the sources will need to be updated in sgl-kernel/build_rocm.py and sgl-kernel/CMakeLists_rocm.txt.
• New Torch Version: If we choose to update the Torch version used, the following changes are required:
  1. Update build_rocm.sh: This file determines the torch version used to build sgl-kernel.
  2. **Update pyproject_rocm.toml from [Feature] Add pyproject_rocm.toml for end-to-end ROCm pip installation support #14802:**The torch version specified in pyproject_rocm.toml is the version installed when user installs SGLang using the wheel. The versions for torchvision and pytorch-triton-rocm also need to be updated. To determine these, you can manually install the desired torch version, which would install compatible versions of torchvision and pytorch-triton-rocm, you can make a note of the compatible torchvision and pytorch-triton-rocm from here. Simply replace the versions of torch, 'torchvision, and pytorch-triton-rocm` with the new versions.
• New ROCm version: If we want to update the Torch version and the new Torch version is built for a new ROCm version, a new wheel must be built by modifying .github/workflows/release-whl-kernel.yml, sgl-kernel/rename_wheels_rocm.sh, and sgl-kernel/build_rocm.sh.
• New Architecture: To add support for a new architecture in the future, the AMDGPU_TARGET variable needs to be updated in sgl-kernel/build_rocm.sh. And any compiler flags specific to the architecture can be set in sgl-kernel/CMakeLists_rocm.txt.

Testing

Process

For testing the sgl-kernel, we build sgl-kernel wheel using desired ROCm and Torch version. Then, install sglang along with the same ROCm and Torch versions. Now, install sgl-kernel using the wheel built, and then run the following suite of tests (which is directly taken from the pr-test-amd.yml github workflow):

python3 -m pytest test_moe_align.py test_moe_topk_softmax.py speculative/test_eagle_utils.py test_apply_token_bitmask_inplace.py test_activation.py test_kvcacheio.py

Environments

The above tests were run on the following matrix of environments:
ROCm Versions: [7.0]
Python Versions: [3.10, 3.11, 3.12]
Hardware: [MI300x, MI350x]

Results

✅ All the unit tests run pass successfully.

Checklist

• Format your code according to the Format code with pre-commit.
• Add code support to build ROCm wheels for sgl-kernel
• Test pip install functionality for ROCm 6.3 and ROCm 7.0 on MI300x and MI350x.
• Add wheel release as part of the Github Workflow.
• Test functionality on Python 3.10, 3.11, 3.12 with install from TestPyPI.
• Update documentation according to Write documentations.