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

refine

Browse Source
This commit is contained in:
kavioyu
2025-03-13 07:04:56 +00:00
parent 6e53c6613d
commit 094d0421ec
3 changed files with 4 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
deep_gemm/jit_kernels/gemm_bw.py
View File

@@ -75,7 +75,7 @@ def get_best_configs(m: int, n: int, k: int, num_groups: int, num_sms: int,
block_ms = (64 if m <= 64 else 128, )
else:
block_ms = (get_m_alignment_for_contiguous_layout(), )
block_ns = tuple(range(32, 129, 32))
block_ns = (128, )
fix_wave_saturate = lambda x: num_sms if x == 0 else x
get_num_waves = lambda bm, bn: (ceil_div(ceil_div(m, bm) * ceil_div(n, bn) * num_groups, num_sms) if bm else None)