feat: implement sm90 megamoe phase2 dispatch-only
This commit is contained in:
Xinyi Liu
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user