[AMD] Enable ROCm kvcache JIT path and add AMD CI coverage. (#18992)

Co-authored-by: Cursor <cursoragent@cursor.com>

.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,

.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,

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 <format> 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 && \

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 <format> 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 && \

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]

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<kDLCUDA>();
+    device.set_options<kDLCUDA, kDLROCM>();
 
     TensorMatcher({B, D})  //
         .with_strides({KS, 1})

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<uint32_t>(threadIdx.x % warp_threads), static_cast<uint32_t>(warp_threads)};
   }
   SGL_DEVICE static Memory cta(int cta_threads = blockDim.x) {
-    return Memory{threadIdx.x, cta_threads};
+    return Memory{static_cast<uint32_t>(threadIdx.x), static_cast<uint32_t>(cta_threads)};
   }
   SGL_DEVICE T load(const void* ptr, int64_t offset = 0) const {
     return static_cast<const T*>(ptr)[tid + offset * tsize];

python/sglang/jit_kernel/include/sgl_kernel/utils.cuh

@@ -7,11 +7,29 @@
 
 #include <concepts>
 #include <cstddef>
+#include <type_traits>
+#ifndef USE_ROCM
 #include <cuda_bf16.h>
 #include <cuda_fp16.h>
 #include <cuda_fp8.h>
 #include <cuda_runtime.h>
-#include <type_traits>
+#else
+#include <hip/hip_bf16.h>
+#include <hip/hip_fp16.h>
+#include <hip/hip_runtime.h>
+#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 <typename T, typename... Args>
   auto operator()(T&& kernel, Args&&... args) const -> void {
+#ifdef USE_ROCM
+    hipLaunchKernelGGL(
+        std::forward<T>(kernel),
+        m_config.gridDim,
+        m_config.blockDim,
+        m_config.dynamicSmemBytes,
+        m_config.stream,
+        std::forward<Args>(args)...);
+    RuntimeDeviceCheck(m_location);
+#else
     RuntimeDeviceCheck(::cudaLaunchKernelEx(&m_config, kernel, std::forward<Args>(args)...), m_location);
+#endif
   }
 
  private:

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,

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),