Skip to content

misc: cache is_hopper_arch #6799

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zhyncs merged 3 commits into from
Jun 1, 2025
Merged

misc: cache is_hopper_arch #6799

Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3ffcb8a
fix
Edenzzzz Jun 1, 2025
7b61380
use lru cache
Edenzzzz Jun 1, 2025
d64cce9
Merge branch 'main' into hopper_arch
Edenzzzz Jun 1, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 10 additions & 4 deletions sgl-kernel/python/sgl_kernel/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,8 +41,14 @@ def _to_tensor_scalar_tuple(x):
return (None, x)


_IS_HOPPER_ARCH = None
Copy link
Copy Markdown
Collaborator

@zhyncs zhyncs Jun 1, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it would be better to cache it after this line

sglang/sgl-kernel/python/sgl_kernel/elementwise.py

Line 4 in 1da8d23

from sgl_kernel.utils import get_cuda_stream, is_hopper_arch

Copy link
Copy Markdown
Collaborator

@zhyncs zhyncs Jun 1, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

But implementing in the function also has its advantages, except for element wise ops, everything else can be used.

Copy link
Copy Markdown
Contributor Author

@Edenzzzz Edenzzzz Jun 1, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yeah maybe keep it here because other functions may call it in the future?



def is_hopper_arch() -> bool:
# Hopper arch's compute capability == 9.0
device = torch.cuda.current_device()
major, minor = torch.cuda.get_device_capability(device)
return major == 9
global _IS_HOPPER_ARCH
if _IS_HOPPER_ARCH is None:
# Hopper arch's compute capability == 9.0
device = torch.cuda.current_device()
major, _ = torch.cuda.get_device_capability(device)
_IS_HOPPER_ARCH = major == 9
return _IS_HOPPER_ARCH
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot Jun 1, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The current implementation for caching is_hopper_arch using a global variable is functional. However, Python's functools.lru_cache decorator offers a more idiomatic and often cleaner way to achieve the same result with less boilerplate.

Using @functools.lru_cache(maxsize=1) (or maxsize=None if Python 3.9+ and you prefer cache) would:

  • Eliminate the need for the explicit _IS_HOPPER_ARCH global variable and the manual check-and-set logic.
  • Handle thread-safety aspects of cache access automatically (Python's lru_cache is thread-safe).
  • Make the intent of caching immediately clear through the decorator.

To implement this, you would need to add import functools at the top of the file.

What are your thoughts on refactoring to use lru_cache for this?

@functools.lru_cache(maxsize=1)
def is_hopper_arch() -> bool:
    # Hopper arch's compute capability == 9.0
    device = torch.cuda.current_device()
    major, _ = torch.cuda.get_device_capability(device)
    return major == 9