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

feat: implement sm90 megamoe phase6 combine

Browse Source
This commit is contained in:
Xinyi Liu
2026-06-18 17:40:49 +08:00
parent a0b3bb0017
commit 453fc7b046
3 changed files with 316 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
csrc/apis/sm90_mega.hpp
View File

@@ -217,6 +217,8 @@ static void fp8_mega_moe(
DG_HOST_ASSERT(num_experts_per_rank == num_experts_per_rank_);
DG_HOST_ASSERT(hidden == hidden_);
DG_HOST_ASSERT(intermediate_hidden_2 == 2 * intermediate_hidden);
DG_HOST_ASSERT(y.scalar_type() == torch::kBFloat16);
DG_HOST_ASSERT(y.dim() == 2 and y.size(1) == hidden and y.is_contiguous());
DG_HOST_ASSERT(l1_weights.is_contiguous() and l2_weights.is_contiguous());
// Check weight SF layout: float, natural MN-major, per-128-N and per-128-K.