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

Use 1D TMA store

Browse Source
This commit is contained in:
Chenggang Zhao
2025-04-11 10:42:01 +08:00
parent a77009cb14
commit 8041ed7164
5 changed files with 31 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
deep_gemm/jit_kernels/gemm.py
View File

@@ -28,10 +28,9 @@ using gemm_t = Gemm<N, K, BLOCK_M, BLOCK_N, BLOCK_K, BLOCK_N_PADDING, 1, kNumSta
auto tma_a_desc = gemm_t::make_2d_tma_a_desc(lhs, m);
auto tma_b_desc = gemm_t::make_2d_tma_b_desc(rhs);
auto tma_scales_a_desc = gemm_t::make_2d_tma_scales_a_desc(lhs_scales, m);
auto tma_d_desc = gemm_t::make_3d_tma_d_desc(out, m);
gemm_t::run(out, rhs_scales, nullptr,
m,
tma_a_desc, tma_b_desc, tma_scales_a_desc, tma_d_desc,
tma_a_desc, tma_b_desc, tma_scales_a_desc,
stream, num_sms, smem_size);
"""