DeepSeek/DeepGEMM
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepGEMM synced 2025-06-26 23:15:49 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Lower TMA requirement

Browse Source
This commit is contained in:
Chenggang Zhao 2025-03-25 17:18:53 +08:00
parent ddccb230ca
commit 1999d553e5
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
deep_gemm/jit_kernels/gemm.py
View File

@ -106,7 +106,7 @@ def get_best_configs(m: int, n: int, k: int, num_groups: int, num_sms: int,
'B': is_tma_multicast_legal(m, best_block_m, 2, num_sms),
}
for i in ('A', 'B') if best_block_m > best_block_n else ('B', 'A'):
if m >= 1024 and is_multicast_legal[i] and num_groups == 1:
if m >= 512 and is_multicast_legal[i] and num_groups == 1:
best_tma_multicast_config = (2, i == 'A')
break