79fcfd6abf
Allow SM100 FP4 scale layout transforms to accept group16 and thread weight granularity through the MegaMoE Python wrapper, API checks, and synthetic benchmark entrypoint. Keep fused SM100 MegaMoE compute behind an explicit group16 capability gate until the SFB/TMEM/MMA scale path is updated and validated. Tested: PYTHONPYCACHEPREFIX=/private/tmp/deepgemm_pycache python3 -m py_compile deep_gemm/mega/__init__.py tests/test_mega_moe.py tests/generators.py Tested: git diff --check Not-tested: CUDA build and SM100/B300 runtime validation are not available locally.
226 lines
10 KiB
C++
226 lines
10 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 SM100FP8FP4MegaMoERuntime final : public LaunchRuntime<SM100FP8FP4MegaMoERuntime> {
|
|
public:
|
|
struct Args {
|
|
// Templated arguments
|
|
int num_max_tokens_per_rank;
|
|
int hidden, intermediate_hidden;
|
|
int num_experts, num_topk;
|
|
int weight_gran_k;
|
|
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_fp4_mega_moe.cuh>
|
|
|
|
using namespace deep_gemm;
|
|
|
|
static void __instantiate_kernel() {{
|
|
auto ptr = reinterpret_cast<void*>(&sm100_fp8_fp4_mega_moe_impl<
|
|
{},
|
|
{}, {},
|
|
{}, {},
|
|
{},
|
|
{},
|
|
{}, {}, {},
|
|
{},
|
|
{}, {},
|
|
{},
|
|
{},
|
|
{},
|
|
{}, {}, {},
|
|
{}, {},
|
|
{},
|
|
{}
|
|
>);
|
|
}};
|
|
)", args.num_max_tokens_per_rank,
|
|
args.hidden, args.intermediate_hidden,
|
|
args.num_experts, args.num_topk,
|
|
args.weight_gran_k,
|
|
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) {
|
|
// TODO: optimize `args` copy
|
|
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_fp4_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 int& weight_gran_k,
|
|
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);
|
|
|
|
// Make tensormap
|
|
constexpr int kActivationGranK = 32;
|
|
const auto tensor_map_l1_acts = make_tma_2d_desc(l1_acts,
|
|
hidden, config.num_max_pool_tokens,
|
|
config.block_k, config.load_block_m,
|
|
static_cast<int>(l1_acts.stride(-2)),
|
|
config.swizzle_acts_mode);
|
|
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, kActivationGranK,
|
|
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.load_block_n,
|
|
static_cast<int>(l1_weights.stride(-2)),
|
|
config.swizzle_weights_mode);
|
|
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, weight_gran_k,
|
|
num_experts_per_rank, 0);
|
|
// NOTES: L1 output and L2 activations are essentially the same tensor.
|
|
// Post-SwiGLU output has half the N width (`BLOCK_N / 2` per input tile),
|
|
// so the swizzle mode is also halved (128 -> 64).
|
|
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)),
|
|
config.swizzle_acts_mode / 2);
|
|
const auto tensor_map_l2_acts = make_tma_2d_desc(l2_acts,
|
|
intermediate_hidden, config.num_max_pool_tokens,
|
|
config.block_k, config.load_block_m,
|
|
static_cast<int>(l2_acts.stride(-2)),
|
|
config.swizzle_acts_mode);
|
|
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, kActivationGranK,
|
|
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.load_block_n,
|
|
static_cast<int>(l2_weights.stride(-2)),
|
|
config.swizzle_weights_mode);
|
|
const auto tensor_map_l2_weights_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l2_weights_sf,
|
|
hidden, intermediate_hidden,
|
|
config.block_n, weight_gran_k,
|
|
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 SM100FP8FP4MegaMoERuntime::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,
|
|
.weight_gran_k = weight_gran_k,
|
|
.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 = SM100FP8FP4MegaMoERuntime::generate(args);
|
|
const auto runtime = compiler->build("sm100_fp8_fp4_mega_moe", code);
|
|
SM100FP8FP4MegaMoERuntime::launch(runtime, args);
|
|
}
|
|
|
|
} // namespace deep_gemm
|