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