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

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-02-24 22:15:05 -08:00
parent aff2f130ec
commit 8bd644765f
10 changed files with 168 additions and 12 deletions
--- 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<kDLCUDA>();
+    device.set_options<kDLCUDA, kDLROCM>();

    TensorMatcher({B, D})  //
        .with_strides({KS, 1})