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

22
csrc/apis/sm90_mega.hpp
View File

@@ -15,7 +15,8 @@ using SM90MegaMoEBufferViews = std::tuple<
torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
torch::Tensor, torch::Tensor, torch::Tensor,
torch::Tensor, torch::Tensor,
torch::Tensor, torch::Tensor, torch::Tensor>;
torch::Tensor, torch::Tensor, torch::Tensor,
torch::Tensor>;
static int get_token_alignment_for_sm90_mega_moe() {
return layout::kLCMCandidateBlockM;
@@ -85,10 +86,16 @@ get_symm_buffer_size_for_sm90_mega_moe(
fp8_intermediate_sf_layout, 1, num_max_padded_sf_pool_tokens,
l2_token_buffer.get_end_ptr());
// Phase 3 verification output: one scaled FP32 L1 accumulator tile.
const auto l1_accum_debug_layout = layout::Data(128 * static_cast<int>(sizeof(float)), false);
const auto l1_accum_debug_buffer = layout::Buffer(
l1_accum_debug_layout, 1, 128,
l2_sf_buffer.get_end_ptr());
// Combine input buffer
const auto combine_token_buffer = layout::Buffer(
bf16_token_layout, num_topk, num_max_tokens_per_rank,
l2_sf_buffer.get_end_ptr());
l1_accum_debug_buffer.get_end_ptr());
// Check SF buffer requirements
DG_HOST_ASSERT(hidden % 128 == 0 and intermediate_hidden % 128 == 0);
@@ -149,10 +156,15 @@ get_symm_buffer_size_for_sm90_mega_moe(
reinterpret_cast<int32_t*>(runtime_workspace.get_token_src_metadata_ptr()),
{num_max_pool_tokens, 3},
torch::TensorOptions().dtype(torch::kInt).device(buffer.device()));
auto l1_accum_debug = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(l1_accum_debug_buffer.base)),
{128, 128},
torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device()));
return std::make_tuple(x, x_sf, topk_idx, topk_weights,
l1_acts, l1_acts_sf, l1_topk_weights,
l2_acts, l2_acts_sf,
expert_recv_count_sum, l1_arrival_count, token_src_metadata);
expert_recv_count_sum, l1_arrival_count, token_src_metadata,
l1_accum_debug);
};
return {reinterpret_cast<int64_t>(combine_token_buffer.get_end_ptr()), slice_input_buffers};
}
@@ -228,7 +240,8 @@ static void fp8_mega_moe(
const auto [x, x_sf, topk_idx, topk_weights,
l1_acts, l1_acts_sf, l1_topk_weights,
l2_acts, l2_acts_sf,
expert_recv_count_sum, l1_arrival_count, token_src_metadata] = slice(sym_buffer);
expert_recv_count_sum, l1_arrival_count, token_src_metadata,
l1_accum_debug] = slice(sym_buffer);
// Dispatch into SM90 path
DG_HOST_ASSERT(arch_major == 9);
@@ -237,6 +250,7 @@ static void fp8_mega_moe(
l2_acts, l2_acts_sf,
l1_weights, l2_weights,
l1_weights_sf, l2_weights_sf,
l1_accum_debug,
cumulative_local_expert_recv_stats,
sym_buffer_ptrs,
rank_idx, num_max_tokens_per_rank,

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

@@ -29,6 +29,7 @@ public:
// Runtime arguments
void* y;
void* l1_accum_debug;
int* cumulative_local_expert_recv_stats;
int num_tokens;
layout::SymBuffer<> sym_buffer_ptrs;
@@ -101,7 +102,8 @@ static void __instantiate_kernel() {{
args.tensor_map_l2_acts,
args.tensor_map_l2_acts_sf,
args.tensor_map_l2_weights,
args.l2_weights_sf
args.l2_weights_sf,
args.l1_accum_debug
));
}
};
@@ -112,6 +114,7 @@ static void sm90_fp8_mega_moe(
const torch::Tensor& l2_acts, const torch::Tensor& l2_acts_sf,
const torch::Tensor& l1_weights, const torch::Tensor& l2_weights,
const torch::Tensor& l1_weights_sf, const torch::Tensor& l2_weights_sf,
const torch::Tensor& l1_accum_debug,
const std::optional<torch::Tensor> cumulative_local_expert_recv_stats,
const std::vector<int64_t>& sym_buffer_ptrs,
const int& rank_idx, const int& num_max_tokens_per_rank,
@@ -184,6 +187,7 @@ static void sm90_fp8_mega_moe(
.fast_math = fast_math,
.config = config,
.y = y.data_ptr(),
.l1_accum_debug = l1_accum_debug.data_ptr(),
.cumulative_local_expert_recv_stats = cumulative_local_expert_recv_stats_ptr,
.num_tokens = num_tokens,
.sym_buffer_ptrs = layout::SymBuffer<>(sym_buffer_ptrs, rank_idx),

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,

169
megamoe_dev_test_scripts/phase3/l1_wgmma_single_tile.py Normal file
View File

@@ -0,0 +1,169 @@
import argparse
import inspect
import os
import pathlib
import random
import sys
from typing import Tuple
import torch
import torch.distributed as dist
REPO_ROOT = pathlib.Path(__file__).resolve().parents[2]
if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
import deep_gemm
from deep_gemm.utils.math import ceil_div
def init_test_dist(local_rank_arg: int = None) -> Tuple[int, int, dist.ProcessGroup]:
local_rank = local_rank_arg if local_rank_arg is not None else int(os.environ.get('LOCAL_RANK', '0'))
rank = int(os.environ.get('RANK', '0'))
world_size = int(os.environ.get('WORLD_SIZE', '1'))
master_addr = os.environ.get('MASTER_ADDR', '127.0.0.1')
master_port = int(os.environ.get('MASTER_PORT', '8362'))
torch.cuda.set_device(local_rank)
sig = inspect.signature(dist.init_process_group)
params = {
'backend': 'nccl',
'init_method': f'tcp://{master_addr}:{master_port}',
'world_size': world_size,
'rank': rank,
}
if 'device_id' in sig.parameters:
params['device_id'] = torch.device(f'cuda:{local_rank}')
dist.init_process_group(**params)
torch.set_default_device('cuda')
return rank, world_size, dist.new_group(list(range(world_size)))
def interleaved_to_natural_n(n: torch.Tensor, half_n: int, gran: int = 8) -> torch.Tensor:
pair_group = n // (2 * gran)
offset = n - pair_group * (2 * gran)
gate_n = pair_group * gran + offset
up_n = half_n + pair_group * gran + offset - gran
return torch.where(offset < gran, gate_n, up_n)
def make_weights(num_experts_per_rank: int, hidden: int, intermediate_hidden: int):
torch.manual_seed(2027 + dist.get_rank())
l1_weights = (torch.randn(
(num_experts_per_rank, intermediate_hidden * 2, hidden),
dtype=torch.float32, device='cuda') * 0.5).to(torch.float8_e4m3fn)
l2_weights = (torch.randn(
(num_experts_per_rank, hidden, intermediate_hidden),
dtype=torch.float32, device='cuda') * 0.5).to(torch.float8_e4m3fn)
l1_weights_sf = torch.empty(
(num_experts_per_rank, ceil_div(intermediate_hidden * 2, 128), hidden // 128),
dtype=torch.float32, device='cuda')
for expert in range(num_experts_per_rank):
for n_group in range(l1_weights_sf.shape[1]):
for k_group in range(l1_weights_sf.shape[2]):
l1_weights_sf[expert, n_group, k_group] = 0.75 + 0.125 * n_group + 0.0625 * k_group
l2_weights_sf = torch.ones(
(num_experts_per_rank, ceil_div(hidden, 128), intermediate_hidden // 128),
dtype=torch.float32, device='cuda')
return deep_gemm.transform_weights_for_mega_moe(
(l1_weights, l1_weights_sf), (l2_weights, l2_weights_sf))
def reference_l1_accum(buffer: deep_gemm.SymmBuffer,
l1_weights: torch.Tensor,
l1_weights_sf: torch.Tensor,
hidden: int,
intermediate_hidden: int) -> torch.Tensor:
block_m = 128
block_n = 128
x_fp8 = buffer.l1_acts[:block_m]
x_sf = buffer.l1_acts_sf[:block_m, :hidden // 128]
w_fp8 = l1_weights[0, :block_n]
n = torch.arange(block_n, device='cuda')
natural_n = interleaved_to_natural_n(n, intermediate_hidden)
n_groups = natural_n // 128
expected = torch.zeros((block_m, block_n), dtype=torch.float32, device='cuda')
for k_group in range(hidden // 128):
k_start = k_group * 128
k_end = k_start + 128
x = x_fp8[:, k_start:k_end].to(torch.float32)
w = w_fp8[:, k_start:k_end].to(torch.float32)
sfa = x_sf[:, k_group].to(torch.float32)
sfb = l1_weights_sf[0, n_groups, k_group].to(torch.float32)
expected += (x @ w.t()) * sfa[:, None] * sfb[None, :]
return expected
def run_case(args: argparse.Namespace, group: dist.ProcessGroup, rank_idx: int, num_ranks: int) -> None:
hidden = args.hidden
intermediate_hidden = args.intermediate_hidden
num_tokens = args.num_tokens
num_topk = 1
num_experts = args.num_experts if args.num_experts is not None else num_ranks
num_experts_per_rank = num_experts // num_ranks
assert num_experts % num_ranks == 0
assert num_experts_per_rank >= 1
buffer = deep_gemm.get_symm_buffer_for_mega_moe(
group, num_experts, args.num_max_tokens_per_rank, num_topk,
hidden, intermediate_hidden)
weights = make_weights(num_experts_per_rank, hidden, intermediate_hidden)
torch.manual_seed(1234 + rank_idx)
x = (torch.randn((num_tokens, hidden), dtype=torch.float32, device='cuda') * 0.5).to(torch.float8_e4m3fn)
x_sf = torch.rand((num_tokens, hidden // 128), dtype=torch.float32, device='cuda') + 0.75
local_global_expert = rank_idx * num_experts_per_rank
topk_idx = torch.full((num_tokens, num_topk), local_global_expert, dtype=torch.long, device='cuda')
topk_weights = torch.ones((num_tokens, num_topk), dtype=torch.float32, device='cuda')
buffer.x[:num_tokens].copy_(x)
buffer.x_sf[:num_tokens].copy_(x_sf)
buffer.topk_idx[:num_tokens].copy_(topk_idx)
buffer.topk_weights[:num_tokens].copy_(topk_weights)
torch.cuda.synchronize()
dist.barrier(group=group)
y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda')
deep_gemm.fp8_mega_moe(y, weights[0], weights[1], buffer)
torch.cuda.synchronize()
expected = reference_l1_accum(buffer, weights[0][0], weights[0][1], hidden, intermediate_hidden)
actual = buffer.l1_accum_debug
torch.testing.assert_close(actual.cpu(), expected.cpu(), rtol=args.rtol, atol=args.atol)
dist.barrier(group=group)
if rank_idx == 0:
max_diff = (actual - expected).abs().max().item()
print(f'[PASSED] Phase 3 L1 WGMMA single tile, max_diff={max_diff:.6f}', flush=True)
buffer.destroy()
def main() -> None:
parser = argparse.ArgumentParser(description='SM90 MegaMoE Phase 3 L1 WGMMA single-tile correctness')
parser.add_argument('--num-tokens', type=int, default=128)
parser.add_argument('--num-max-tokens-per-rank', type=int, default=384)
parser.add_argument('--hidden', type=int, default=256)
parser.add_argument('--intermediate-hidden', type=int, default=128)
parser.add_argument('--num-experts', type=int, default=None)
parser.add_argument('--local-rank', type=int, default=None)
parser.add_argument('--atol', type=float, default=1e-2)
parser.add_argument('--rtol', type=float, default=5e-2)
args = parser.parse_args()
rank_idx, num_ranks, group = init_test_dist(args.local_rank)
assert torch.cuda.get_device_capability(torch.cuda.current_device())[0] == 9
assert args.num_tokens == 128
assert args.hidden % 128 == 0 and args.intermediate_hidden % 128 == 0
random.seed(5678 + rank_idx)
run_case(args, group, rank_idx, num_ranks)
dist.destroy_process_group()
if __name__ == '__main__':
main()