EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Cache the result of is_blackwell platform check (#10498)

Browse Source
This commit is contained in:
b8zhong
2025-09-15 22:30:28 -07:00
committed by GitHub
parent 5fe39e85a2
commit b2435be682
2 changed files with 3 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
python/sglang/srt/layers/quantization/deep_gemm_wrapper/configurer.py
View File

@@ -2,7 +2,7 @@ import logging
import torch
from sglang.srt.utils import get_bool_env_var, get_device_sm
from sglang.srt.utils import get_bool_env_var, get_device_sm, is_blackwell
logger = logging.getLogger(__name__)
@@ -21,12 +21,7 @@ def _compute_enable_deep_gemm():
return get_bool_env_var("SGL_ENABLE_JIT_DEEPGEMM", default="true")
def _is_blackwell_arch() -> bool:
major, minor = torch.cuda.get_device_capability(torch.cuda.current_device())
return major == 10
ENABLE_JIT_DEEPGEMM = _compute_enable_deep_gemm()
DEEPGEMM_BLACKWELL = ENABLE_JIT_DEEPGEMM and _is_blackwell_arch()
DEEPGEMM_BLACKWELL = ENABLE_JIT_DEEPGEMM and is_blackwell()
DEEPGEMM_SCALE_UE8M0 = DEEPGEMM_BLACKWELL