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

feat: add sm90 megamoe phase1 interfaces

Browse Source
This commit is contained in:
Xinyi Liu
2026-06-17 23:54:49 +08:00
parent 062cb160cf
commit 1dafd0f173
12 changed files with 1520 additions and 130 deletions
Download Patch File Download Diff File Expand all files Collapse all files

58
deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh Normal file
View File

@@ -0,0 +1,58 @@
#pragma once
#include <cutlass/cutlass.h>
#include <deep_gemm/common/exception.cuh>
#include <deep_gemm/common/tma_utils.cuh>
#include <deep_gemm/common/types.cuh>
#include <deep_gemm/layout/mega_moe.cuh>
#include <deep_gemm/layout/sym_buffer.cuh>
#include <deep_gemm/scheduler/sm90_mega_moe.cuh>
namespace deep_gemm {
template <
uint32_t kNumMaxTokensPerRank,
uint32_t kHidden, uint32_t kIntermediateHidden,
uint32_t kNumExperts, uint32_t kNumTopk,
uint32_t kNumExpertsPerWave,
uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K,
uint32_t kStoreBlockM,
uint32_t kNumMaxPoolTokens,
uint32_t kNumPaddedSFPoolTokens,
uint32_t kNumStages,
uint32_t kNumDispatchThreads, uint32_t kNumTMAThreads,
uint32_t kNumMathThreads,
bool kCooperativeMode, bool kUseNMajorL2,
uint32_t kNumSMs, uint32_t kNumRanks,
float kActivationClamp, bool kFastMath,
uint32_t L1_SHAPE_N = kIntermediateHidden * 2,
uint32_t L1_SHAPE_K = kHidden,
uint32_t L2_SHAPE_N = kHidden,
uint32_t L2_SHAPE_K = kIntermediateHidden,
uint32_t kNumThreads = kNumDispatchThreads + kNumTMAThreads + kNumMathThreads,
uint32_t kNumExpertsPerRank = kNumExperts / kNumRanks>
CUTLASS_GLOBAL __launch_bounds__(kNumThreads, 1) void
sm90_fp8_mega_moe_impl(void* y,
int* cumulative_local_expert_recv_stats,
const uint32_t num_tokens,
const __grid_constant__ layout::SymBuffer<kNumRanks> sym_buffer,
const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts,
const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts_sf,
const __grid_constant__ cute::TmaDescriptor tensor_map_l1_weights,
const void* l1_weights_sf,
const __grid_constant__ cute::TmaDescriptor tensor_map_l1_output,
const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts,
const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts_sf,
const __grid_constant__ cute::TmaDescriptor tensor_map_l2_weights,
const void* l2_weights_sf) {
DG_STATIC_ASSERT(kNumThreads == 384, "SM90 MegaMoE expects 384 threads");
DG_STATIC_ASSERT(BLOCK_N == 128, "SM90 MegaMoE expects BLOCK_N=128");
DG_STATIC_ASSERT(BLOCK_K == 128, "SM90 MegaMoE expects BLOCK_K=128");
DG_STATIC_ASSERT(kNumExperts % kNumRanks == 0, "Invalid number of experts or ranks");
// Phase 1 only validates the host/JIT/API path and launches an empty kernel.
return;
}
} // namespace deep_gemm

20
deep_gemm/include/deep_gemm/layout/mega_moe.cuh
View File

@@ -108,46 +108,46 @@ struct Workspace {
static constexpr uint32_t kNumMaxGridSyncCounters = 4;
template <uint32_t kIndex = 0>
CUTLASS_DEVICE
CUTLASS_HOST_DEVICE
uint32_t* get_grid_sync_count_ptr() const {
DG_STATIC_ASSERT(kIndex < kNumMaxGridSyncCounters, "Grid sync index out of bounds");
return static_cast<uint32_t*>(base) + kIndex;
}
CUTLASS_DEVICE
CUTLASS_HOST_DEVICE
uint32_t* get_nvl_barrier_counter_ptr() const {
return static_cast<uint32_t*>(base) + kNumMaxGridSyncCounters;
}
CUTLASS_DEVICE
CUTLASS_HOST_DEVICE
int* get_nvl_barrier_signal_ptr(const uint32_t& phase) const {
// NOTES: the signal is signed, as we may minus
return math::advance_ptr<int>(base, (kNumMaxGridSyncCounters + 1) * sizeof(uint32_t) + phase * sizeof(int));
}
CUTLASS_DEVICE
CUTLASS_HOST_DEVICE
uint64_t* get_expert_send_count_ptr(const uint32_t& expert_idx = 0) const {
return math::advance_ptr<uint64_t>(base, kNumBarrierSignalBytes) + expert_idx;
}
CUTLASS_DEVICE
CUTLASS_HOST_DEVICE
uint64_t* get_expert_recv_count_ptr(
const uint32_t& rank_idx = 0, const uint32_t& expert_idx = 0) const {
return get_expert_send_count_ptr(num_experts) + rank_idx * num_experts_per_rank + expert_idx;
}
CUTLASS_DEVICE
CUTLASS_HOST_DEVICE
uint64_t* get_expert_recv_count_sum_ptr(const uint32_t& expert_idx = 0) const {
return get_expert_send_count_ptr(num_experts * 2) + expert_idx;
}
CUTLASS_DEVICE
CUTLASS_HOST_DEVICE
uint32_t* get_l1_arrival_count_ptr(const uint32_t& pool_block_idx = 0) const {
const auto base = get_expert_recv_count_sum_ptr(num_experts_per_rank);
return reinterpret_cast<uint32_t*>(base) + pool_block_idx;
}
CUTLASS_DEVICE
CUTLASS_HOST_DEVICE
uint64_t* get_l2_arrival_mask_ptr(const uint32_t& pool_block_idx = 0) const {
// Pad L1 entry count to even so that the `l2_arrival_mask` is 8-byte aligned
const auto base = get_l1_arrival_count_ptr(math::align(num_max_pool_blocks, 2u));
@@ -155,7 +155,7 @@ struct Workspace {
}
// For dispatch pulling
CUTLASS_DEVICE
CUTLASS_HOST_DEVICE
uint32_t* get_src_token_topk_idx_ptr(
const uint32_t& expert_idx = 0, const uint32_t& rank_idx = 0, const uint32_t& token_idx = 0) const {
const auto base = get_l2_arrival_mask_ptr(num_max_pool_blocks);
@@ -165,7 +165,7 @@ struct Workspace {
}
// For combine usages
CUTLASS_DEVICE
CUTLASS_HOST_DEVICE
TokenSrcMetadata* get_token_src_metadata_ptr(const uint32_t& pool_token_idx = 0) const {
const auto base = reinterpret_cast<TokenSrcMetadata*>(get_src_token_topk_idx_ptr(num_experts_per_rank));
return base + pool_token_idx;

199
deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh Normal file
View File

@@ -0,0 +1,199 @@
#pragma once
#include <deep_gemm/common/cute_tie.cuh>
#include <deep_gemm/common/math.cuh>
#include <deep_gemm/common/types.cuh>
#include <deep_gemm/layout/mega_moe.cuh>
#include <deep_gemm/ptx/ld_st.cuh>
#include <deep_gemm/ptx/utils.cuh>
namespace deep_gemm::sched {
enum class SM90BlockPhase {
None = 0,
Linear1 = 1,
Linear2 = 2
};
template <uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K,
uint32_t L1_SHAPE_N, uint32_t L1_SHAPE_K,
uint32_t L2_SHAPE_N, uint32_t L2_SHAPE_K,
uint32_t kNumExpertsPerRank,
uint32_t kNumExpertsPerWave,
uint32_t kNumSMs, uint32_t kNumRanks,
bool kUseNMajorL2,
uint32_t kNumExpertsPerLane = math::constexpr_ceil_div(kNumExpertsPerRank, 32u),
uint32_t kNumL1BlockNs = L1_SHAPE_N / BLOCK_N,
uint32_t kNumL2BlockNs = L2_SHAPE_N / BLOCK_N,
uint32_t kNumL1BlockKs = L1_SHAPE_K / BLOCK_K,
uint32_t kNumL2BlockKs = L2_SHAPE_K / BLOCK_K>
struct SM90MegaMoEScheduler {
DG_STATIC_ASSERT(L1_SHAPE_N % BLOCK_N == 0, "Invalid shape");
DG_STATIC_ASSERT(L2_SHAPE_N % BLOCK_N == 0, "Invalid shape");
DG_STATIC_ASSERT(L1_SHAPE_K % BLOCK_K == 0, "Invalid shape");
DG_STATIC_ASSERT(L2_SHAPE_K % BLOCK_K == 0, "Invalid shape");
DG_STATIC_ASSERT(kNumExpertsPerRank % kNumExpertsPerWave == 0, "Invalid wave config");
const layout::Workspace& workspace;
SM90BlockPhase next_phase = SM90BlockPhase::Linear1;
uint32_t current_local_expert_idx = 0;
uint32_t current_num_tokens = 0;
uint32_t current_pool_block_offset = 0;
uint32_t block_idx = 0;
uint32_t m_block_idx = 0;
uint32_t n_block_idx = 0;
uint32_t stored_num_tokens_per_expert[kNumExpertsPerLane] = {};
CUTLASS_DEVICE explicit SM90MegaMoEScheduler(const layout::Workspace& workspace): workspace(workspace) {
block_idx = blockIdx.x;
}
CUTLASS_DEVICE uint32_t get_wave_expert_end_idx() const {
return math::align(current_local_expert_idx + 1, kNumExpertsPerWave);
}
CUTLASS_DEVICE uint32_t get_num_tokens(const uint32_t& expert_idx) const {
uint32_t valid_value;
#pragma unroll
for (uint32_t i = 0; i < kNumExpertsPerLane; ++ i) {
valid_value = (expert_idx == i * 32 + ptx::get_lane_idx()) ?
stored_num_tokens_per_expert[i] : valid_value;
}
return ptx::exchange(valid_value, expert_idx % 32);
}
CUTLASS_DEVICE uint32_t get_pool_block_offset(const uint32_t& expert_idx) {
uint32_t num_blocks = 0;
#pragma unroll
for (uint32_t i = 0; i < kNumExpertsPerLane; ++ i) {
if (i * 32 + ptx::get_lane_idx() < expert_idx)
num_blocks += math::ceil_div(stored_num_tokens_per_expert[i], BLOCK_M);
}
return __reduce_add_sync(0xffffffff, num_blocks);
}
CUTLASS_DEVICE void advance_expert_idx() {
current_pool_block_offset += get_current_num_m_blocks();
current_local_expert_idx += 1;
current_num_tokens = get_num_tokens(current_local_expert_idx);
}
CUTLASS_DEVICE void set_expert_idx(const uint32_t& expert_idx) {
current_local_expert_idx = expert_idx;
current_num_tokens = get_num_tokens(expert_idx);
current_pool_block_offset = get_pool_block_offset(expert_idx);
}
CUTLASS_DEVICE uint32_t get_current_pool_block_offset() const {
return current_pool_block_offset;
}
CUTLASS_DEVICE uint32_t get_current_num_m_blocks() const {
return math::ceil_div(current_num_tokens, BLOCK_M);
}
CUTLASS_DEVICE uint32_t get_valid_m() const {
return cute::min(current_num_tokens - m_block_idx * BLOCK_M, BLOCK_M);
}
CUTLASS_DEVICE bool fetch_next_l1_block() {
const auto wave_end_expert_idx = get_wave_expert_end_idx();
while (current_local_expert_idx < wave_end_expert_idx) {
const auto num_m_blocks = get_current_num_m_blocks();
m_block_idx = block_idx / kNumL1BlockNs;
if (m_block_idx < num_m_blocks)
return true;
block_idx -= num_m_blocks * kNumL1BlockNs;
advance_expert_idx();
}
return false;
}
CUTLASS_DEVICE bool fetch_next_l2_block() {
const auto wave_end_expert_idx = get_wave_expert_end_idx();
while (current_local_expert_idx < wave_end_expert_idx) {
const auto num_m_blocks = get_current_num_m_blocks();
if (block_idx < num_m_blocks * kNumL2BlockNs) {
if constexpr (kUseNMajorL2) {
n_block_idx = block_idx / num_m_blocks;
m_block_idx = block_idx - n_block_idx * num_m_blocks;
} else {
m_block_idx = block_idx / kNumL2BlockNs;
n_block_idx = block_idx - m_block_idx * kNumL2BlockNs;
}
return true;
}
block_idx -= num_m_blocks * kNumL2BlockNs;
advance_expert_idx();
}
return false;
}
CUTLASS_DEVICE cute::tuple<SM90BlockPhase, uint32_t, uint32_t, uint32_t> get_next_block() {
while (true) {
if (current_local_expert_idx >= kNumExpertsPerRank)
break;
if (next_phase == SM90BlockPhase::Linear1) {
if (fetch_next_l1_block()) {
n_block_idx = block_idx - m_block_idx * kNumL1BlockNs;
block_idx += kNumSMs;
return {SM90BlockPhase::Linear1, current_local_expert_idx, m_block_idx, n_block_idx};
} else {
next_phase = SM90BlockPhase::Linear2;
set_expert_idx(math::align<uint32_t, false>(current_local_expert_idx - 1, kNumExpertsPerWave));
}
} else {
if (fetch_next_l2_block()) {
if constexpr (not kUseNMajorL2) {
n_block_idx = block_idx - m_block_idx * kNumL2BlockNs;
}
block_idx += kNumSMs;
return {SM90BlockPhase::Linear2, current_local_expert_idx, m_block_idx, n_block_idx};
} else {
next_phase = SM90BlockPhase::Linear1;
}
}
}
return {SM90BlockPhase::None, 0, 0, 0};
}
CUTLASS_DEVICE void fetch_expert_recv_count() {
#pragma unroll
for (uint32_t i = 0; i < kNumExpertsPerLane; ++ i) {
const auto expert_idx = i * 32 + ptx::get_lane_idx();
uint64_t value = 0;
if (expert_idx < kNumExpertsPerRank) {
do {
value = ptx::ld_volatile(workspace.get_expert_recv_count_sum_ptr(expert_idx));
} while (static_cast<uint32_t>(value >> 32) != kNumSMs * kNumRanks);
}
stored_num_tokens_per_expert[i] = static_cast<uint32_t>(value);
}
__syncwarp();
}
template <typename Func>
CUTLASS_DEVICE void for_each_block(Func&& func) {
fetch_expert_recv_count();
set_expert_idx(0);
while (true) {
CUTE_TIE_DECL(get_next_block(), block_phase, current_local_expert_idx, m_block_idx, n_block_idx);
if (block_phase == SM90BlockPhase::None)
break;
func(block_phase, current_local_expert_idx,
block_phase == SM90BlockPhase::Linear2 ? kNumL2BlockKs : kNumL1BlockKs,
m_block_idx, n_block_idx);
}
}
};
} // namespace deep_gemm::sched