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

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

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