DeepSeek/DeepGEMM
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepGEMM synced 2025-04-18 21:14:37 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Update get_best_configs

Browse Source 
a better strategy to choose config
This commit is contained in:
Liang 2025-02-27 23:18:52 +08:00 committed by GitHub
parent 461427ecd0
commit fbec9e5eee
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
deep_gemm/jit_kernels/gemm.py
View File

@ -79,10 +79,12 @@ def get_best_configs(m: int, n: int, k: int, num_groups: int, num_sms: int,
elif num_waves < best_num_waves: elif num_waves < best_num_waves:
success = True success = True
elif num_waves == best_num_waves: elif num_waves == best_num_waves:
div_n = bool(128 % block_n)
best_div_n = bool(128 % best_block_n)
# Check last wave utilization # Check last wave utilization
util = get_last_wave_util(block_m, block_n) util = get_last_wave_util(block_m, block_n)
best_util = get_last_wave_util(best_block_m, best_block_n) best_util = get_last_wave_util(best_block_m, best_block_n)
success = util > best_util or (util == best_util and (block_n >= best_block_n and block_m <= best_block_m)) success = util > best_util or (util == best_util and (block_m > best_block_m or block_m == best_block_m and (div_n < best_div_n or div_n == best_div_n and block_n < best_block_n)))
best_block_m, best_block_n = (block_m, block_n) if success else (best_block_m, best_block_n) best_block_m, best_block_n = (block_m, block_n) if success else (best_block_m, best_block_n)
assert best_block_m is not None and best_block_n is not None assert best_block_m is not None and best_block_n is not None