diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py b/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py index 65261e33c..b467a153a 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py @@ -18,9 +18,9 @@ def _compute_enable_deep_gemm(): sm_version = get_device_sm() if (_is_cuda and sm_version < 90) or (_is_musa and sm_version < 31): return False - # SM120 support (mma.sync block-scale, no TMEM) landed in DeepGEMM#324; + # SM120/SM121 support (including GB10) landed in DeepGEMM#324; # probe the entry point since installed builds may predate it. - if sm_version == 120: + if sm_version in (120, 121): try: from deep_gemm import m_grouped_fp8_fp4_gemm_nt_contiguous # noqa: F401 except (ImportError, AttributeError): @@ -40,7 +40,7 @@ ENABLE_JIT_DEEPGEMM = _compute_enable_deep_gemm() DEEPGEMM_BLACKWELL = ENABLE_JIT_DEEPGEMM and get_platform().is_sm100 DEEPGEMM_SCALE_UE8M0 = ENABLE_JIT_DEEPGEMM and ( - get_platform().is_sm100 or get_device_sm() == 120 + get_platform().is_sm100 or get_device_sm() in (120, 121) ) DEEPGEMM_NEED_TMA_ALIGNED_SCALES = not (DEEPGEMM_SCALE_UE8M0 or _is_musa)