feat: add sm90 megamoe phase1 interfaces
This commit is contained in:
Xinyi Liu
@@ -0,0 +1,212 @@
|
||||
#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
|
||||
@@ -0,0 +1,207 @@
|
||||
#pragma once
|
||||
|
||||
#include <torch/python.h>
|
||||
|
||||
#include "../../jit/compiler.hpp"
|
||||
#include "../../jit/kernel_runtime.hpp"
|
||||
#include "../../utils/exception.hpp"
|
||||
#include "../../utils/format.hpp"
|
||||
#include "runtime_utils.hpp"
|
||||
|
||||
#include <deep_gemm/layout/mega_moe.cuh>
|
||||
#include <deep_gemm/layout/sym_buffer.cuh>
|
||||
|
||||
#include "../heuristics/sm90_mega_moe.hpp"
|
||||
|
||||
namespace deep_gemm {
|
||||
|
||||
class SM90FP8MegaMoERuntime final : public LaunchRuntime<SM90FP8MegaMoERuntime> {
|
||||
public:
|
||||
struct Args {
|
||||
// Templated arguments
|
||||
int num_max_tokens_per_rank;
|
||||
int hidden, intermediate_hidden;
|
||||
int num_experts, num_topk;
|
||||
int num_ranks;
|
||||
float activation_clamp;
|
||||
bool fast_math;
|
||||
SM90MegaMoEConfig config;
|
||||
|
||||
// Runtime arguments
|
||||
void* y;
|
||||
int* cumulative_local_expert_recv_stats;
|
||||
int num_tokens;
|
||||
layout::SymBuffer<> sym_buffer_ptrs;
|
||||
|
||||
// Tensormap
|
||||
CUtensorMap tensor_map_l1_acts;
|
||||
CUtensorMap tensor_map_l1_acts_sf;
|
||||
CUtensorMap tensor_map_l1_weights;
|
||||
CUtensorMap tensor_map_l1_output;
|
||||
CUtensorMap tensor_map_l2_acts;
|
||||
CUtensorMap tensor_map_l2_acts_sf;
|
||||
CUtensorMap tensor_map_l2_weights;
|
||||
void* l1_weights_sf;
|
||||
void* l2_weights_sf;
|
||||
|
||||
// Launch configs
|
||||
LaunchArgs launch_args;
|
||||
};
|
||||
|
||||
static std::string generate_impl(const Args& args) {
|
||||
return fmt::format(R"(
|
||||
#include <deep_gemm/impls/sm90_fp8_mega_moe.cuh>
|
||||
|
||||
using namespace deep_gemm;
|
||||
|
||||
static void __instantiate_kernel() {{
|
||||
auto ptr = reinterpret_cast<void*>(&sm90_fp8_mega_moe_impl<
|
||||
{},
|
||||
{}, {},
|
||||
{}, {},
|
||||
{},
|
||||
{}, {}, {},
|
||||
{},
|
||||
{}, {},
|
||||
{},
|
||||
{}, {}, {},
|
||||
{}, {},
|
||||
{}, {},
|
||||
{}, {}
|
||||
>);
|
||||
}};
|
||||
)", args.num_max_tokens_per_rank,
|
||||
args.hidden, args.intermediate_hidden,
|
||||
args.num_experts, args.num_topk,
|
||||
args.config.num_experts_per_wave,
|
||||
args.config.block_m, args.config.block_n, args.config.block_k,
|
||||
args.config.store_block_m,
|
||||
args.config.num_max_pool_tokens,
|
||||
args.config.num_padded_sf_pool_tokens,
|
||||
args.config.num_stages,
|
||||
args.config.num_dispatch_threads, args.config.num_tma_threads, args.config.num_math_threads,
|
||||
args.config.cooperative ? "true" : "false",
|
||||
args.config.use_n_major_l2 ? "true" : "false",
|
||||
args.launch_args.grid_dim.first, args.num_ranks,
|
||||
to_string(args.activation_clamp),
|
||||
args.fast_math ? "true" : "false");
|
||||
}
|
||||
|
||||
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
|
||||
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
|
||||
args.y,
|
||||
args.cumulative_local_expert_recv_stats,
|
||||
args.num_tokens,
|
||||
args.sym_buffer_ptrs,
|
||||
args.tensor_map_l1_acts,
|
||||
args.tensor_map_l1_acts_sf,
|
||||
args.tensor_map_l1_weights,
|
||||
args.l1_weights_sf,
|
||||
args.tensor_map_l1_output,
|
||||
args.tensor_map_l2_acts,
|
||||
args.tensor_map_l2_acts_sf,
|
||||
args.tensor_map_l2_weights,
|
||||
args.l2_weights_sf
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
static void sm90_fp8_mega_moe(
|
||||
const torch::Tensor& y,
|
||||
const torch::Tensor& l1_acts, const torch::Tensor& l1_acts_sf,
|
||||
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 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,
|
||||
const int& num_experts_per_rank,
|
||||
const int& num_tokens, const int& num_topk,
|
||||
const int& hidden, const int& intermediate_hidden,
|
||||
const float& activation_clamp,
|
||||
const bool& fast_math
|
||||
) {
|
||||
const auto num_ranks = static_cast<int>(sym_buffer_ptrs.size());
|
||||
const auto num_experts = num_experts_per_rank * num_ranks;
|
||||
const auto num_padded_sf_pool_tokens = static_cast<int>(l1_acts_sf.size(0));
|
||||
|
||||
// Heuristics
|
||||
const auto config = get_sm90_mega_moe_config(
|
||||
num_ranks, num_experts, num_experts_per_rank,
|
||||
num_max_tokens_per_rank, num_tokens, num_topk, hidden, intermediate_hidden, num_padded_sf_pool_tokens);
|
||||
|
||||
// Make tensormap
|
||||
constexpr int kGranK = 128;
|
||||
const auto tensor_map_l1_acts = make_tma_2d_desc(l1_acts,
|
||||
hidden, config.num_max_pool_tokens,
|
||||
config.block_k, config.block_m,
|
||||
static_cast<int>(l1_acts.stride(-2)),
|
||||
128);
|
||||
const auto tensor_map_l1_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l1_acts_sf,
|
||||
config.num_padded_sf_pool_tokens, hidden,
|
||||
config.block_m, kGranK,
|
||||
1, 0);
|
||||
const auto tensor_map_l1_weights = make_tma_2d_desc(l1_weights,
|
||||
hidden, num_experts_per_rank * intermediate_hidden * 2,
|
||||
config.block_k, config.block_n,
|
||||
static_cast<int>(l1_weights.stride(-2)),
|
||||
128);
|
||||
// L1 output SwiGLU has half N width
|
||||
const auto tensor_map_l1_output = make_tma_2d_desc(l2_acts,
|
||||
intermediate_hidden, config.num_max_pool_tokens,
|
||||
config.block_n / 2, config.store_block_m,
|
||||
static_cast<int>(l2_acts.stride(-2)),
|
||||
64);
|
||||
const auto tensor_map_l2_acts = make_tma_2d_desc(l2_acts,
|
||||
intermediate_hidden, config.num_max_pool_tokens,
|
||||
config.block_k, config.block_m,
|
||||
static_cast<int>(l2_acts.stride(-2)),
|
||||
128);
|
||||
const auto tensor_map_l2_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l2_acts_sf,
|
||||
config.num_padded_sf_pool_tokens, intermediate_hidden,
|
||||
config.block_m, kGranK,
|
||||
1, 0);
|
||||
const auto tensor_map_l2_weights = make_tma_2d_desc(l2_weights,
|
||||
intermediate_hidden, num_experts_per_rank * hidden,
|
||||
config.block_k, config.block_n,
|
||||
static_cast<int>(l2_weights.stride(-2)),
|
||||
128);
|
||||
|
||||
// Stats can be optional
|
||||
int* cumulative_local_expert_recv_stats_ptr = nullptr;
|
||||
if (cumulative_local_expert_recv_stats.has_value())
|
||||
cumulative_local_expert_recv_stats_ptr = cumulative_local_expert_recv_stats->data_ptr<int>();
|
||||
|
||||
// Launch
|
||||
const auto num_sms = device_runtime->get_num_sms();
|
||||
const int num_threads = config.num_dispatch_threads + config.num_tma_threads + config.num_math_threads;
|
||||
const SM90FP8MegaMoERuntime::Args args = {
|
||||
.num_max_tokens_per_rank = num_max_tokens_per_rank,
|
||||
.hidden = hidden, .intermediate_hidden = intermediate_hidden,
|
||||
.num_experts = num_experts, .num_topk = num_topk,
|
||||
.num_ranks = num_ranks,
|
||||
.activation_clamp = activation_clamp,
|
||||
.fast_math = fast_math,
|
||||
.config = config,
|
||||
.y = y.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),
|
||||
.tensor_map_l1_acts = tensor_map_l1_acts,
|
||||
.tensor_map_l1_acts_sf = tensor_map_l1_acts_sf,
|
||||
.tensor_map_l1_weights = tensor_map_l1_weights,
|
||||
.tensor_map_l1_output = tensor_map_l1_output,
|
||||
.tensor_map_l2_acts = tensor_map_l2_acts,
|
||||
.tensor_map_l2_acts_sf = tensor_map_l2_acts_sf,
|
||||
.tensor_map_l2_weights = tensor_map_l2_weights,
|
||||
.l1_weights_sf = l1_weights_sf.data_ptr(),
|
||||
.l2_weights_sf = l2_weights_sf.data_ptr(),
|
||||
.launch_args = LaunchArgs(num_sms, num_threads, config.smem_size)
|
||||
};
|
||||
|
||||
const auto code = SM90FP8MegaMoERuntime::generate(args);
|
||||
const auto runtime = compiler->build("sm90_fp8_mega_moe", code);
|
||||
SM90FP8MegaMoERuntime::launch(runtime, args);
|
||||
}
|
||||
|
||||
} // namespace deep_gemm
|
||||
Reference in New Issue
Block a user