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

Fix TMA multicast bugs

Browse Source
This commit is contained in:
Chenggang Zhao
2025-04-07 14:34:42 +08:00
parent bff5724ded
commit b4ecf9c3ff
2 changed files with 6 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
deep_gemm/jit_kernels/m_grouped_gemm.py
View File

@@ -166,7 +166,7 @@ def m_grouped_gemm_fp8_fp8_bf16_nt_masked(lhs: Tuple[torch.Tensor, torch.Tensor]
# Auto-tuning with compilation
global includes, template
num_sms = get_num_sms()
num_sms, block_m, block_n, num_stages, tma_multicast_config, smem_size = get_best_configs(expected_m, n, k, num_groups, num_sms)
num_sms, block_m, block_n, num_stages, tma_multicast_config, smem_size = get_best_configs(expected_m, n, k, num_groups, num_sms, is_grouped_masked=True)
# Extra checks for TMA store
if num_groups > 1 and m > block_m: