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

41
deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh
View File

@@ -94,6 +94,7 @@ sm90_fp8_mega_moe_impl(void* y,
constexpr uint32_t L1_OUT_BLOCK_N = BLOCK_N / 2;
constexpr uint32_t kL2ActsGranK = 64;
constexpr uint32_t kMathBarrierIdx = 2;
constexpr uint32_t kDispatchWithMathBarrierIdx = 3;
DG_STATIC_ASSERT(kNumTokensPerWarp > 0, "Invalid number of top-k experts");
DG_STATIC_ASSERT(kNumPaddedSFPoolTokens % SF_BLOCK_M == 0, "Invalid padded SF pool size");
DG_STATIC_ASSERT(BLOCK_N == WGMMA::N and BLOCK_K % WGMMA::K == 0, "Invalid WGMMA tile shape");
@@ -118,6 +119,7 @@ sm90_fp8_mega_moe_impl(void* y,
constexpr uint32_t kDispatchGridSyncIndex = 0;
constexpr uint32_t kAfterWorkspaceCleanBarrierTag = 1;
constexpr uint32_t kBeforeDispatchPullBarrierTag = 2;
constexpr uint32_t kBeforeCombineReduceBarrierTag = 3;
const auto dispatch_sync = []() {
ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);
};
@@ -454,6 +456,9 @@ sm90_fp8_mega_moe_impl(void* y,
}
__syncwarp();
}
if constexpr (BLOCK_M == 128)
ptx::sync_unaligned(kNumDispatchThreads + kNumMathThreads, kDispatchWithMathBarrierIdx);
} else if (thread_idx < kNumDispatchThreads + kNumTMAThreads) {
if constexpr (BLOCK_M == 128) {
if (warp_idx == kNumDispatchWarps) {
@@ -525,6 +530,7 @@ sm90_fp8_mega_moe_impl(void* y,
}
} else {
if constexpr (BLOCK_M == 128) {
const uint32_t math_thread_idx = thread_idx - kNumDispatchThreads - kNumTMAThreads;
const uint32_t math_warp_idx = warp_idx - kMathWarpStart;
const uint32_t math_wg_idx = math_warp_idx / 4;
const uint32_t warp_idx_in_wg = math_warp_idx % 4;
@@ -707,7 +713,7 @@ sm90_fp8_mega_moe_impl(void* y,
*l1_topk_weights_buffer.get_data_buffer(pool_token_idx_1).get_base_ptr<float>() : 0.0f;
const auto apply_swiglu = [&](float gate, float up, const float& topk_weight) {
if constexpr (kActivationClamp != cute::numeric_limits<float>::infinity()) {
if constexpr (kActivationClamp < 1.0e30f) {
gate = fminf(fmaxf(gate, -kActivationClamp), kActivationClamp);
up = fminf(fmaxf(up, -kActivationClamp), kActivationClamp);
}
@@ -786,6 +792,39 @@ sm90_fp8_mega_moe_impl(void* y,
}
}
});
ptx::sync_unaligned(kNumDispatchThreads + kNumMathThreads, kDispatchWithMathBarrierIdx);
__threadfence_system();
const auto math_sync = []() {
ptx::sync_aligned(kNumMathThreads, kMathBarrierIdx);
};
comm::nvlink_barrier<kNumRanks, kNumSMs, kNumMathThreads,
kDispatchGridSyncIndex, kBeforeCombineReduceBarrierTag>(
workspace, sym_buffer, sm_idx, math_thread_idx, math_sync);
auto y_ptr = reinterpret_cast<nv_bfloat16*>(y);
const uint64_t num_output_values = static_cast<uint64_t>(num_tokens) * kHidden;
const uint64_t output_stride = static_cast<uint64_t>(kNumSMs) * kNumMathThreads;
for (uint64_t elem_idx = static_cast<uint64_t>(sm_idx) * kNumMathThreads + math_thread_idx;
elem_idx < num_output_values;
elem_idx += output_stride) {
const uint32_t token_idx = static_cast<uint32_t>(elem_idx / kHidden);
const uint32_t hidden_idx = static_cast<uint32_t>(elem_idx - static_cast<uint64_t>(token_idx) * kHidden);
float reduced = 0.0f;
#pragma unroll 1
for (uint32_t topk_slot = 0; topk_slot < kNumTopk; ++ topk_slot) {
const auto expert_idx = static_cast<int64_t>(__ldg(
input_topk_idx_buffer.get_base_ptr<int64_t>() + token_idx * kNumTopk + topk_slot));
if (expert_idx >= 0 and expert_idx < static_cast<int64_t>(kNumExperts)) {
const auto src_ptr = combine_token_buffer.get_rank_buffer(topk_slot)
.get_data_buffer(token_idx).get_base_ptr<nv_bfloat16>();
reduced += __bfloat162float(src_ptr[hidden_idx]);
}
}
y_ptr[token_idx * kHidden + hidden_idx] = __float2bfloat16(reduced);
}
}
}
#endif