chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Hubert Lu
2026-02-24 22:15:05 -08:00
committed by GitHub
parent aff2f130ec
commit 8bd644765f
10 changed files with 168 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
python/sglang/jit_kernel/benchmark/bench_store_cache.py
View File

@@ -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]