biondizzle/DeepGEMM
1
0
Fork 0
Code Issues Pull Requests 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
indexing/main.cu
View File

@@ -20,10 +20,9 @@ int main() {
auto tma_a_desc = gemm_t::make_2d_tma_a_desc(reinterpret_cast<__nv_fp8_e4m3*>(0), m);
auto tma_b_desc = gemm_t::make_2d_tma_b_desc(reinterpret_cast<__nv_fp8_e4m3*>(0));
auto tma_scales_a_desc = gemm_t::make_2d_tma_scales_a_desc(reinterpret_cast<float*>(0), m);
auto tma_d_desc = gemm_t::make_3d_tma_d_desc(reinterpret_cast<nv_bfloat16*>(0), m);
gemm_t::run(nullptr, nullptr, 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,
nullptr, 132, 0);
return 0;
}