[Bugfix] Include SM121 in DeepGEMM packed-scale selection (#39482)

Signed-off-by: bojiang-li <327132355+bojiang-li@users.noreply.github.com>
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
This commit is contained in:
Bojiang Li
2026-09-18 13:57:14 +08:00
committed by GitHub
co-authored by Xiaoyu Zhang
parent 0dad91d50f
commit 3d1b9e7549
@@ -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)