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

feat: implement sm90 megamoe phase3 l1 wgmma

Browse Source
This commit is contained in:
Xinyi Liu
2026-06-18 00:46:41 +08:00
parent ed6067c7a5
commit 842124b1c5
5 changed files with 479 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,16 +1,22 @@
#pragma once
#include <cutlass/arch/barrier.h>
#include <cutlass/cutlass.h>
#include <deep_gemm/common/math.cuh>
#include <deep_gemm/common/exception.cuh>
#include <deep_gemm/common/tma_copy.cuh>
#include <deep_gemm/common/tma_utils.cuh>
#include <deep_gemm/common/types.cuh>
#include <deep_gemm/common/utils.cuh>
#include <deep_gemm/comm/barrier.cuh>
#include <deep_gemm/layout/mega_moe.cuh>
#include <deep_gemm/layout/sym_buffer.cuh>
#include <deep_gemm/mma/sm90.cuh>
#include <deep_gemm/ptx/ld_st.cuh>
#include <deep_gemm/ptx/tma.cuh>
#include <deep_gemm/ptx/utils.cuh>
#include <deep_gemm/ptx/wgmma.cuh>
#include <deep_gemm/scheduler/sm90_mega_moe.cuh>
namespace deep_gemm {
@@ -43,44 +49,63 @@ sm90_fp8_mega_moe_impl(void* y,
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) {
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,
void* l1_accum_debug) {
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");
#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900)) or defined(__CLION_IDE__)
DG_STATIC_ASSERT(kNumDispatchThreads == 64, "SM90 dispatch-only path expects 64 dispatch threads");
DG_STATIC_ASSERT(kNumDispatchThreads == 64, "SM90 MegaMoE expects 64 dispatch threads");
DG_STATIC_ASSERT(kNumTMAThreads == 64, "SM90 MegaMoE expects 64 TMA threads");
DG_STATIC_ASSERT(kNumMathThreads == 256, "SM90 MegaMoE expects 256 math threads");
DG_STATIC_ASSERT(kNumTopk <= 32, "Invalid number of top-k experts");
DG_STATIC_ASSERT(kHidden % 128 == 0, "SM90 activation SF expects per-128K groups");
DG_STATIC_ASSERT(kHidden % sizeof(uint4) == 0, "Token copy expects 16-byte alignment");
using Barrier = cutlass::arch::ClusterTransactionBarrier;
using a_dtype_t = __nv_fp8_e4m3;
using b_dtype_t = __nv_fp8_e4m3;
using WGMMA = typename mma::sm90::FP8MMASelector<BLOCK_N>::type;
constexpr uint32_t kNumDispatchWarps = kNumDispatchThreads / 32;
constexpr uint32_t kNumTMAWarps = kNumTMAThreads / 32;
constexpr uint32_t kNumMathWarps = kNumMathThreads / 32;
constexpr uint32_t kMathWarpStart = kNumDispatchWarps + kNumTMAWarps;
constexpr uint32_t kNumTokensPerWarp = 32 / kNumTopk;
constexpr uint32_t kNumActivateLanes = kNumTokensPerWarp * kNumTopk;
constexpr uint32_t SF_BLOCK_M = math::constexpr_align(BLOCK_M, 128u);
constexpr uint32_t kNumMaxPoolBlocks = kNumMaxPoolTokens / layout::kMinCandidateBlockM;
constexpr uint32_t kNumTokenUint4 = kHidden / sizeof(uint4);
constexpr uint32_t kNumSFValues = kHidden / 128;
constexpr uint32_t kSwizzleAMode = BLOCK_K * sizeof(a_dtype_t);
constexpr uint32_t kSwizzleBMode = BLOCK_K * sizeof(b_dtype_t);
constexpr uint32_t kNumL1WeightSFGroupsN = L1_SHAPE_N / 128;
constexpr uint32_t kNumL1WeightSFGroupsK = L1_SHAPE_K / 128;
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");
// Only the first two warps participate in Phase 2 dispatch. TMA/math warps stay idle.
const uint32_t thread_idx = threadIdx.x;
if (thread_idx >= kNumDispatchThreads)
return;
const uint32_t sm_idx = blockIdx.x;
const uint32_t warp_idx = thread_idx / 32;
const uint32_t lane_idx = ptx::get_lane_idx();
constexpr uint32_t kDispatchBarrierIdx = 0;
if (warp_idx == kNumDispatchWarps and cute::elect_one_sync()) {
cute::prefetch_tma_descriptor(&tensor_map_l1_acts);
cute::prefetch_tma_descriptor(&tensor_map_l1_acts_sf);
cute::prefetch_tma_descriptor(&tensor_map_l1_weights);
}
__syncwarp();
constexpr uint32_t kDispatchBarrierIdx = 1;
constexpr uint32_t kDispatchGridSyncIndex = 0;
constexpr uint32_t kAfterWorkspaceCleanBarrierTag = 1;
constexpr uint32_t kBeforeDispatchPullBarrierTag = 2;
@@ -131,23 +156,78 @@ sm90_fp8_mega_moe_impl(void* y,
constexpr uint32_t kSharedMemoryAlignment = 1024;
extern __shared__ __align__(kSharedMemoryAlignment) uint8_t smem_buffer[];
const auto smem_expert_count = reinterpret_cast<uint32_t*>(smem_buffer);
// Clean local dispatch workspace from any previous call. Barrier state is intentionally
// persistent and must not be reset because grid/NVLink barriers use phase counters.
for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads)
*workspace.get_expert_send_count_ptr(i) = 0;
for (uint32_t i = thread_idx; i < kNumRanks * kNumExpertsPerRank; i += kNumDispatchThreads)
*workspace.get_expert_recv_count_ptr(i / kNumExpertsPerRank, i % kNumExpertsPerRank) = 0;
for (uint32_t i = thread_idx; i < kNumExpertsPerRank; i += kNumDispatchThreads)
*workspace.get_expert_recv_count_sum_ptr(i) = 0;
for (uint32_t i = thread_idx; i < kNumMaxPoolBlocks; i += kNumDispatchThreads) {
*workspace.get_l1_arrival_count_ptr(i) = 0;
*workspace.get_l2_arrival_mask_ptr(i) = 0;
constexpr uint32_t SMEM_EXPERT_COUNT_SIZE =
math::constexpr_align<uint32_t>(kNumExperts * sizeof(uint32_t), kSharedMemoryAlignment);
constexpr uint32_t SMEM_SEND_BUFFER_SIZE =
math::constexpr_align<uint32_t>(fp8_token_layout.get_num_bytes() * kNumDispatchWarps, kSharedMemoryAlignment);
constexpr uint32_t SMEM_CD_L1_SIZE =
(kCooperativeMode ? 2u : 1u) * kStoreBlockM * (BLOCK_N / 2) * 2;
constexpr uint32_t SMEM_CD_L2_SIZE = BLOCK_M * BLOCK_N * static_cast<uint32_t>(sizeof(nv_bfloat16));
constexpr uint32_t SMEM_CD_SIZE = SMEM_CD_L1_SIZE > SMEM_CD_L2_SIZE ? SMEM_CD_L1_SIZE : SMEM_CD_L2_SIZE;
constexpr uint32_t SMEM_A_SIZE_PER_STAGE = BLOCK_M * BLOCK_K * sizeof(a_dtype_t);
constexpr uint32_t SMEM_B_SIZE_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(b_dtype_t);
constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = SF_BLOCK_M * sizeof(float);
DG_STATIC_ASSERT(SMEM_A_SIZE_PER_STAGE % kSharedMemoryAlignment == 0, "Invalid A smem alignment");
DG_STATIC_ASSERT(SMEM_B_SIZE_PER_STAGE % kSharedMemoryAlignment == 0, "Invalid B smem alignment");
const auto smem_expert_count = reinterpret_cast<uint32_t*>(smem_buffer);
auto smem_gemm_base = math::advance_ptr<uint8_t>(
smem_buffer, SMEM_EXPERT_COUNT_SIZE + SMEM_SEND_BUFFER_SIZE);
auto smem_a = utils::PatternVisitor([=](const uint32_t& i) {
return math::advance_ptr<a_dtype_t>(smem_gemm_base, SMEM_CD_SIZE + i * SMEM_A_SIZE_PER_STAGE);
});
auto smem_b = utils::PatternVisitor([=](const uint32_t& i) {
return math::advance_ptr<b_dtype_t>(smem_gemm_base, SMEM_CD_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE);
});
auto smem_sfa = utils::PatternVisitor([=](const uint32_t& i) {
return math::advance_ptr<float>(smem_gemm_base,
SMEM_CD_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE) + i * SMEM_SFA_SIZE_PER_STAGE);
});
auto barrier_start_ptr = reinterpret_cast<Barrier*>(math::advance_ptr<uint8_t>(
smem_gemm_base,
SMEM_CD_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE)));
auto full_barriers = utils::PatternVisitor([=](const uint32_t& i) {
return barrier_start_ptr + i;
});
auto empty_barriers = utils::PatternVisitor([=](const uint32_t& i) {
return barrier_start_ptr + kNumStages + i;
});
if constexpr (BLOCK_M == 128) {
if (sm_idx == 0) {
auto debug_ptr = reinterpret_cast<float*>(l1_accum_debug);
for (uint32_t i = thread_idx; i < 128 * 128; i += kNumThreads)
debug_ptr[i] = 0.0f;
}
if (warp_idx == kNumDispatchWarps + 1 and cute::elect_one_sync()) {
#pragma unroll
for (uint32_t i = 0; i < kNumStages; ++ i) {
full_barriers[i]->init(1);
empty_barriers[i]->init(kNumMathWarps);
}
cutlass::arch::fence_barrier_init();
}
__syncthreads();
}
comm::nvlink_barrier<kNumRanks, kNumSMs, kNumDispatchThreads,
kDispatchGridSyncIndex, kAfterWorkspaceCleanBarrierTag>(
workspace, sym_buffer, sm_idx, thread_idx, dispatch_sync);
if (thread_idx < kNumDispatchThreads) {
// Clean local dispatch workspace from any previous call. Barrier state is intentionally
// persistent and must not be reset because grid/NVLink barriers use phase counters.
for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads)
*workspace.get_expert_send_count_ptr(i) = 0;
for (uint32_t i = thread_idx; i < kNumRanks * kNumExpertsPerRank; i += kNumDispatchThreads)
*workspace.get_expert_recv_count_ptr(i / kNumExpertsPerRank, i % kNumExpertsPerRank) = 0;
for (uint32_t i = thread_idx; i < kNumExpertsPerRank; i += kNumDispatchThreads)
*workspace.get_expert_recv_count_sum_ptr(i) = 0;
for (uint32_t i = thread_idx; i < kNumMaxPoolBlocks; i += kNumDispatchThreads) {
*workspace.get_l1_arrival_count_ptr(i) = 0;
*workspace.get_l2_arrival_mask_ptr(i) = 0;
}
comm::nvlink_barrier<kNumRanks, kNumSMs, kNumDispatchThreads,
kDispatchGridSyncIndex, kAfterWorkspaceCleanBarrierTag>(
workspace, sym_buffer, sm_idx, thread_idx, dispatch_sync);
for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads)
smem_expert_count[i] = 0;
@@ -358,6 +438,181 @@ sm90_fp8_mega_moe_impl(void* y,
}
__syncwarp();
}
} else if (thread_idx < kNumDispatchThreads + kNumTMAThreads) {
if constexpr (BLOCK_M == 128) {
if (warp_idx == kNumDispatchWarps) {
auto scheduler = sched::SM90MegaMoEScheduler<
BLOCK_M, BLOCK_N, BLOCK_K,
L1_SHAPE_N, L1_SHAPE_K,
L2_SHAPE_N, L2_SHAPE_K,
kNumExpertsPerRank,
kNumExpertsPerWave,
kNumSMs, kNumRanks,
kUseNMajorL2>(workspace);
uint32_t stage_idx = 0, phase = 0;
auto advance_pipeline = [&](uint32_t& k_block_idx) {
++ k_block_idx;
stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1;
phase ^= stage_idx == 0;
};
scheduler.for_each_block([&](const sched::SM90BlockPhase& block_phase,
const uint32_t& local_expert_idx,
const uint32_t& num_k_blocks,
const uint32_t& m_block_idx,
const uint32_t& n_block_idx) {
if (block_phase != sched::SM90BlockPhase::Linear1)
return;
const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx;
const uint32_t valid_m = scheduler.get_valid_m();
const auto arrival_ptr = workspace.get_l1_arrival_count_ptr(pool_block_idx);
while (ptx::ld_acq(arrival_ptr) != valid_m);
#pragma unroll 1
for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) {
empty_barriers[stage_idx]->wait(phase ^ 1);
if (cute::elect_one_sync()) {
auto& full_barrier = *full_barriers[stage_idx];
const uint32_t k_idx = k_block_idx * BLOCK_K;
const uint32_t m_idx = pool_block_idx * BLOCK_M;
const uint32_t n_idx = local_expert_idx * L1_SHAPE_N + n_block_idx * BLOCK_N;
const uint32_t sfa_m_idx = pool_block_idx * SF_BLOCK_M;
tma::copy<BLOCK_K, BLOCK_M, kSwizzleAMode>(
&tensor_map_l1_acts, &full_barrier, smem_a[stage_idx], k_idx, m_idx);
tma::copy<SF_BLOCK_M, 1, 0>(
&tensor_map_l1_acts_sf, &full_barrier, smem_sfa[stage_idx], sfa_m_idx, k_block_idx);
tma::copy<BLOCK_K, BLOCK_N, kSwizzleBMode>(
&tensor_map_l1_weights, &full_barrier, smem_b[stage_idx], k_idx, n_idx);
full_barrier.arrive_and_expect_tx(
SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE);
}
__syncwarp();
}
});
}
}
} else {
if constexpr (BLOCK_M == 128) {
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;
const uint32_t row_idx = lane_idx / 4;
const uint32_t col_idx = lane_idx % 4;
const uint32_t r_0 = warp_idx_in_wg * 16 + row_idx;
const uint32_t r_1 = r_0 + 8;
const auto l1_weights_sf_ptr = reinterpret_cast<const float*>(l1_weights_sf);
auto l1_accum_debug_ptr = reinterpret_cast<float*>(l1_accum_debug);
const auto get_l1_weight_sf_group = [](const uint32_t& interleaved_n) {
constexpr uint32_t kInterleaveGran = 8;
const uint32_t pair_group = interleaved_n / (2 * kInterleaveGran);
const uint32_t offset = interleaved_n - pair_group * (2 * kInterleaveGran);
const uint32_t natural_n = offset < kInterleaveGran ?
pair_group * kInterleaveGran + offset :
(L1_SHAPE_N / 2) + pair_group * kInterleaveGran + offset - kInterleaveGran;
return natural_n / 128;
};
auto scheduler = sched::SM90MegaMoEScheduler<
BLOCK_M, BLOCK_N, BLOCK_K,
L1_SHAPE_N, L1_SHAPE_K,
L2_SHAPE_N, L2_SHAPE_K,
kNumExpertsPerRank,
kNumExpertsPerWave,
kNumSMs, kNumRanks,
kUseNMajorL2>(workspace);
uint32_t stage_idx = 0, phase = 0;
auto advance_pipeline = [&](uint32_t& k_block_idx) {
++ k_block_idx;
stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1;
phase ^= stage_idx == 0;
};
scheduler.for_each_block([&](const sched::SM90BlockPhase& block_phase,
const uint32_t& local_expert_idx,
const uint32_t& num_k_blocks,
const uint32_t& m_block_idx,
const uint32_t& n_block_idx) {
if (block_phase != sched::SM90BlockPhase::Linear1)
return;
const uint32_t valid_m = scheduler.get_valid_m();
float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum] = {0};
const auto empty_barrier_arrive = [&](const uint32_t& s) {
if (lane_idx == 0)
empty_barriers[s]->arrive();
};
#pragma unroll 1
for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) {
full_barriers[stage_idx]->wait(phase);
const float scale_a_0 = ptx::ld_shared(smem_sfa[stage_idx] + math_wg_idx * WGMMA::M + r_0);
const float scale_a_1 = ptx::ld_shared(smem_sfa[stage_idx] + math_wg_idx * WGMMA::M + r_1);
#pragma unroll
for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i)
ptx::warpgroup_fence_operand(accum[i]);
ptx::warpgroup_arrive();
#pragma unroll
for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) {
auto desc_a = mma::sm90::make_smem_desc(
smem_a[stage_idx] + math_wg_idx * WGMMA::M * BLOCK_K + k * WGMMA::K, 1);
auto desc_b = mma::sm90::make_smem_desc(
smem_b[stage_idx] + k * WGMMA::K, 1);
WGMMA::wgmma(desc_a, desc_b, accum, k);
}
ptx::warpgroup_commit_batch();
#pragma unroll
for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i)
ptx::warpgroup_fence_operand(accum[i]);
ptx::warpgroup_wait<0>();
empty_barrier_arrive(stage_idx);
#pragma unroll
for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
const uint32_t col_0 = i * 8 + col_idx * 2;
const uint32_t col_1 = col_0 + 1;
const uint32_t n_0 = n_block_idx * BLOCK_N + col_0;
const uint32_t n_1 = n_block_idx * BLOCK_N + col_1;
const auto sf_base = l1_weights_sf_ptr +
local_expert_idx * kNumL1WeightSFGroupsN * kNumL1WeightSFGroupsK;
const float scale_b_0 = __ldg(
sf_base + get_l1_weight_sf_group(n_0) * kNumL1WeightSFGroupsK + k_block_idx);
const float scale_b_1 = __ldg(
sf_base + get_l1_weight_sf_group(n_1) * kNumL1WeightSFGroupsK + k_block_idx);
final_accum[i * 4 + 0] += scale_a_0 * scale_b_0 * accum[i * 4 + 0];
final_accum[i * 4 + 1] += scale_a_0 * scale_b_1 * accum[i * 4 + 1];
final_accum[i * 4 + 2] += scale_a_1 * scale_b_0 * accum[i * 4 + 2];
final_accum[i * 4 + 3] += scale_a_1 * scale_b_1 * accum[i * 4 + 3];
}
}
if (local_expert_idx == 0 and m_block_idx == 0 and n_block_idx == 0) {
const uint32_t row_0 = math_wg_idx * WGMMA::M + r_0;
const uint32_t row_1 = math_wg_idx * WGMMA::M + r_1;
#pragma unroll
for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
const uint32_t col = i * 8 + col_idx * 2;
if (row_0 < valid_m) {
l1_accum_debug_ptr[row_0 * BLOCK_N + col + 0] = final_accum[i * 4 + 0];
l1_accum_debug_ptr[row_0 * BLOCK_N + col + 1] = final_accum[i * 4 + 1];
}
if (row_1 < valid_m) {
l1_accum_debug_ptr[row_1 * BLOCK_N + col + 0] = final_accum[i * 4 + 2];
l1_accum_debug_ptr[row_1 * BLOCK_N + col + 1] = final_accum[i * 4 + 3];
}
}
}
});
}
}
#endif
return;
}

5
deep_gemm/mega/__init__.py
View File

@@ -68,7 +68,8 @@ class SymmBuffer:
self.l2_acts, self.l2_acts_sf,
self.expert_recv_count_sum,
self.l1_arrival_count,
self.token_src_metadata) = buffer_views
self.token_src_metadata,
self.l1_accum_debug) = buffer_views
else:
(self.x, self.x_sf,
self.topk_idx, self.topk_weights,
@@ -78,6 +79,7 @@ class SymmBuffer:
self.expert_recv_count_sum = None
self.l1_arrival_count = None
self.token_src_metadata = None
self.l1_accum_debug = None
def destroy(self):
self.handle = None
@@ -95,6 +97,7 @@ class SymmBuffer:
self.expert_recv_count_sum = None
self.l1_arrival_count = None
self.token_src_metadata = None
self.l1_accum_debug = None
def get_symm_buffer_for_mega_moe(group: dist.ProcessGroup,