feat: add sm90 megamoe phase1 interfaces

2026-06-17 23:54:49 +08:00
parent 062cb160cf
commit 1dafd0f173
12 changed files with 1520 additions and 130 deletions
@@ -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