chenchenghao/DeepGEMM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: implement sm90 megamoe phase4 l1 epilogue

Browse Source
This commit is contained in:
Xinyi Liu
2026-06-18 01:09:45 +08:00
parent 2bb1756787
commit f3553f976c
6 changed files with 330 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp
View File

@@ -135,6 +135,7 @@ static void sm90_fp8_mega_moe(
// Make tensormap
constexpr int kGranK = 128;
constexpr int kL2ActsGranK = 64;
const auto tensor_map_l1_acts = make_tma_2d_desc(l1_acts,
hidden, config.num_max_pool_tokens,
config.block_k, config.block_m,
@@ -162,7 +163,7 @@ static void sm90_fp8_mega_moe(
128);
const auto tensor_map_l2_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l2_acts_sf,
config.num_padded_sf_pool_tokens, intermediate_hidden,
config.block_m, kGranK,
config.block_m, kL2ActsGranK,
1, 0);
const auto tensor_map_l2_weights = make_tma_2d_desc(l2_weights,
intermediate_hidden, num_experts_per_rank * hidden,