feat: implement sm90 megamoe phase2 dispatch-only
This commit is contained in:
Xinyi Liu
@@ -13,7 +13,9 @@ namespace deep_gemm::mega {
|
||||
|
||||
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;
|
||||
@@ -92,8 +94,11 @@ get_symm_buffer_size_for_sm90_mega_moe(
|
||||
DG_HOST_ASSERT(hidden % 128 == 0 and intermediate_hidden % 128 == 0);
|
||||
DG_HOST_ASSERT(num_max_padded_sf_pool_tokens % 4 == 0);
|
||||
|
||||
// Slice function: creates input and L1/L2 pool tensor views.
|
||||
// Slice function: creates input, L1/L2 pool, and Phase 2 dispatch-result tensor views.
|
||||
auto slice_input_buffers = [=](const torch::Tensor& buffer) {
|
||||
static_assert(sizeof(layout::TokenSrcMetadata) == 3 * sizeof(uint32_t));
|
||||
const auto runtime_workspace = layout::Workspace(
|
||||
buffer.data_ptr(), num_ranks, num_experts, num_max_tokens_per_rank, num_topk);
|
||||
auto x = torch::from_blob(
|
||||
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(input_token_buffer.base)),
|
||||
{num_max_tokens_per_rank, hidden},
|
||||
@@ -119,6 +124,10 @@ get_symm_buffer_size_for_sm90_mega_moe(
|
||||
{num_max_padded_sf_pool_tokens, hidden / 128},
|
||||
{1, num_max_padded_sf_pool_tokens},
|
||||
torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device()));
|
||||
auto l1_topk_weights = torch::from_blob(
|
||||
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(l1_topk_weights_buffer.base)),
|
||||
{num_max_pool_tokens},
|
||||
torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device()));
|
||||
auto l2_acts = torch::from_blob(
|
||||
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(l2_token_buffer.base)),
|
||||
{num_max_pool_tokens, intermediate_hidden},
|
||||
@@ -128,8 +137,22 @@ get_symm_buffer_size_for_sm90_mega_moe(
|
||||
{num_max_padded_sf_pool_tokens, intermediate_hidden / 128},
|
||||
{1, num_max_padded_sf_pool_tokens},
|
||||
torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device()));
|
||||
auto expert_recv_count_sum = torch::from_blob(
|
||||
runtime_workspace.get_expert_recv_count_sum_ptr(),
|
||||
{num_experts / num_ranks},
|
||||
torch::TensorOptions().dtype(torch::kInt64).device(buffer.device()));
|
||||
auto l1_arrival_count = torch::from_blob(
|
||||
runtime_workspace.get_l1_arrival_count_ptr(),
|
||||
{static_cast<int>(runtime_workspace.num_max_pool_blocks)},
|
||||
torch::TensorOptions().dtype(torch::kInt).device(buffer.device()));
|
||||
auto token_src_metadata = torch::from_blob(
|
||||
reinterpret_cast<int32_t*>(runtime_workspace.get_token_src_metadata_ptr()),
|
||||
{num_max_pool_tokens, 3},
|
||||
torch::TensorOptions().dtype(torch::kInt).device(buffer.device()));
|
||||
return std::make_tuple(x, x_sf, topk_idx, topk_weights,
|
||||
l1_acts, l1_acts_sf, l2_acts, l2_acts_sf);
|
||||
l1_acts, l1_acts_sf, l1_topk_weights,
|
||||
l2_acts, l2_acts_sf,
|
||||
expert_recv_count_sum, l1_arrival_count, token_src_metadata);
|
||||
};
|
||||
return {reinterpret_cast<int64_t>(combine_token_buffer.get_end_ptr()), slice_input_buffers};
|
||||
}
|
||||
@@ -202,7 +225,10 @@ static void fp8_mega_moe(
|
||||
DG_HOST_ASSERT(num_experts == num_experts_);
|
||||
|
||||
// Already registered tensors
|
||||
const auto [x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf] = slice(sym_buffer);
|
||||
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);
|
||||
|
||||
// Dispatch into SM90 path
|
||||
DG_HOST_ASSERT(arch_major == 9);
|
||||
|
||||
@@ -2,11 +2,15 @@
|
||||
|
||||
#include <cutlass/cutlass.h>
|
||||
|
||||
#include <deep_gemm/common/math.cuh>
|
||||
#include <deep_gemm/common/exception.cuh>
|
||||
#include <deep_gemm/common/tma_utils.cuh>
|
||||
#include <deep_gemm/common/types.cuh>
|
||||
#include <deep_gemm/comm/barrier.cuh>
|
||||
#include <deep_gemm/layout/mega_moe.cuh>
|
||||
#include <deep_gemm/layout/sym_buffer.cuh>
|
||||
#include <deep_gemm/ptx/ld_st.cuh>
|
||||
#include <deep_gemm/ptx/utils.cuh>
|
||||
#include <deep_gemm/scheduler/sm90_mega_moe.cuh>
|
||||
|
||||
namespace deep_gemm {
|
||||
@@ -51,7 +55,310 @@ sm90_fp8_mega_moe_impl(void* y,
|
||||
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.
|
||||
#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(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");
|
||||
|
||||
constexpr uint32_t kNumDispatchWarps = kNumDispatchThreads / 32;
|
||||
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;
|
||||
DG_STATIC_ASSERT(kNumTokensPerWarp > 0, "Invalid number of top-k experts");
|
||||
DG_STATIC_ASSERT(kNumPaddedSFPoolTokens % SF_BLOCK_M == 0, "Invalid padded SF pool size");
|
||||
|
||||
// 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;
|
||||
constexpr uint32_t kDispatchGridSyncIndex = 0;
|
||||
constexpr uint32_t kAfterWorkspaceCleanBarrierTag = 1;
|
||||
constexpr uint32_t kBeforeDispatchPullBarrierTag = 2;
|
||||
const auto dispatch_sync = []() {
|
||||
ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);
|
||||
};
|
||||
|
||||
const auto workspace = layout::Workspace(
|
||||
sym_buffer.get_base_ptr(), kNumRanks, kNumExperts, kNumMaxTokensPerRank, kNumTopk);
|
||||
|
||||
constexpr auto fp8_token_layout = layout::Data(kHidden);
|
||||
constexpr auto fp8_sf_layout = layout::Data(kHidden / 128 * static_cast<uint32_t>(sizeof(float)), false);
|
||||
constexpr auto fp8_intermediate_token_layout = layout::Data(kIntermediateHidden);
|
||||
constexpr auto fp8_intermediate_sf_layout = layout::Data(kIntermediateHidden / 128 * static_cast<uint32_t>(sizeof(float)), false);
|
||||
constexpr auto input_topk_idx_layout = layout::Data(kNumTopk * static_cast<uint32_t>(sizeof(int64_t)), false);
|
||||
constexpr auto input_topk_weights_layout = layout::Data(kNumTopk * static_cast<uint32_t>(sizeof(float)), false);
|
||||
constexpr auto l1_topk_weights_layout = layout::Data(static_cast<uint32_t>(sizeof(float)), false);
|
||||
|
||||
const auto input_token_buffer = layout::Buffer(
|
||||
fp8_token_layout, 1, kNumMaxTokensPerRank,
|
||||
workspace.get_end_ptr());
|
||||
const auto input_sf_buffer = layout::Buffer(
|
||||
fp8_sf_layout, 1, kNumMaxTokensPerRank,
|
||||
input_token_buffer.get_end_ptr());
|
||||
const auto input_topk_idx_buffer = layout::Buffer(
|
||||
input_topk_idx_layout, 1, kNumMaxTokensPerRank,
|
||||
input_sf_buffer.get_end_ptr());
|
||||
const auto input_topk_weights_buffer = layout::Buffer(
|
||||
input_topk_weights_layout, 1, kNumMaxTokensPerRank,
|
||||
input_topk_idx_buffer.get_end_ptr());
|
||||
|
||||
const auto l1_token_buffer = layout::Buffer(
|
||||
fp8_token_layout, 1, kNumMaxPoolTokens,
|
||||
input_topk_weights_buffer.get_end_ptr());
|
||||
const auto l1_sf_buffer = layout::Buffer(
|
||||
fp8_sf_layout, 1, kNumPaddedSFPoolTokens,
|
||||
l1_token_buffer.get_end_ptr());
|
||||
const auto l1_topk_weights_buffer = layout::Buffer(
|
||||
l1_topk_weights_layout, 1, kNumMaxPoolTokens,
|
||||
l1_sf_buffer.get_end_ptr());
|
||||
const auto l2_token_buffer = layout::Buffer(
|
||||
fp8_intermediate_token_layout, 1, kNumMaxPoolTokens,
|
||||
l1_topk_weights_buffer.get_end_ptr());
|
||||
const auto l2_sf_buffer = layout::Buffer(
|
||||
fp8_intermediate_sf_layout, 1, kNumPaddedSFPoolTokens,
|
||||
l2_token_buffer.get_end_ptr());
|
||||
(void)l2_sf_buffer;
|
||||
|
||||
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;
|
||||
}
|
||||
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;
|
||||
ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);
|
||||
|
||||
const auto read_topk_idx = [&](const auto& process) {
|
||||
#pragma unroll 1
|
||||
for (uint32_t i = (sm_idx * kNumDispatchWarps + warp_idx) * kNumTokensPerWarp;
|
||||
i < num_tokens;
|
||||
i += kNumSMs * kNumDispatchWarps * kNumTokensPerWarp) {
|
||||
const uint32_t token_offset = lane_idx / kNumTopk;
|
||||
const uint32_t token_idx = i + token_offset;
|
||||
int expert_idx = -1;
|
||||
if (token_idx < num_tokens and lane_idx < kNumActivateLanes) {
|
||||
expert_idx = static_cast<int>(
|
||||
__ldg(input_topk_idx_buffer.get_base_ptr<int64_t>() + i * kNumTopk + lane_idx));
|
||||
if (expert_idx >= 0 and expert_idx < static_cast<int>(kNumExperts))
|
||||
process(i * kNumTopk + lane_idx, static_cast<uint32_t>(expert_idx));
|
||||
}
|
||||
__syncwarp();
|
||||
}
|
||||
};
|
||||
|
||||
// Count local outgoing token-topk entries per global expert.
|
||||
read_topk_idx([&](const uint32_t&, const uint32_t& expert_idx) {
|
||||
atomicAdd_block(smem_expert_count + expert_idx, 1);
|
||||
});
|
||||
ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);
|
||||
|
||||
// Convert per-SM counts into global per-rank offsets. High 32 bits count arrived SMs.
|
||||
for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) {
|
||||
const uint64_t send_value = (1ull << 32) | static_cast<uint64_t>(smem_expert_count[i]);
|
||||
smem_expert_count[i] = static_cast<uint32_t>(
|
||||
ptx::atomic_add(workspace.get_expert_send_count_ptr(i), send_value));
|
||||
}
|
||||
ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);
|
||||
|
||||
// Write source token-topk indices into destination ranks' local-expert tables.
|
||||
read_topk_idx([&](const uint32_t& token_topk_idx, const uint32_t& expert_idx) {
|
||||
const uint32_t dst_rank_idx = expert_idx / kNumExpertsPerRank;
|
||||
const uint32_t dst_local_expert_idx = expert_idx - dst_rank_idx * kNumExpertsPerRank;
|
||||
const uint32_t dst_slot_idx = atomicAdd_block(smem_expert_count + expert_idx, 1);
|
||||
const auto dst_ptr = workspace.get_src_token_topk_idx_ptr(
|
||||
dst_local_expert_idx, sym_buffer.rank_idx, dst_slot_idx);
|
||||
*sym_buffer.map(dst_ptr, dst_rank_idx) = token_topk_idx;
|
||||
});
|
||||
|
||||
// Wait until all local SMs have finished filling local send-count/source-index data.
|
||||
comm::grid_sync<kNumSMs, kDispatchGridSyncIndex>(
|
||||
workspace, sm_idx, thread_idx, dispatch_sync);
|
||||
|
||||
// Publish per-rank expert counts and summed ready/count words to destination ranks.
|
||||
if (sm_idx == 0) {
|
||||
for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) {
|
||||
const uint32_t dst_rank_idx = i / kNumExpertsPerRank;
|
||||
const uint32_t dst_local_expert_idx = i - dst_rank_idx * kNumExpertsPerRank;
|
||||
const uint64_t expert_status = *workspace.get_expert_send_count_ptr(i);
|
||||
*sym_buffer.map(
|
||||
workspace.get_expert_recv_count_ptr(sym_buffer.rank_idx, dst_local_expert_idx),
|
||||
dst_rank_idx) = expert_status & 0xffffffffu;
|
||||
ptx::atomic_add_sys(
|
||||
sym_buffer.map(workspace.get_expert_recv_count_sum_ptr(dst_local_expert_idx), dst_rank_idx),
|
||||
expert_status);
|
||||
}
|
||||
}
|
||||
ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);
|
||||
|
||||
// All ranks must observe the complete count/source-index tables before pulling.
|
||||
comm::nvlink_barrier<kNumRanks, kNumSMs, kNumDispatchThreads,
|
||||
kDispatchGridSyncIndex, kBeforeDispatchPullBarrierTag>(
|
||||
workspace, sym_buffer, sm_idx, thread_idx, dispatch_sync,
|
||||
/* sync_prologue */ false, /* sync_epilogue */ true);
|
||||
|
||||
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);
|
||||
scheduler.fetch_expert_recv_count();
|
||||
|
||||
if (sm_idx == 0 and cumulative_local_expert_recv_stats != nullptr) {
|
||||
for (uint32_t i = thread_idx; i < kNumExpertsPerRank; i += kNumDispatchThreads) {
|
||||
const auto num_recv_tokens = static_cast<uint32_t>(*workspace.get_expert_recv_count_sum_ptr(i));
|
||||
ptx::red_add(cumulative_local_expert_recv_stats + i, static_cast<int>(num_recv_tokens));
|
||||
}
|
||||
}
|
||||
ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);
|
||||
|
||||
constexpr uint32_t kNumRanksPerLane = math::constexpr_ceil_div(kNumRanks, 32u);
|
||||
int current_expert_idx = -1;
|
||||
uint32_t stored_rank_count[kNumRanksPerLane] = {};
|
||||
uint32_t expert_start_idx = 0, expert_end_idx = 0;
|
||||
uint32_t expert_pool_block_offset = 0;
|
||||
|
||||
constexpr uint32_t kNumGlobalDispatchWarps = kNumSMs * kNumDispatchWarps;
|
||||
for (uint32_t token_idx = sm_idx * kNumDispatchWarps + warp_idx; ; token_idx += kNumGlobalDispatchWarps) {
|
||||
int old_expert_idx = current_expert_idx;
|
||||
while (token_idx >= expert_end_idx) {
|
||||
if (++ current_expert_idx >= static_cast<int>(kNumExpertsPerRank))
|
||||
break;
|
||||
|
||||
expert_pool_block_offset += math::ceil_div(expert_end_idx - expert_start_idx, BLOCK_M);
|
||||
expert_start_idx = expert_end_idx;
|
||||
expert_end_idx += scheduler.get_num_tokens(static_cast<uint32_t>(current_expert_idx));
|
||||
}
|
||||
|
||||
if (current_expert_idx >= static_cast<int>(kNumExpertsPerRank))
|
||||
break;
|
||||
|
||||
if (old_expert_idx != current_expert_idx) {
|
||||
old_expert_idx = current_expert_idx;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) {
|
||||
const uint32_t j = i * 32 + lane_idx;
|
||||
stored_rank_count[i] = j < kNumRanks ?
|
||||
static_cast<uint32_t>(*workspace.get_expert_recv_count_ptr(j, static_cast<uint32_t>(current_expert_idx))) : 0;
|
||||
}
|
||||
}
|
||||
|
||||
uint32_t current_rank_in_expert_idx = 0;
|
||||
uint32_t remaining[kNumRanksPerLane];
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumRanksPerLane; ++ i)
|
||||
remaining[i] = stored_rank_count[i];
|
||||
|
||||
uint32_t offset = 0;
|
||||
const uint32_t token_idx_in_expert = token_idx - expert_start_idx;
|
||||
uint32_t slot_idx = token_idx_in_expert;
|
||||
uint32_t token_idx_in_rank = 0;
|
||||
while (true) {
|
||||
uint32_t num_actives_in_lane = 0;
|
||||
uint32_t min_in_lane = 0xffffffffu;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) {
|
||||
num_actives_in_lane += remaining[i] > 0;
|
||||
if (remaining[i] > 0)
|
||||
min_in_lane = cute::min(min_in_lane, remaining[i]);
|
||||
}
|
||||
const uint32_t num_active_ranks = __reduce_add_sync(0xffffffff, num_actives_in_lane);
|
||||
const uint32_t length = __reduce_min_sync(0xffffffff, min_in_lane);
|
||||
const uint32_t num_round_tokens = length * num_active_ranks;
|
||||
|
||||
if (slot_idx < num_round_tokens) {
|
||||
const uint32_t slot_idx_in_round = slot_idx % num_active_ranks;
|
||||
uint32_t num_seen_ranks = 0;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) {
|
||||
const uint32_t mask = __ballot_sync(0xffffffff, remaining[i] > 0);
|
||||
const uint32_t num_active_lanes = __popc(mask);
|
||||
if (slot_idx_in_round >= num_seen_ranks and slot_idx_in_round < num_seen_ranks + num_active_lanes)
|
||||
current_rank_in_expert_idx = i * 32 + __fns(mask, 0, slot_idx_in_round - num_seen_ranks + 1);
|
||||
num_seen_ranks += num_active_lanes;
|
||||
}
|
||||
token_idx_in_rank = offset + slot_idx / num_active_ranks;
|
||||
break;
|
||||
}
|
||||
|
||||
slot_idx -= num_round_tokens;
|
||||
offset += length;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumRanksPerLane; ++ i)
|
||||
remaining[i] -= cute::min(remaining[i], length);
|
||||
}
|
||||
|
||||
const uint32_t src_token_topk_idx = *workspace.get_src_token_topk_idx_ptr(
|
||||
static_cast<uint32_t>(current_expert_idx), current_rank_in_expert_idx, token_idx_in_rank);
|
||||
const uint32_t src_token_idx = src_token_topk_idx / kNumTopk;
|
||||
const uint32_t src_topk_idx = src_token_topk_idx - src_token_idx * kNumTopk;
|
||||
const uint32_t pool_token_idx = expert_pool_block_offset * BLOCK_M + token_idx_in_expert;
|
||||
const uint32_t sf_pool_token_idx =
|
||||
(expert_pool_block_offset + token_idx_in_expert / BLOCK_M) * SF_BLOCK_M +
|
||||
(token_idx_in_expert % BLOCK_M);
|
||||
|
||||
const auto remote_token_ptr = sym_buffer.map(
|
||||
input_token_buffer.get_data_buffer(src_token_idx).get_base_ptr<uint4>(),
|
||||
current_rank_in_expert_idx);
|
||||
const auto local_token_ptr = l1_token_buffer.get_data_buffer(pool_token_idx).get_base_ptr<uint4>();
|
||||
#pragma unroll
|
||||
for (uint32_t i = lane_idx; i < kNumTokenUint4; i += 32)
|
||||
local_token_ptr[i] = remote_token_ptr[i];
|
||||
|
||||
const auto remote_sf_ptr = sym_buffer.map(
|
||||
input_sf_buffer.get_data_buffer(src_token_idx).get_base_ptr<float>(),
|
||||
current_rank_in_expert_idx);
|
||||
const auto local_sf_ptr = l1_sf_buffer.get_base_ptr<float>();
|
||||
#pragma unroll
|
||||
for (uint32_t i = lane_idx; i < kNumSFValues; i += 32)
|
||||
local_sf_ptr[i * kNumPaddedSFPoolTokens + sf_pool_token_idx] = remote_sf_ptr[i];
|
||||
|
||||
__syncwarp();
|
||||
if (lane_idx == 0) {
|
||||
const float weight = *sym_buffer.map(
|
||||
input_topk_weights_buffer.get_base_ptr<float>() + src_token_topk_idx,
|
||||
current_rank_in_expert_idx);
|
||||
*l1_topk_weights_buffer.get_data_buffer(pool_token_idx).get_base_ptr<float>() = weight;
|
||||
*workspace.get_token_src_metadata_ptr(pool_token_idx) = {
|
||||
current_rank_in_expert_idx, src_token_idx, src_topk_idx};
|
||||
}
|
||||
__syncwarp();
|
||||
__threadfence();
|
||||
__syncwarp();
|
||||
if (lane_idx == 0) {
|
||||
ptx::red_add_rel(
|
||||
workspace.get_l1_arrival_count_ptr(expert_pool_block_offset + token_idx_in_expert / BLOCK_M), 1);
|
||||
}
|
||||
__syncwarp();
|
||||
}
|
||||
#endif
|
||||
return;
|
||||
}
|
||||
|
||||
|
||||
@@ -60,14 +60,24 @@ class SymmBuffer:
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# Create input buffer views
|
||||
(self.x, self.x_sf,
|
||||
self.topk_idx, self.topk_weights,
|
||||
self.l1_acts, self.l1_acts_sf,
|
||||
self.l2_acts, self.l2_acts_sf) = slice_input_buffers(self.buffer)
|
||||
self.l1_topk_weights = None
|
||||
self.expert_recv_count_sum = None
|
||||
self.l1_arrival_count = None
|
||||
self.token_src_metadata = None
|
||||
buffer_views = slice_input_buffers(self.buffer)
|
||||
if _is_sm90():
|
||||
(self.x, self.x_sf,
|
||||
self.topk_idx, self.topk_weights,
|
||||
self.l1_acts, self.l1_acts_sf, self.l1_topk_weights,
|
||||
self.l2_acts, self.l2_acts_sf,
|
||||
self.expert_recv_count_sum,
|
||||
self.l1_arrival_count,
|
||||
self.token_src_metadata) = buffer_views
|
||||
else:
|
||||
(self.x, self.x_sf,
|
||||
self.topk_idx, self.topk_weights,
|
||||
self.l1_acts, self.l1_acts_sf,
|
||||
self.l2_acts, self.l2_acts_sf) = buffer_views
|
||||
self.l1_topk_weights = None
|
||||
self.expert_recv_count_sum = None
|
||||
self.l1_arrival_count = None
|
||||
self.token_src_metadata = None
|
||||
|
||||
def destroy(self):
|
||||
self.handle = None
|
||||
|
||||
266
megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py
Normal file
266
megamoe_dev_test_scripts/phase2/dispatch_only_correctness.py
Normal file
@@ -0,0 +1,266 @@
|
||||
import argparse
|
||||
import inspect
|
||||
import os
|
||||
import pathlib
|
||||
import random
|
||||
import sys
|
||||
from typing import List, 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 parse_tokens_list(value: str) -> List[int]:
|
||||
return [int(item) for item in value.split(',') if item]
|
||||
|
||||
|
||||
def gather_same_shape(tensor: torch.Tensor, group: dist.ProcessGroup) -> torch.Tensor:
|
||||
gathered = [torch.empty_like(tensor) for _ in range(dist.get_world_size(group))]
|
||||
dist.all_gather(gathered, tensor.contiguous(), group=group)
|
||||
return torch.stack(gathered, dim=0)
|
||||
|
||||
|
||||
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', '8361'))
|
||||
|
||||
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 get_block_m(num_tokens: int, num_ranks: int, num_topk: int, num_experts: int) -> int:
|
||||
expected = num_tokens * num_ranks * num_topk / num_experts
|
||||
if expected <= 16.5:
|
||||
return 32
|
||||
if expected <= 64.5:
|
||||
return 64
|
||||
return 128
|
||||
|
||||
|
||||
def make_topk(num_tokens: int, num_experts: int, num_topk: int, rank_idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
token_idx = torch.arange(num_tokens, device='cuda', dtype=torch.long).unsqueeze(1)
|
||||
topk_slot = torch.arange(num_topk, device='cuda', dtype=torch.long).unsqueeze(0)
|
||||
topk_idx = (token_idx * num_topk + topk_slot + rank_idx * 3) % num_experts
|
||||
topk_weights = (token_idx.float() * 0.125 + topk_slot.float() * 0.25 + rank_idx + 1.0).contiguous()
|
||||
return topk_idx.contiguous(), topk_weights.contiguous()
|
||||
|
||||
|
||||
def make_weights(num_experts_per_rank: int, hidden: int, intermediate_hidden: int):
|
||||
l1_weights = torch.randn(
|
||||
(num_experts_per_rank, intermediate_hidden * 2, hidden),
|
||||
dtype=torch.float32, device='cuda').to(torch.float8_e4m3fn)
|
||||
l2_weights = torch.randn(
|
||||
(num_experts_per_rank, hidden, intermediate_hidden),
|
||||
dtype=torch.float32, device='cuda').to(torch.float8_e4m3fn)
|
||||
l1_weights_sf = torch.ones(
|
||||
(num_experts_per_rank, ceil_div(intermediate_hidden * 2, 128), hidden // 128),
|
||||
dtype=torch.float32, device='cuda')
|
||||
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 build_expected_entries(all_topk_idx: torch.Tensor,
|
||||
rank_idx: int,
|
||||
num_experts_per_rank: int) -> Tuple[List[List[Tuple[int, int, int]]], torch.Tensor]:
|
||||
num_ranks, num_tokens, num_topk = all_topk_idx.shape
|
||||
local_start = rank_idx * num_experts_per_rank
|
||||
local_end = local_start + num_experts_per_rank
|
||||
entries: List[List[Tuple[int, int, int]]] = [[] for _ in range(num_experts_per_rank)]
|
||||
counts = [0 for _ in range(num_experts_per_rank)]
|
||||
|
||||
all_topk_idx_cpu = all_topk_idx.cpu()
|
||||
for src_rank in range(num_ranks):
|
||||
for token_idx in range(num_tokens):
|
||||
for topk_idx in range(num_topk):
|
||||
expert_idx = int(all_topk_idx_cpu[src_rank, token_idx, topk_idx])
|
||||
if local_start <= expert_idx < local_end:
|
||||
local_expert = expert_idx - local_start
|
||||
entries[local_expert].append((src_rank, token_idx, topk_idx))
|
||||
counts[local_expert] += 1
|
||||
return entries, torch.tensor(counts, dtype=torch.int64, device='cuda')
|
||||
|
||||
|
||||
def verify_case(buffer: deep_gemm.SymmBuffer,
|
||||
all_x_f32: torch.Tensor,
|
||||
all_x_sf: torch.Tensor,
|
||||
all_topk_idx: torch.Tensor,
|
||||
all_topk_weights: torch.Tensor,
|
||||
cumulative_stats: torch.Tensor,
|
||||
rank_idx: int,
|
||||
num_ranks: int,
|
||||
num_experts: int,
|
||||
num_topk: int,
|
||||
block_m: int) -> None:
|
||||
del num_topk
|
||||
num_experts_per_rank = num_experts // num_ranks
|
||||
sf_block_m = ((block_m + 127) // 128) * 128
|
||||
num_sms = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count
|
||||
entries_by_expert, counts = build_expected_entries(all_topk_idx, rank_idx, num_experts_per_rank)
|
||||
|
||||
torch.testing.assert_close(cumulative_stats.cpu(), counts.to(torch.int32).cpu(), rtol=0, atol=0)
|
||||
|
||||
expert_status = buffer.expert_recv_count_sum[:num_experts_per_rank].detach().cpu().tolist()
|
||||
arrival = buffer.l1_arrival_count.detach().cpu()
|
||||
metadata = buffer.token_src_metadata.detach().cpu()
|
||||
|
||||
pool_block_offset = 0
|
||||
local_start = rank_idx * num_experts_per_rank
|
||||
for local_expert, expected_entries in enumerate(entries_by_expert):
|
||||
count = len(expected_entries)
|
||||
status = int(expert_status[local_expert])
|
||||
assert (status & 0xffffffff) == count, (local_expert, status, count)
|
||||
assert (status >> 32) == num_sms * num_ranks, (local_expert, status, num_sms, num_ranks)
|
||||
|
||||
expected_set = set(expected_entries)
|
||||
actual_set = set()
|
||||
base_pool_token = pool_block_offset * block_m
|
||||
expected_global_expert = local_start + local_expert
|
||||
|
||||
for token_in_expert in range(count):
|
||||
pool_token_idx = base_pool_token + token_in_expert
|
||||
src_rank, src_token_idx, src_topk_idx = [int(v) for v in metadata[pool_token_idx].tolist()]
|
||||
actual_entry = (src_rank, src_token_idx, src_topk_idx)
|
||||
assert actual_entry in expected_set, (local_expert, token_in_expert, actual_entry, expected_set)
|
||||
assert actual_entry not in actual_set, (local_expert, token_in_expert, actual_entry)
|
||||
actual_set.add(actual_entry)
|
||||
|
||||
expert_idx = int(all_topk_idx[src_rank, src_token_idx, src_topk_idx].item())
|
||||
assert expert_idx == expected_global_expert, (pool_token_idx, expert_idx, expected_global_expert)
|
||||
|
||||
torch.testing.assert_close(
|
||||
buffer.l1_acts[pool_token_idx].to(torch.float32).cpu(),
|
||||
all_x_f32[src_rank, src_token_idx].cpu(),
|
||||
rtol=0, atol=0)
|
||||
torch.testing.assert_close(
|
||||
buffer.l1_topk_weights[pool_token_idx].cpu(),
|
||||
all_topk_weights[src_rank, src_token_idx, src_topk_idx].cpu(),
|
||||
rtol=0, atol=0)
|
||||
|
||||
sf_pool_token_idx = (pool_block_offset + token_in_expert // block_m) * sf_block_m + token_in_expert % block_m
|
||||
torch.testing.assert_close(
|
||||
buffer.l1_acts_sf[sf_pool_token_idx].cpu(),
|
||||
all_x_sf[src_rank, src_token_idx].cpu(),
|
||||
rtol=0, atol=0)
|
||||
|
||||
assert actual_set == expected_set, (local_expert, actual_set, expected_set)
|
||||
|
||||
for block_idx in range(ceil_div(count, block_m)):
|
||||
expected_arrivals = min(block_m, count - block_idx * block_m)
|
||||
actual_arrivals = int(arrival[pool_block_offset + block_idx].item())
|
||||
assert actual_arrivals == expected_arrivals, (local_expert, block_idx, actual_arrivals, expected_arrivals)
|
||||
|
||||
pool_block_offset += ceil_div(count, block_m)
|
||||
|
||||
|
||||
def run_case(args: argparse.Namespace,
|
||||
group: dist.ProcessGroup,
|
||||
rank_idx: int,
|
||||
num_ranks: int,
|
||||
buffer: deep_gemm.SymmBuffer,
|
||||
weights,
|
||||
num_tokens: int) -> None:
|
||||
hidden = args.hidden
|
||||
num_topk = args.num_topk
|
||||
num_experts = args.num_experts
|
||||
|
||||
x_f32 = torch.randn((num_tokens, hidden), dtype=torch.float32, device='cuda')
|
||||
x_fp8 = x_f32.to(torch.float8_e4m3fn)
|
||||
x_sf = torch.rand((num_tokens, hidden // 128), dtype=torch.float32, device='cuda') + 0.5
|
||||
topk_idx, topk_weights = make_topk(num_tokens, num_experts, num_topk, rank_idx)
|
||||
|
||||
all_x_f32 = gather_same_shape(x_fp8.to(torch.float32), group)
|
||||
all_x_sf = gather_same_shape(x_sf, group)
|
||||
all_topk_idx = gather_same_shape(topk_idx, group)
|
||||
all_topk_weights = gather_same_shape(topk_weights, group)
|
||||
|
||||
buffer.x[:num_tokens].copy_(x_fp8)
|
||||
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)
|
||||
|
||||
cumulative_stats = torch.zeros((num_experts // num_ranks,), dtype=torch.int32, device='cuda')
|
||||
y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda')
|
||||
deep_gemm.fp8_mega_moe(y, weights[0], weights[1], buffer,
|
||||
cumulative_local_expert_recv_stats=cumulative_stats)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
block_m = get_block_m(num_tokens, num_ranks, num_topk, num_experts)
|
||||
verify_case(buffer, all_x_f32, all_x_sf, all_topk_idx, all_topk_weights,
|
||||
cumulative_stats, rank_idx, num_ranks, num_experts, num_topk, block_m)
|
||||
dist.barrier(group=group)
|
||||
if rank_idx == 0:
|
||||
print(f'[PASSED] tokens={num_tokens}, block_m={block_m}', flush=True)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description='SM90 MegaMoE Phase 2 dispatch-only correctness')
|
||||
parser.add_argument('--tokens-list', type=str, default='0,8,48,192')
|
||||
parser.add_argument('--num-max-tokens-per-rank', type=int, default=384)
|
||||
parser.add_argument('--hidden', type=int, default=512)
|
||||
parser.add_argument('--intermediate-hidden', type=int, default=256)
|
||||
parser.add_argument('--num-experts', type=int, default=16)
|
||||
parser.add_argument('--num-topk', type=int, default=6)
|
||||
parser.add_argument('--local-rank', type=int, default=None)
|
||||
args = parser.parse_args()
|
||||
|
||||
local_rank = args.local_rank if args.local_rank is not None else int(os.environ.get('LOCAL_RANK', '0'))
|
||||
rank_idx, num_ranks, group = init_test_dist(local_rank)
|
||||
assert torch.cuda.get_device_capability(torch.cuda.current_device())[0] == 9
|
||||
assert args.num_experts % num_ranks == 0
|
||||
assert args.hidden % 128 == 0 and args.intermediate_hidden % 128 == 0
|
||||
assert args.num_topk <= args.num_experts
|
||||
|
||||
torch.manual_seed(1234 + rank_idx)
|
||||
random.seed(1234 + rank_idx)
|
||||
|
||||
tokens_list = parse_tokens_list(args.tokens_list)
|
||||
max_tokens = max([args.num_max_tokens_per_rank] + tokens_list)
|
||||
buffer = deep_gemm.get_symm_buffer_for_mega_moe(
|
||||
group, args.num_experts, max_tokens, args.num_topk,
|
||||
args.hidden, args.intermediate_hidden)
|
||||
weights = make_weights(args.num_experts // num_ranks, args.hidden, args.intermediate_hidden)
|
||||
|
||||
if rank_idx == 0:
|
||||
print(f'[Phase 2] ranks={num_ranks}, tokens_list={tokens_list}, '
|
||||
f'hidden={args.hidden}, intermediate={args.intermediate_hidden}, '
|
||||
f'experts={args.num_experts}, topk={args.num_topk}', flush=True)
|
||||
|
||||
for num_tokens in tokens_list:
|
||||
run_case(args, group, rank_idx, num_ranks, buffer, weights, num_tokens)
|
||||
|
||||
dist.barrier(group=group)
|
||||
buffer.destroy()
|
||||
dist.destroy_process_group()
|
||||
if rank_idx == 0:
|
||||
print('[PASSED] SM90 MegaMoE Phase 2 dispatch-only correctness', flush=True)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in New Issue
Block a user