[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:
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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user