235 lines
11 KiB
C++
235 lines
11 KiB
C++
#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/mega_moe.hpp"
|
||
|
||
namespace deep_gemm {
|
||
|
||
class SM100FP8NVFP4MegaMoERuntime final : public LaunchRuntime<SM100FP8NVFP4MegaMoERuntime> {
|
||
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;
|
||
MegaMoEConfig 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_weights_sf;
|
||
CUtensorMap tensor_map_l1_output;
|
||
CUtensorMap tensor_map_l2_acts;
|
||
CUtensorMap tensor_map_l2_acts_sf;
|
||
CUtensorMap tensor_map_l2_weights;
|
||
CUtensorMap tensor_map_l2_weights_sf;
|
||
|
||
// Launch configs
|
||
LaunchArgs launch_args;
|
||
};
|
||
|
||
static std::string generate_impl(const Args& args) {
|
||
return fmt::format(R"(
|
||
#include <deep_gemm/impls/sm100_fp8_nvfp4_mega_moe.cuh>
|
||
|
||
using namespace deep_gemm;
|
||
|
||
static void __instantiate_kernel() {{
|
||
auto ptr = reinterpret_cast<void*>(&sm100_fp8_nvfp4_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.sf_block_m, args.config.sf_block_n,
|
||
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_non_epilogue_threads, args.config.num_epilogue_threads,
|
||
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) {
|
||
// Debug: sync before launch to flush TMA creation errors
|
||
if (get_env<int>("DG_JIT_DEBUG")) {
|
||
auto err = cudaDeviceSynchronize();
|
||
printf("[MEGA_MOE_LAUNCH_DEBUG] Pre-launch sync: %d, about to launch kernel, y=%p, num_tokens=%d\n", (int)err, args.y, args.num_tokens);
|
||
fflush(stdout);
|
||
}
|
||
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.tensor_map_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.tensor_map_l2_weights_sf
|
||
));
|
||
}
|
||
};
|
||
|
||
static void sm100_fp8_nvfp4_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_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);
|
||
|
||
// NVFP4: kGranK=16 for group_size=16
|
||
constexpr int kGranK = 16;
|
||
|
||
// Make tensormap — NVFP4 packed E2M1 (2 per byte), so K-dim is hidden/2 bytes
|
||
// L1 activations: packed E2M1 (4 bits/element), K-dim = hidden/2
|
||
const auto tensor_map_l1_acts = make_tma_2d_desc(l1_acts,
|
||
hidden / 2, config.num_max_pool_tokens,
|
||
config.block_k / 2, config.load_block_m,
|
||
static_cast<int>(l1_acts.stride(-2)),
|
||
config.swizzle_acts_mode / 2,
|
||
0, false, false); // fp4_unpacked_smem=false
|
||
// NVFP4 SF TMA: kGranK=16, so SF K-dim is hidden/16, packed as hidden/64 int32
|
||
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.sf_block_m, kGranK,
|
||
1, 0);
|
||
const auto tensor_map_l1_weights = make_tma_2d_desc(l1_weights,
|
||
hidden / 2, num_experts_per_rank * intermediate_hidden * 2,
|
||
config.block_k / 2, config.load_block_n,
|
||
static_cast<int>(l1_weights.stride(-2)),
|
||
config.swizzle_weights_mode / 2,
|
||
0, false, false); // fp4_unpacked_smem=false (packed!)
|
||
const auto tensor_map_l1_weights_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l1_weights_sf,
|
||
intermediate_hidden * 2, hidden,
|
||
config.block_n, kGranK,
|
||
num_experts_per_rank, 0);
|
||
// L1 output: packed E2M1, K-dim = intermediate_hidden/2, inner = block_n/4 bytes (SwiGLU halving × FP4 packing), no swizzle (v1)
|
||
const auto tensor_map_l1_output = make_tma_2d_desc(l2_acts,
|
||
intermediate_hidden / 2, config.num_max_pool_tokens,
|
||
config.block_n / 4, config.store_block_m,
|
||
static_cast<int>(l2_acts.stride(-2)),
|
||
0, 0, // no swizzle
|
||
false, // allow_tf32
|
||
false); // fp4_unpacked_smem=false (packed!)
|
||
// L2 activations: packed E2M1, K-dim = intermediate_hidden/2
|
||
const auto tensor_map_l2_acts = make_tma_2d_desc(l2_acts,
|
||
intermediate_hidden / 2, config.num_max_pool_tokens,
|
||
config.block_k / 2, config.load_block_m,
|
||
static_cast<int>(l2_acts.stride(-2)),
|
||
config.swizzle_acts_mode / 2,
|
||
0, false, false); // fp4_unpacked_smem=false
|
||
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.sf_block_m, kGranK,
|
||
1, 0);
|
||
const auto tensor_map_l2_weights = make_tma_2d_desc(l2_weights,
|
||
intermediate_hidden / 2, num_experts_per_rank * hidden,
|
||
config.block_k / 2, config.load_block_n,
|
||
static_cast<int>(l2_weights.stride(-2)),
|
||
config.swizzle_weights_mode / 2,
|
||
0, false, false); // fp4_unpacked_smem=false (packed!)
|
||
const auto tensor_map_l2_weights_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l2_weights_sf,
|
||
hidden, intermediate_hidden,
|
||
config.block_n, kGranK,
|
||
num_experts_per_rank, 0);
|
||
|
||
// 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 SM100FP8NVFP4MegaMoERuntime::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_weights_sf = tensor_map_l1_weights_sf,
|
||
.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,
|
||
.tensor_map_l2_weights_sf = tensor_map_l2_weights_sf,
|
||
.launch_args = LaunchArgs(num_sms,
|
||
config.num_dispatch_threads + config.num_non_epilogue_threads + config.num_epilogue_threads,
|
||
config.smem_size, 2)
|
||
};
|
||
|
||
const auto code = SM100FP8NVFP4MegaMoERuntime::generate(args);
|
||
const auto runtime = compiler->build("sm100_fp8_nvfp4_mega_moe", code);
|
||
SM100FP8NVFP4MegaMoERuntime::launch(runtime, args);
|
||
}
|
||
|
||
} // namespace deep_gemm
|