From 8bd644765fc69219022eb5e48d8b4c693146a643 Mon Sep 17 00:00:00 2001 From: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com> Date: Tue, 24 Feb 2026 22:15:05 -0800 Subject: [PATCH] [AMD] Enable ROCm kvcache JIT path and add AMD CI coverage. (#18992) Co-authored-by: Cursor --- .github/workflows/pr-test-amd-rocm720.yml | 44 +++++++++++++++++ .github/workflows/pr-test-amd.yml | 44 +++++++++++++++++ docker/rocm.Dockerfile | 2 +- docker/rocm720.Dockerfile | 2 +- .../jit_kernel/benchmark/bench_store_cache.py | 17 +++++-- .../jit_kernel/csrc/elementwise/kvcache.cuh | 2 +- .../jit_kernel/include/sgl_kernel/tile.cuh | 4 +- .../jit_kernel/include/sgl_kernel/utils.cuh | 48 ++++++++++++++++++- python/sglang/jit_kernel/utils.py | 15 +++++- python/sglang/srt/mem_cache/memory_pool.py | 2 +- 10 files changed, 168 insertions(+), 12 deletions(-) diff --git a/.github/workflows/pr-test-amd-rocm720.yml b/.github/workflows/pr-test-amd-rocm720.yml index d47168a18..42bcad325 100644 --- a/.github/workflows/pr-test-amd-rocm720.yml +++ b/.github/workflows/pr-test-amd-rocm720.yml @@ -65,6 +65,7 @@ jobs: outputs: main_package: ${{ steps.filter.outputs.main_package || steps.run-mode.outputs.run_all_tests }} sgl_kernel: ${{ steps.filter.outputs.sgl_kernel || steps.run-mode.outputs.run_all_tests }} + jit_kernel: ${{ steps.filter.outputs.jit_kernel || steps.run-mode.outputs.run_all_tests }} multimodal_gen: ${{ steps.filter.outputs.multimodal_gen || steps.run-mode.outputs.run_all_tests }} steps: - name: Checkout code @@ -102,6 +103,9 @@ jobs: sgl_kernel: - "sgl-kernel/**" - ".github/workflows/pr-test-amd-rocm720.yml" + jit_kernel: + - "python/sglang/jit_kernel/**" + - ".github/workflows/pr-test-amd-rocm720.yml" multimodal_gen: - "python/sglang/multimodal_gen/**" - "python/sglang/cli/**" @@ -238,6 +242,45 @@ jobs: run: | bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-a-test-1-amd --continue-on-error + jit-kernel-unit-test-amd: + needs: [check-changes] + if: | + always() && + ( + (inputs.target_stage == 'jit-kernel-unit-test-amd') || + ( + !inputs.target_stage && + needs.check-changes.outputs.jit_kernel == 'true' + ) + ) + strategy: + fail-fast: false + matrix: + runner: [linux-mi325-gpu-1] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} + + - name: Ensure VRAM is clear + run: bash scripts/ensure_vram_clear.sh rocm + + - name: Start CI container + run: bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: | + bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-aiter-build + + - name: Run JIT kernel unit tests + timeout-minutes: 10 + run: | + bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout" python3 -m pytest -q python/sglang/jit_kernel/tests/test_store_cache.py + stage-b-test-small-1-gpu-amd: needs: [check-changes] if: | @@ -756,6 +799,7 @@ jobs: multimodal-gen-test-2-gpu-amd, stage-a-test-1-amd, + jit-kernel-unit-test-amd, stage-b-test-small-1-gpu-amd, stage-b-test-small-1-gpu-amd-mi35x, stage-b-test-large-1-gpu-amd, diff --git a/.github/workflows/pr-test-amd.yml b/.github/workflows/pr-test-amd.yml index 5458de7fb..454397c78 100644 --- a/.github/workflows/pr-test-amd.yml +++ b/.github/workflows/pr-test-amd.yml @@ -62,6 +62,7 @@ jobs: outputs: main_package: ${{ steps.filter.outputs.main_package || steps.run-mode.outputs.run_all_tests }} sgl_kernel: ${{ steps.filter.outputs.sgl_kernel || steps.run-mode.outputs.run_all_tests }} + jit_kernel: ${{ steps.filter.outputs.jit_kernel || steps.run-mode.outputs.run_all_tests }} multimodal_gen: ${{ steps.filter.outputs.multimodal_gen || steps.run-mode.outputs.run_all_tests }} steps: - name: Checkout code @@ -99,6 +100,9 @@ jobs: sgl_kernel: - "sgl-kernel/**" - ".github/workflows/pr-test-amd.yml" + jit_kernel: + - "python/sglang/jit_kernel/**" + - ".github/workflows/pr-test-amd.yml" multimodal_gen: - "python/sglang/multimodal_gen/**" - "python/sglang/cli/**" @@ -235,6 +239,45 @@ jobs: run: | bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-a-test-1-amd + jit-kernel-unit-test-amd: + needs: [check-changes] + if: | + always() && + ( + (inputs.target_stage == 'jit-kernel-unit-test-amd') || + ( + !inputs.target_stage && + needs.check-changes.outputs.jit_kernel == 'true' + ) + ) + strategy: + fail-fast: false + matrix: + runner: [linux-mi325-gpu-1] + runs-on: ${{matrix.runner}} + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} + + - name: Ensure VRAM is clear + run: bash scripts/ensure_vram_clear.sh rocm + + - name: Start CI container + run: bash scripts/ci/amd/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: | + bash scripts/ci/amd/amd_ci_install_dependency.sh + + - name: Run JIT kernel unit tests + timeout-minutes: 10 + run: | + bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout" python3 -m pytest -q python/sglang/jit_kernel/tests/test_store_cache.py + stage-b-test-small-1-gpu-amd: needs: [check-changes, stage-a-test-1-amd] if: | @@ -845,6 +888,7 @@ jobs: multimodal-gen-test-2-gpu-amd, stage-a-test-1-amd, + jit-kernel-unit-test-amd, stage-b-test-small-1-gpu-amd, stage-b-test-small-1-gpu-amd-mi35x, stage-b-test-large-1-gpu-amd, diff --git a/docker/rocm.Dockerfile b/docker/rocm.Dockerfile index e9f7b1160..232098812 100644 --- a/docker/rocm.Dockerfile +++ b/docker/rocm.Dockerfile @@ -280,7 +280,7 @@ RUN /bin/bash -lc 'set -euo pipefail; \ \ # TVM Python bits need Cython + z3 before configure. # Pin z3-solver==4.15.4.0: 4.15.4.0 has a manylinux wheel; 4.15.5.0 has no wheel and builds from source (fails: C++20 needs GCC 14+, image has GCC 11). - "$VENV_PIP" install --no-cache-dir "cython>=0.29.36,<3.0" "apache-tvm-ffi>=0.1.6" "z3-solver==4.15.4.0"; \ + "$VENV_PIP" install --no-cache-dir "cython>=0.29.36,<3.0" "apache-tvm-ffi @ git+https://github.com/apache/tvm-ffi.git@v0.1.9-rc1" "z3-solver==4.15.4.0"; \ \ # Clone + pin TileLang (bundled TVM), then build git clone --recursive "${TILELANG_REPO}" /opt/tilelang && \ diff --git a/docker/rocm720.Dockerfile b/docker/rocm720.Dockerfile index d7dae79a0..75b28c401 100644 --- a/docker/rocm720.Dockerfile +++ b/docker/rocm720.Dockerfile @@ -309,7 +309,7 @@ RUN /bin/bash -lc 'set -euo pipefail; \ \ # TVM Python bits need Cython + z3 before configure. # Pin z3-solver==4.15.4.0: 4.15.4.0 has a manylinux wheel; 4.15.5.0 has no wheel and builds from source (fails: C++20 needs GCC 14+, image has GCC 11). - "$VENV_PIP" install --no-cache-dir "cython>=0.29.36,<3.0" "apache-tvm-ffi>=0.1.6" "z3-solver==4.15.4.0"; \ + "$VENV_PIP" install --no-cache-dir "cython>=0.29.36,<3.0" "apache-tvm-ffi @ git+https://github.com/apache/tvm-ffi.git@v0.1.9-rc1" "z3-solver==4.15.4.0"; \ \ # Clone + pin TileLang (bundled TVM), then build git clone --recursive "${TILELANG_REPO}" /opt/tilelang && \ diff --git a/python/sglang/jit_kernel/benchmark/bench_store_cache.py b/python/sglang/jit_kernel/benchmark/bench_store_cache.py index 2aeb7708d..1f179242e 100644 --- a/python/sglang/jit_kernel/benchmark/bench_store_cache.py +++ b/python/sglang/jit_kernel/benchmark/bench_store_cache.py @@ -14,6 +14,9 @@ from sglang.jit_kernel.benchmark.utils import ( ) from sglang.jit_kernel.kvcache import store_cache +_is_hip = bool(torch.version.hip) +HAS_AOT_STORE_CACHE = hasattr(torch.ops.sgl_kernel, "store_kv_cache") + def sglang_aot_store_cache( k: torch.Tensor, @@ -77,9 +80,14 @@ ITEM_SIZE = get_benchmark_range( ci_range=[1024], ) -LINE_VALS = ["aot", "jit", "torch_compile", "torch_streams"] -LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "PyTorch Compile", "PyTorch 2 Stream"] -STYLES = [("orange", "-"), ("blue", "--"), ("red", ":"), ("green", "-.")] +LINE_VALS = ["jit", "torch_compile", "torch_streams"] +LINE_NAMES = ["SGL JIT Kernel", "PyTorch Compile", "PyTorch 2 Stream"] +STYLES = [("blue", "--"), ("red", ":"), ("green", "-.")] +# Keep non-HIP benchmark lines unchanged; only HIP tolerates missing AOT op. +if (not _is_hip) or HAS_AOT_STORE_CACHE: + LINE_VALS = ["aot"] + LINE_VALS + LINE_NAMES = ["SGL AOT Kernel"] + LINE_NAMES + STYLES = [("orange", "-")] + STYLES X_NAMES = ["item_size", "batch_size"] CONFIGS = list(itertools.product(ITEM_SIZE, BS_RANGE)) @@ -116,11 +124,12 @@ def benchmark( torch.cuda.synchronize() FN_MAP = { - "aot": sglang_aot_store_cache, "jit": sglang_jit_store_cache, "torch_compile": torch_compile_store_cache, "torch_streams": torch_streams_store_cache, } + if (not _is_hip) or HAS_AOT_STORE_CACHE: + FN_MAP["aot"] = sglang_aot_store_cache def fn(): impl = FN_MAP[provider] diff --git a/python/sglang/jit_kernel/csrc/elementwise/kvcache.cuh b/python/sglang/jit_kernel/csrc/elementwise/kvcache.cuh index e5a1336b0..fa17cbf88 100644 --- a/python/sglang/jit_kernel/csrc/elementwise/kvcache.cuh +++ b/python/sglang/jit_kernel/csrc/elementwise/kvcache.cuh @@ -149,7 +149,7 @@ struct StoreKVCacheKernel { auto dtype = SymbolicDType{}; auto device = SymbolicDevice{}; auto indice_dtype = SymbolicDType{}; - device.set_options(); + device.set_options(); TensorMatcher({B, D}) // .with_strides({KS, 1}) diff --git a/python/sglang/jit_kernel/include/sgl_kernel/tile.cuh b/python/sglang/jit_kernel/include/sgl_kernel/tile.cuh index d227a5585..1328c64f9 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/tile.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/tile.cuh @@ -13,10 +13,10 @@ struct Memory { return Memory{0, 1}; } SGL_DEVICE static Memory warp(int warp_threads = kWarpThreads) { - return Memory{threadIdx.x % warp_threads, warp_threads}; + return Memory{static_cast(threadIdx.x % warp_threads), static_cast(warp_threads)}; } SGL_DEVICE static Memory cta(int cta_threads = blockDim.x) { - return Memory{threadIdx.x, cta_threads}; + return Memory{static_cast(threadIdx.x), static_cast(cta_threads)}; } SGL_DEVICE T load(const void* ptr, int64_t offset = 0) const { return static_cast(ptr)[tid + offset * tsize]; diff --git a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh index 56e38da76..eae04e68d 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh @@ -7,11 +7,29 @@ #include #include +#include +#ifndef USE_ROCM #include #include #include #include -#include +#else +#include +#include +#include +#ifndef __grid_constant__ +#define __grid_constant__ +#endif +using cudaError_t = hipError_t; +using cudaStream_t = hipStream_t; +using cudaLaunchConfig_t = hipLaunchConfig_t; +using cudaLaunchAttribute = hipLaunchAttribute; +inline constexpr auto cudaSuccess = hipSuccess; +#define cudaStreamPerThread hipStreamPerThread +#define cudaGetErrorString hipGetErrorString +#define cudaGetLastError hipGetLastError +#define cudaLaunchKernel hipLaunchKernel +#endif #ifndef USE_ROCM using fp32_t = float; @@ -26,6 +44,18 @@ using bf16x2_t = __nv_bfloat162; using fp8x2_e4m3_t = __nv_fp8x2_e4m3; using fp8x2_e5m2_t = __nv_fp8x2_e5m2; +using fp32x4_t = float4; +#else +using fp32_t = float; +using fp16_t = __half; +using bf16_t = __hip_bfloat16; +using fp8_e4m3_t = uint8_t; +using fp8_e5m2_t = uint8_t; +using fp32x2_t = float2; +using fp16x2_t = half2; +using bf16x2_t = __hip_bfloat162; +using fp8x2_e4m3_t = uint16_t; +using fp8x2_e5m2_t = uint16_t; using fp32x4_t = float4; #endif @@ -146,6 +176,10 @@ struct LaunchKernel { } auto enable_pdl(bool enabled = true) -> LaunchKernel& { +#ifdef USE_ROCM + (void)enabled; + m_config.numAttrs = 0; +#else if (enabled) { m_attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; m_attrs[0].val.programmaticStreamSerializationAllowed = true; @@ -154,12 +188,24 @@ struct LaunchKernel { } else { m_config.numAttrs = 0; } +#endif return *this; } template auto operator()(T&& kernel, Args&&... args) const -> void { +#ifdef USE_ROCM + hipLaunchKernelGGL( + std::forward(kernel), + m_config.gridDim, + m_config.blockDim, + m_config.dynamicSmemBytes, + m_config.stream, + std::forward(args)...); + RuntimeDeviceCheck(m_location); +#else RuntimeDeviceCheck(::cudaLaunchKernelEx(&m_config, kernel, std::forward(args)...), m_location); +#endif } private: diff --git a/python/sglang/jit_kernel/utils.py b/python/sglang/jit_kernel/utils.py index ede47831a..69c84dcb7 100644 --- a/python/sglang/jit_kernel/utils.py +++ b/python/sglang/jit_kernel/utils.py @@ -61,6 +61,9 @@ KERNEL_PATH = _resolve_kernel_path() DEFAULT_INCLUDE = [str(KERNEL_PATH / "include")] DEFAULT_CFLAGS = ["-std=c++20", "-O3"] DEFAULT_CUDA_CFLAGS = ["-std=c++20", "-O3", "--expt-relaxed-constexpr"] +DEFAULT_HIP_CFLAGS = [ + flag for flag in DEFAULT_CUDA_CFLAGS if flag != "--expt-relaxed-constexpr" +] DEFAULT_LDFLAGS = [] CPP_TEMPLATE_TYPE: TypeAlias = Union[int, float, bool, torch.dtype] @@ -77,6 +80,12 @@ CPP_DTYPE_MAP = { } +# AMD/ROCm note: +@cache_once +def is_hip_runtime() -> bool: + return bool(torch.version.hip) + + def make_cpp_args(*args: CPP_TEMPLATE_TYPE) -> CPPArgList: def _convert(arg: CPP_TEMPLATE_TYPE) -> str: if isinstance(arg, bool): @@ -156,6 +165,10 @@ def load_jit( # Override TVM_FFI_CUDA_ARCH_LIST if it does not exist. env_key = "TVM_FFI_CUDA_ARCH_LIST" env_existed = env_key in os.environ + selected_cuda_cflags = DEFAULT_CUDA_CFLAGS + if is_hip_runtime(): + selected_cuda_cflags = DEFAULT_HIP_CFLAGS + extra_cuda_cflags = ["-DUSE_ROCM"] + extra_cuda_cflags if not env_existed: os.environ[env_key] = _get_cuda_arch_list() try: @@ -164,7 +177,7 @@ def load_jit( cpp_sources=cpp_sources, cuda_sources=cuda_sources, extra_cflags=DEFAULT_CFLAGS + extra_cflags, - extra_cuda_cflags=DEFAULT_CUDA_CFLAGS + extra_cuda_cflags, + extra_cuda_cflags=selected_cuda_cflags + extra_cuda_cflags, extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, build_directory=build_directory, diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 1d917137c..317be57a3 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -99,7 +99,7 @@ def _set_kv_buffer_impl( same_kv_dim: bool = True, ) -> None: row_bytes = row_dim * store_dtype.itemsize - if _is_cuda and same_kv_dim and can_use_store_cache(row_bytes): + if (_is_cuda or _is_hip) and same_kv_dim and can_use_store_cache(row_bytes): return store_cache( k.view(-1, row_dim), v.view(-1, row_dim),