From 222eda1598be6fb9c6dcaa1adb8d9b5f70364a38 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Tue, 14 Apr 2026 15:22:10 -0700 Subject: [PATCH] [Misc] Use cache_once for is_arch_support_pdl in sgl-kernel (#22725) --- sgl-kernel/python/sgl_kernel/utils.py | 29 +++++++++++++++++++++++---- 1 file changed, 25 insertions(+), 4 deletions(-) diff --git a/sgl-kernel/python/sgl_kernel/utils.py b/sgl-kernel/python/sgl_kernel/utils.py index d03476eff..7e98c6cc3 100644 --- a/sgl-kernel/python/sgl_kernel/utils.py +++ b/sgl-kernel/python/sgl_kernel/utils.py @@ -37,9 +37,30 @@ def _to_tensor_scalar_tuple(x): return (None, x) -@functools.lru_cache(maxsize=1) +def cache_once(fn): + """ + NOTE: `functools.lru_cache` is not compatible with `torch.compile` + So we manually implement a simple cache_once decorator to replace it. + """ + result_map = {} + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + key = (args, tuple(sorted(kwargs.items()))) + if key not in result_map: + result_map[key] = fn(*args, **kwargs) + return result_map[key] + + return wrapper + + +@cache_once def is_arch_support_pdl() -> bool: - # Hopper arch's compute capability == 9.0 - device = torch.cuda.current_device() - major, minor = torch.cuda.get_device_capability(device) + if bool(torch.version.hip): + return False + try: + device = torch.cuda.current_device() + major, _ = torch.cuda.get_device_capability(device) + except Exception: + return False return major >= 9