Skip to content

GitLab

Unverified Commit c429919d authored by Wenxuan Tan's avatar Wenxuan Tan Committed by GitHub
Browse files

misc: cache is_hopper_arch (#6799)

parent 1da8d230
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 0 deletions
sgl-kernel/python/sgl_kernel/utils.py
View file @ c429919d
...@@ -13,6 +13,7 @@ ...@@ -13,6 +13,7 @@
# limitations under the License. # limitations under the License.
# ============================================================================== # ==============================================================================
import functools
from typing import Dict, Tuple from typing import Dict, Tuple
import torch import torch
...@@ -41,6 +42,7 @@ def _to_tensor_scalar_tuple(x): ...@@ -41,6 +42,7 @@ def _to_tensor_scalar_tuple(x):
return (None, x) return (None, x)
@functools.lru_cache(maxsize=1)
def is_hopper_arch() -> bool: def is_hopper_arch() -> bool:
# Hopper arch's compute capability == 9.0 # Hopper arch's compute capability == 9.0
device = torch.cuda.current_device() device = torch.cuda.current_device()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment