Refactor DeepGEMM integration (#7150)

This commit is contained in:
fzyzcjy
2025-06-13 20:41:03 -07:00
committed by GitHub
parent 8b8f2e7463
commit b4c41f7276
12 changed files with 207 additions and 147 deletions
@@ -0,0 +1,26 @@
import logging
from sglang.srt.utils import get_bool_env_var, get_device_sm
logger = logging.getLogger(__name__)
def _compute_enable_deep_gemm():
try:
import deep_gemm
except ImportError:
logger.warning("Failed to import deep_gemm, disable ENABLE_JIT_DEEPGEMM.")
return False
sm_version = get_device_sm()
if sm_version < 90:
return False
return get_bool_env_var("SGL_ENABLE_JIT_DEEPGEMM", default="true")
ENABLE_JIT_DEEPGEMM = _compute_enable_deep_gemm()
DEEPGEMM_V202506 = False
DEEPGEMM_SCALE_UE8M0 = DEEPGEMM_V202506