213 lines
8.5 KiB
C++
213 lines
8.5 KiB
C++
#pragma once
|
|
|
|
#include <algorithm>
|
|
#include <unordered_set>
|
|
|
|
#include <deep_gemm/layout/mega_moe.cuh>
|
|
|
|
#include "../../utils/exception.hpp"
|
|
#include "../../utils/math.hpp"
|
|
#include "../../utils/system.hpp"
|
|
|
|
namespace deep_gemm {
|
|
|
|
struct SM90MegaMoEConfig {
|
|
// Block tiling
|
|
int block_m, block_n, block_k;
|
|
int store_block_m;
|
|
|
|
// Pool capacity and SF-padded token count
|
|
int num_max_pool_tokens;
|
|
int num_padded_sf_pool_tokens;
|
|
|
|
// Number of experts to process per wave
|
|
int num_experts_per_wave;
|
|
|
|
// Pipeline stages and shared memory
|
|
int num_stages, smem_size;
|
|
|
|
// Thread layout (384 total: 64 dispatch + 64 TMA + 256 math)
|
|
int num_dispatch_threads, num_tma_threads, num_math_threads;
|
|
|
|
// Mode flags
|
|
bool cooperative;
|
|
bool use_n_major_l2;
|
|
|
|
friend std::ostream& operator << (std::ostream& os, const SM90MegaMoEConfig& config) {
|
|
os << "SM90MegaMoEConfig("
|
|
<< "block_m=" << config.block_m << ", block_n=" << config.block_n << ", block_k=" << config.block_k
|
|
<< ", store_block_m=" << config.store_block_m
|
|
<< ", num_max_pool_tokens=" << config.num_max_pool_tokens
|
|
<< ", num_padded_sf_pool_tokens=" << config.num_padded_sf_pool_tokens
|
|
<< ", num_experts_per_wave=" << config.num_experts_per_wave
|
|
<< ", num_stages=" << config.num_stages << ", smem_size=" << config.smem_size
|
|
<< ", num_dispatch_threads=" << config.num_dispatch_threads
|
|
<< ", num_tma_threads=" << config.num_tma_threads
|
|
<< ", num_math_threads=" << config.num_math_threads
|
|
<< ", cooperative=" << config.cooperative
|
|
<< ", use_n_major_l2=" << config.use_n_major_l2 << ")";
|
|
return os;
|
|
}
|
|
};
|
|
|
|
static std::tuple<int, int, int, bool> get_block_config_for_sm90_mega_moe(
|
|
const int& num_ranks, const int& num_experts,
|
|
const int& num_max_tokens_per_rank, const int& num_topk,
|
|
const int& num_tokens) {
|
|
|
|
float num_expected_tokens_per_expert =
|
|
static_cast<float>(num_tokens) * num_ranks * num_topk / num_experts;
|
|
|
|
if (num_expected_tokens_per_expert <= 16.5) {
|
|
// Extreme decode: RL long-tail, large EP
|
|
return {32, 16, 256, false};
|
|
} else if (num_expected_tokens_per_expert <= 64.5) {
|
|
// Medium decode
|
|
return {64, 32, 256, false};
|
|
} else {
|
|
// Large M prefill / large EP decode
|
|
return {128, 32, 256, true};
|
|
}
|
|
}
|
|
|
|
static int get_num_experts_per_wave_for_sm90_mega_moe(
|
|
const int& num_experts_per_rank, const int& num_tokens, const int& num_topk,
|
|
const int& intermediate_hidden, const int& block_m, const int& block_n, const int& num_sms) {
|
|
|
|
float expected_tokens_per_expert = static_cast<float>(num_tokens) * num_topk / num_experts_per_rank;
|
|
if (expected_tokens_per_expert < 1) {
|
|
return num_experts_per_rank;
|
|
}
|
|
|
|
constexpr int kImbalanceFactor = 2;
|
|
|
|
const int num_m_blocks = ceil_div(static_cast<int>(std::ceil(expected_tokens_per_expert)), block_m);
|
|
const int num_n_blocks = (2 * intermediate_hidden) / block_n;
|
|
const int num_l1_blocks_per_expert = num_m_blocks * num_n_blocks;
|
|
|
|
int num_experts_per_wave = num_l1_blocks_per_expert > 0
|
|
? ceil_div(kImbalanceFactor * num_sms, num_l1_blocks_per_expert) : 1;
|
|
num_experts_per_wave = std::min(num_experts_per_wave, num_experts_per_rank);
|
|
|
|
while (num_experts_per_wave < num_experts_per_rank and num_experts_per_rank % num_experts_per_wave != 0)
|
|
++ num_experts_per_wave;
|
|
|
|
return num_experts_per_wave;
|
|
}
|
|
|
|
static std::pair<int, int> get_pipeline_config_for_sm90_mega_moe(
|
|
const int& smem_capacity,
|
|
const int& num_experts, const int& hidden,
|
|
const int& block_m, const int& block_n, const int& block_k, const int& store_block_m,
|
|
const int& num_dispatch_threads, const int& num_math_threads,
|
|
const bool& cooperative) {
|
|
constexpr int kSmemAlignment = 1024;
|
|
constexpr int kNumTMAStoreStages = 2;
|
|
|
|
const int num_dispatch_warps = num_dispatch_threads / 32;
|
|
const int num_math_warps = num_math_threads / 32;
|
|
|
|
// Dispatch region
|
|
const int smem_expert_count_size = align(
|
|
num_experts * static_cast<int>(sizeof(uint32_t)), kSmemAlignment);
|
|
const int smem_send_buffers_size = align(
|
|
static_cast<int>(layout::Buffer(layout::Data(hidden), num_dispatch_warps, 1).get_num_bytes()),
|
|
kSmemAlignment);
|
|
const int smem_dispatch_size = smem_expert_count_size + smem_send_buffers_size;
|
|
|
|
// C/D output region: max of L1 FP8 (2 TMA stages) and L2 BF16 staging
|
|
// L1: store_block_m * (block_n / 2) * kNumTMAStoreStages (FP8 = 1 byte)
|
|
// L2: block_m * block_n * sizeof(BF16) (BF16 = 2 bytes)
|
|
const int num_math_warpgroups = cooperative ? 2 : 1;
|
|
const int smem_cd_l1 = num_math_warpgroups * store_block_m * (block_n / 2) * kNumTMAStoreStages;
|
|
const int smem_cd_l2 = block_m * block_n * static_cast<int>(sizeof(nv_bfloat16));
|
|
const int smem_cd = std::max(smem_cd_l1, smem_cd_l2);
|
|
|
|
// Barriers: dispatch + full/empty pipeline (2 per stage) + combine (2 per math warp)
|
|
const int smem_barriers = (num_dispatch_warps + 2 * 2 + num_math_warps * 2) * 8;
|
|
|
|
// Amax reduction
|
|
const int smem_amax_reduction = store_block_m * num_math_warps * static_cast<int>(sizeof(float));
|
|
|
|
// Float SF per stage: align(2 * BLOCK_M * sizeof(float), 128)
|
|
const int smem_sfa_per_stage = align(2 * block_m * static_cast<int>(sizeof(float)), 128);
|
|
|
|
// Per-stage: A tile + B tile + SFA tile + full/empty barriers
|
|
const int smem_per_stage = block_m * block_k + block_n * block_k + smem_sfa_per_stage + 2 * 8;
|
|
|
|
// Fixed total
|
|
const int smem_fixed = smem_dispatch_size + smem_cd + smem_amax_reduction + smem_barriers;
|
|
|
|
const int num_stages = (smem_capacity - smem_fixed) / smem_per_stage;
|
|
DG_HOST_ASSERT(num_stages >= 3);
|
|
|
|
return {num_stages, smem_fixed + num_stages * smem_per_stage};
|
|
}
|
|
|
|
static SM90MegaMoEConfig get_sm90_mega_moe_config(
|
|
const int& num_ranks, const int& num_experts, const int& num_experts_per_rank,
|
|
const int& num_max_tokens_per_rank, const int& num_tokens, const int& num_topk,
|
|
const int& hidden, const int& intermediate_hidden,
|
|
const int& num_padded_sf_pool_tokens) {
|
|
|
|
const auto [block_m, store_block_m, num_math_threads, cooperative] =
|
|
get_block_config_for_sm90_mega_moe(num_ranks, num_experts, num_max_tokens_per_rank, num_topk, num_tokens);
|
|
const int block_n = 128;
|
|
const int block_k = 128;
|
|
|
|
const int num_max_pool_tokens = layout::get_num_max_pool_tokens(
|
|
num_ranks, num_max_tokens_per_rank, num_topk, num_experts_per_rank);
|
|
|
|
// Thread layout: 64 dispatch + 64 TMA + 256 math/epilogue = 384
|
|
const int num_dispatch_threads = 64;
|
|
const int num_tma_threads = 64;
|
|
|
|
// Auto N-major L2: enabled when large M (high tokens per expert)
|
|
const bool use_n_major_l2 = [&]() {
|
|
auto env_val = get_env<int>("DG_SM90_MOE_NMAJOR", -1);
|
|
if (env_val != -1)
|
|
return env_val > 0;
|
|
float expected = static_cast<float>(num_tokens) * num_ranks * num_topk / num_experts;
|
|
return expected >= 256;
|
|
}();
|
|
|
|
// Waves
|
|
const int num_sms = device_runtime->get_num_sms();
|
|
const int num_experts_per_wave = get_num_experts_per_wave_for_sm90_mega_moe(
|
|
num_experts_per_rank, num_tokens, num_topk,
|
|
intermediate_hidden, block_m, block_n, num_sms);
|
|
|
|
// Pipeline
|
|
constexpr int smem_capacity = 232448;
|
|
const auto [num_stages, smem_size] = get_pipeline_config_for_sm90_mega_moe(
|
|
smem_capacity,
|
|
num_experts, hidden,
|
|
block_m, block_n, block_k, store_block_m,
|
|
num_dispatch_threads, num_math_threads,
|
|
cooperative);
|
|
|
|
const auto config = SM90MegaMoEConfig {
|
|
block_m, block_n, block_k,
|
|
store_block_m,
|
|
num_max_pool_tokens, num_padded_sf_pool_tokens,
|
|
num_experts_per_wave,
|
|
num_stages, smem_size,
|
|
num_dispatch_threads, num_tma_threads, num_math_threads,
|
|
cooperative, use_n_major_l2
|
|
};
|
|
|
|
if (get_env<int>("DG_JIT_DEBUG") or get_env<int>("DG_PRINT_CONFIGS")) {
|
|
const auto key = fmt::format(
|
|
"SM90MegaMoEConfig(num_ranks={}, num_experts={}, hidden={}, intermediate_hidden={}, num_max_tokens_per_rank={}, num_tokens={}, num_topk={})",
|
|
num_ranks, num_experts, hidden, intermediate_hidden, num_max_tokens_per_rank, num_tokens, num_topk);
|
|
static std::unordered_set<std::string> printed;
|
|
if (printed.count(key) == 0) {
|
|
std::cout << key << ": " << config << std::endl;
|
|
printed.insert(key);
|
|
}
|
|
}
|
|
return config;
|
|
}
|
|
|
|
} // namespace deep_gemm
|