DeepSeek/DeepGEMM
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepGEMM synced 2025-05-31 18:48:16 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add extra TMA checks

Browse Source
This commit is contained in:
Chenggang Zhao 2025-02-27 18:20:57 +08:00
parent ca13ce0fab
commit 6da94d2d36
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
deep_gemm/jit_kernels/m_grouped_gemm.py
View File

@ -160,6 +160,11 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_masked(lhs: Tuple[torch.Tensor, torch.Tensor]
global includes, template global includes, template
num_sms = get_num_sms() num_sms = get_num_sms()
block_m, block_n, num_stages, num_tma_multicast, smem_size = get_best_configs(expected_m, n, k, num_groups, num_sms) block_m, block_n, num_stages, num_tma_multicast, smem_size = get_best_configs(expected_m, n, k, num_groups, num_sms)
# Extra checks for TMA store
if num_groups > 1 and m > block_m:
assert m % block_m == 0, f'For masked grouped GEMM, shape M should be multiple of the block M (current block M: {block_m})'
args = (lhs, lhs_scales, rhs, rhs_scales, out, args = (lhs, lhs_scales, rhs, rhs_scales, out,
masked_m, m, masked_m, m,
torch.cuda.current_stream(), num_sms, smem_size) torch.cuda.current_stream(), num_sms, smem_size)