[Misc] Use cache_once for is_arch_support_pdl in sgl-kernel (#22725)

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