mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
781 lines
31 KiB
Plaintext
781 lines
31 KiB
Plaintext
/*
|
|
* SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
* SPDX-License-Identifier: Apache-2.0
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
|
|
// Include the fixture with the actual benchmark code
|
|
#include "mixtureOfExpertsBackendBenchmarkFixture.h"
|
|
|
|
/*
|
|
* Below is all the setup for parameterising the benchmarks
|
|
*/
|
|
|
|
template <class DataType_, class WeightType_ = DataType_, class OutputType_ = DataType_>
|
|
struct WeightParams
|
|
{
|
|
using DataType = DataType_;
|
|
using WeightType = WeightType_;
|
|
using OutputType = OutputType_;
|
|
};
|
|
|
|
#define BENCHMARK_BASIC(atype, wtype, otype) \
|
|
BENCHMARK_TEMPLATE_DEFINE_F(MixtureOfExpertsBenchmark, Basic_##atype##_##wtype, WeightParams<atype, wtype, otype>) \
|
|
(benchmark::State & state) \
|
|
{ \
|
|
runBenchmark(state); \
|
|
}
|
|
|
|
#define BENCHMARK_BASIC_DO_REGISTER(atype, wtype, otype) \
|
|
BENCHMARK_REGISTER_F(MixtureOfExpertsBenchmark, Basic_##atype##_##wtype) \
|
|
->Apply(argGen<MixtureOfExpertsBenchmark<WeightParams<atype, wtype, otype>>>)
|
|
|
|
template <class BenchClass>
|
|
auto listAllTactics()
|
|
{
|
|
int const sm = getSMVersion();
|
|
using RunnerType = decltype(BenchClass::mMoERunner);
|
|
return RunnerType::getTactics(sm);
|
|
}
|
|
|
|
template <class BenchClass>
|
|
int parseTacticToId(nlohmann::json tactic_config)
|
|
{
|
|
bool is_tma_warp_specialized = tactic_config.at("is_tma_warp_specialized").get<bool>();
|
|
int tile_shape_id = -1;
|
|
std::array<int, 3> tile_shape;
|
|
if (tactic_config.at("tile_shape").is_array())
|
|
tactic_config.at("tile_shape").get_to(tile_shape);
|
|
else
|
|
tile_shape_id = tactic_config.at("tile_shape").get<int>();
|
|
|
|
std::vector<tensorrt_llm::cutlass_extensions::CutlassGemmConfig> confs = listAllTactics<BenchClass>();
|
|
|
|
try
|
|
{
|
|
for (int i = 0; i < confs.size(); i++)
|
|
{
|
|
auto const& c = confs[i];
|
|
if (c.is_tma_warp_specialized != is_tma_warp_specialized)
|
|
continue;
|
|
|
|
if (!is_tma_warp_specialized)
|
|
{
|
|
int stages = tactic_config.at("stages").get<int>();
|
|
if (c.stages != stages)
|
|
continue;
|
|
}
|
|
|
|
if (tile_shape_id != -1)
|
|
{
|
|
int comp = c.getTileConfigAsInt();
|
|
if (tile_shape_id != comp)
|
|
continue;
|
|
if (is_tma_warp_specialized && (int) c.cluster_shape != tactic_config.at("cluster_shape").get<int>())
|
|
continue;
|
|
|
|
// Found matching config
|
|
return i;
|
|
}
|
|
|
|
// Handle if the user provided a shape instead of the enum value
|
|
if (is_tma_warp_specialized)
|
|
{
|
|
// TODO Add cases for blackwell shapes
|
|
using Kv = uint64_t;
|
|
constexpr static auto K = [](int m, int n) { return (uint64_t(m) << 32) | uint64_t(n); };
|
|
static std::unordered_map<Kv, CutlassTileConfigSM90> const tile_map{
|
|
{K(64, 16), CutlassTileConfigSM90::CtaShape64x16x128B},
|
|
{K(64, 32), CutlassTileConfigSM90::CtaShape64x32x128B},
|
|
{K(64, 64), CutlassTileConfigSM90::CtaShape64x64x128B},
|
|
{K(64, 128), CutlassTileConfigSM90::CtaShape64x128x128B},
|
|
{K(64, 256), CutlassTileConfigSM90::CtaShape64x256x128B},
|
|
|
|
{K(128, 16), CutlassTileConfigSM90::CtaShape128x16x128B},
|
|
{K(128, 32), CutlassTileConfigSM90::CtaShape128x32x128B},
|
|
{K(128, 64), CutlassTileConfigSM90::CtaShape128x64x128B},
|
|
{K(128, 128), CutlassTileConfigSM90::CtaShape128x128x128B},
|
|
{K(128, 256), CutlassTileConfigSM90::CtaShape128x256x128B},
|
|
{K(256, 128), CutlassTileConfigSM90::CtaShape256x128x128B},
|
|
};
|
|
|
|
if (c.getTileConfigAsInt() != (int) tile_map.at(K(tile_shape[0], tile_shape[1])))
|
|
continue;
|
|
|
|
static std::unordered_map<Kv, ClusterShape> const cluster_map{
|
|
// CTA configs for M=64
|
|
{K(1, 1), ClusterShape::ClusterShape_1x1x1},
|
|
{K(2, 1), ClusterShape::ClusterShape_2x1x1},
|
|
{K(1, 2), ClusterShape::ClusterShape_1x2x1},
|
|
{K(2, 2), ClusterShape::ClusterShape_2x2x1},
|
|
};
|
|
|
|
std::array<int, 3> cluster_shape;
|
|
tactic_config.at("cluster_shape").get_to(cluster_shape);
|
|
|
|
if (c.cluster_shape != cluster_map.at(K(cluster_shape[0], cluster_shape[1])))
|
|
continue;
|
|
|
|
// Found matching config
|
|
return i;
|
|
}
|
|
else
|
|
{
|
|
std::array<int, 3> warp_shape;
|
|
tactic_config.at("warp_shape").get_to(warp_shape);
|
|
|
|
using Kv = uint64_t;
|
|
constexpr static auto K = [](std::array<int, 3> a, std::array<int, 3> b)
|
|
{
|
|
uint64_t sum = 0;
|
|
for (auto v : a)
|
|
sum = sum * 512 + v;
|
|
for (auto v : b)
|
|
sum = sum * 256 + v;
|
|
return sum;
|
|
};
|
|
static std::unordered_map<Kv, CutlassTileConfig> tile_map{
|
|
{K({128, 128, 8}, {64, 64, 8}), CutlassTileConfig::CtaShape128x128x8_WarpShape64x64x8},
|
|
|
|
{K({16, 128, 64}, {16, 32, 64}), CutlassTileConfig::CtaShape16x128x64_WarpShape16x32x64},
|
|
{K({32, 128, 64}, {32, 32, 64}), CutlassTileConfig::CtaShape32x128x64_WarpShape32x32x64},
|
|
|
|
{K({64, 128, 64}, {32, 64, 64}), CutlassTileConfig::CtaShape64x128x64_WarpShape32x64x64},
|
|
{K({64, 64, 128}, {32, 64, 64}), CutlassTileConfig::CtaShape64x64x128_WarpShape32x64x64},
|
|
{K({64, 128, 64}, {64, 32, 64}), CutlassTileConfig::CtaShape64x128x64_WarpShape64x32x64},
|
|
|
|
{K({128, 64, 64}, {64, 32, 64}), CutlassTileConfig::CtaShape128x64x64_WarpShape64x32x64},
|
|
{K({128, 128, 64}, {64, 32, 64}), CutlassTileConfig::CtaShape128x128x64_WarpShape64x32x64},
|
|
{K({128, 128, 64}, {64, 64, 64}), CutlassTileConfig::CtaShape128x128x64_WarpShape64x64x64},
|
|
{K({128, 128, 64}, {64, 32, 64}), CutlassTileConfig::CtaShape128x128x64_WarpShape128x32x64},
|
|
{K({128, 256, 64}, {64, 64, 64}), CutlassTileConfig::CtaShape128x256x64_WarpShape64x64x64},
|
|
|
|
{K({256, 128, 64}, {64, 64, 64}), CutlassTileConfig::CtaShape256x128x64_WarpShape64x64x64},
|
|
|
|
{K({16, 256, 64}, {16, 64, 64}), CutlassTileConfig::CtaShape16x256x64_WarpShape16x64x64}
|
|
|
|
};
|
|
if (c.tile_config_sm80 != tile_map.at(K(tile_shape, warp_shape)))
|
|
continue;
|
|
|
|
// Found matching config
|
|
return i;
|
|
}
|
|
}
|
|
}
|
|
catch (std::out_of_range const& e)
|
|
{
|
|
std::cerr << "Warning: error parsing tactic " << tactic_config.dump(2) << std::endl;
|
|
}
|
|
|
|
return -1;
|
|
}
|
|
|
|
template <class BenchClass>
|
|
void parseTacticToVectorID(nlohmann::json& tactic, std::vector<int>& tactic_ids)
|
|
{
|
|
if (tactic.is_number_integer())
|
|
{
|
|
tactic_ids.push_back(tactic.get<int>());
|
|
}
|
|
else if (tactic.is_array())
|
|
{
|
|
for (auto c : tactic)
|
|
{
|
|
parseTacticToVectorID<BenchClass>(c, tactic_ids);
|
|
}
|
|
}
|
|
else if (tactic.is_object())
|
|
{
|
|
tactic_ids.push_back(parseTacticToId<BenchClass>(tactic));
|
|
}
|
|
else if (tactic.is_string())
|
|
{
|
|
assert(tactic.is_string());
|
|
auto tactic_name = tactic.get<std::string>();
|
|
if (tactic_name == "all")
|
|
{
|
|
auto all_tactics = listAllTactics<BenchClass>();
|
|
tactic_ids.resize(all_tactics.size());
|
|
std::iota(tactic_ids.begin(), tactic_ids.end(), 0);
|
|
}
|
|
else
|
|
{
|
|
assert(tactic.get<std::string>() == "auto");
|
|
tactic_ids.push_back(-1);
|
|
}
|
|
}
|
|
else
|
|
{
|
|
throw std::invalid_argument("Invalid tactic format");
|
|
}
|
|
}
|
|
|
|
// This interdependence of globals could be better, but it works ok for this limited case.
|
|
std::unordered_map<std::string, std::pair<int, int>> name_info_map{
|
|
{routingConfigCache[LOAD_BALANCED_ROUTING_CONFIG]->getName(), {-1, LOAD_BALANCED_ROUTING_CONFIG}},
|
|
{routingConfigCache[UNIFORM_ROUTING_CONFIG]->getName(), {-1, UNIFORM_ROUTING_CONFIG}},
|
|
};
|
|
|
|
int getNameCacheIdx(std::string const& name)
|
|
{
|
|
if (name_info_map.find(name) == name_info_map.end())
|
|
{
|
|
return -1;
|
|
}
|
|
return name_info_map.at(name).second;
|
|
}
|
|
|
|
void setNameCacheIdx(std::string const& name, int id)
|
|
{
|
|
name_info_map.at(name).second = id;
|
|
}
|
|
|
|
template <class ConfigType>
|
|
std::optional<int> loadRoutingValues(nlohmann::json entry, int64_t num_experts, int64_t k, std::string config_name)
|
|
{
|
|
std::optional<int> routing_config;
|
|
if (entry.is_string())
|
|
{
|
|
routing_config = getNameCacheIdx(entry.get<std::string>());
|
|
if (routing_config < 0)
|
|
{
|
|
throw std::invalid_argument("Invalid routing value, could not find valid config");
|
|
}
|
|
}
|
|
else
|
|
{
|
|
if (config_name.empty())
|
|
{
|
|
throw std::invalid_argument("Explicit routing configurations must specify a name");
|
|
}
|
|
|
|
std::vector<typename ConfigType::ElementType> values;
|
|
entry.get_to(values);
|
|
|
|
routingConfigCache.push_back(std::make_shared<ConfigType>(std::move(values), num_experts, k, config_name));
|
|
routing_config = routingConfigCache.size() - 1;
|
|
}
|
|
|
|
auto conf = routingConfigCache[*routing_config];
|
|
|
|
bool const is_supported = conf->supportsConfig(num_experts, k, {});
|
|
auto conf_derived = std::dynamic_pointer_cast<ConfigType>(conf);
|
|
auto conf_default = std::dynamic_pointer_cast<std::conditional_t<std::is_same_v<ConfigType, VectoredRoutingConfig>,
|
|
LoadBalancedRoutingConfig, UniformRoutingConfig>>(conf);
|
|
bool const is_valid_type = conf_derived || conf_default;
|
|
|
|
if (!is_supported || !is_valid_type)
|
|
{
|
|
throw std::invalid_argument("Incompatible config selected. "
|
|
+ ((conf_derived) ? "Expected experts: " + std::to_string(num_experts) + " and k: " + std::to_string(k)
|
|
+ " in routing configuration. Found: " + std::to_string(conf_derived->num_experts) + " and "
|
|
+ std::to_string(conf_derived->k)
|
|
: "Found incompatible routing config type"));
|
|
}
|
|
|
|
return routing_config;
|
|
}
|
|
|
|
// This is suboptimal for large benchmark files as we reread it for every data type
|
|
template <class BenchClass>
|
|
void argGenLoadFile(benchmark::internal::Benchmark* benchmark)
|
|
{
|
|
/*
|
|
* See help text for schema description
|
|
*/
|
|
|
|
std::ifstream file{workloadFile};
|
|
std::stringstream buffer;
|
|
buffer << file.rdbuf();
|
|
auto file_contents = buffer.str();
|
|
if (LOG_LEVEL >= INFO)
|
|
std::cout << "Loaded benchmark file: " << file_contents << std::endl;
|
|
auto source_data = nlohmann::json::parse(file_contents);
|
|
|
|
int i = 0;
|
|
for (auto run_config : source_data)
|
|
{
|
|
if (LOG_LEVEL >= VERBOSE)
|
|
std::cout << "Parsing run config: " << run_config.dump(2) << std::endl;
|
|
std::string config_name = "config_" + std::to_string(i);
|
|
|
|
// WARNING: Process the routing configuration immediately, so we can guarantee all configs get processed for all
|
|
// data types. We should not skip any test cases as a later test config may depend on this config
|
|
if (run_config.contains("routing_name"))
|
|
{
|
|
run_config["routing_name"].get_to(config_name);
|
|
if (!run_config.contains("selected_experts") && !run_config.contains("expert_distribution"))
|
|
{
|
|
throw std::invalid_argument("Setting routing value configuration name but missing routing values");
|
|
}
|
|
}
|
|
|
|
if (run_config.contains("routing_values") || run_config.contains("routing_distribution"))
|
|
{
|
|
throw std::invalid_argument(
|
|
"Using deprecated routing_values or routing_distribution. Please use selected_experts or "
|
|
"expert_distribution instead.");
|
|
}
|
|
|
|
std::optional<int> routing_config;
|
|
auto res = name_info_map.emplace(config_name, std::pair{i, -1});
|
|
// We must check i is not equal since this function gets called for each data type
|
|
if (!res.second && res.first->second.first != i)
|
|
{
|
|
throw std::invalid_argument("Redefinition of routing_name " + config_name + " at config "
|
|
+ std::to_string(i) + ". First declared at " + std::to_string(res.first->second.first));
|
|
}
|
|
else if (!res.second)
|
|
{
|
|
// Reuse the existing config from a previous parse
|
|
routing_config = getNameCacheIdx(config_name);
|
|
}
|
|
i++;
|
|
|
|
int num_experts = run_config.at("num_experts").get<int>();
|
|
int k = run_config.at("k").get<int>();
|
|
|
|
if (!routing_config)
|
|
{
|
|
if (run_config.contains("selected_experts"))
|
|
{
|
|
routing_config = loadRoutingValues<VectoredRoutingConfig>(
|
|
run_config["selected_experts"], num_experts, k, config_name);
|
|
}
|
|
else if (run_config.contains("expert_distribution"))
|
|
{
|
|
routing_config = loadRoutingValues<RandomDistributionRoutingConfig>(
|
|
run_config["expert_distribution"], num_experts, k, config_name);
|
|
}
|
|
}
|
|
// Use the selected config or fall back to balanced
|
|
routing_config = routing_config.value_or(LOAD_BALANCED_ROUTING_CONFIG);
|
|
setNameCacheIdx(config_name, *routing_config);
|
|
|
|
// Filter out the types we don't care about testing
|
|
if (run_config.contains("dtypes"))
|
|
{
|
|
std::vector<std::string> dtypes;
|
|
run_config["dtypes"].get_to(dtypes);
|
|
|
|
auto hasDtype = [&](char const* d)
|
|
{ return std::any_of(dtypes.begin(), dtypes.end(), [&](auto const& n) { return n == d; }); };
|
|
|
|
if (BenchClass::FP8 && !hasDtype("fp8"))
|
|
{
|
|
continue;
|
|
}
|
|
else if (BenchClass::NVFP4 && !hasDtype("fp4"))
|
|
{
|
|
continue;
|
|
}
|
|
else if (BenchClass::INT4 && !hasDtype("int4"))
|
|
{
|
|
continue;
|
|
}
|
|
else if (!BenchClass::INT4 && BenchClass::INT_QUANT && !hasDtype("int8"))
|
|
{
|
|
continue;
|
|
}
|
|
else if (std::is_same_v<typename BenchClass::WeightType, float> && !hasDtype("float")
|
|
&& !hasDtype("float32"))
|
|
{
|
|
continue;
|
|
}
|
|
else if (std::is_same_v<typename BenchClass::WeightType, half> && !hasDtype("float16") && !hasDtype("half"))
|
|
{
|
|
continue;
|
|
}
|
|
else if (std::is_same_v<typename BenchClass::WeightType, __nv_bfloat16> && !hasDtype("bfloat16")
|
|
&& !hasDtype("bf16"))
|
|
{
|
|
continue;
|
|
}
|
|
else if (BenchClass::WFP4AFP8 && !hasDtype("wfp4afp8"))
|
|
{
|
|
continue;
|
|
}
|
|
}
|
|
|
|
// Do this after filtering datatypes as tactics only make sense if we know the data type
|
|
bool has_tactic_ids2 = false;
|
|
std::vector<int> tactic_ids1{};
|
|
std::vector<int> tactic_ids2{};
|
|
if (run_config.contains("tactic_id1") || run_config.contains("tactic_id2"))
|
|
{
|
|
if (run_config.contains("tactic_id"))
|
|
{
|
|
throw std::invalid_argument("Cannot use tactic_id and tactic_idX");
|
|
}
|
|
has_tactic_ids2 = true;
|
|
parseTacticToVectorID<BenchClass>(run_config["tactic_id1"], tactic_ids1);
|
|
parseTacticToVectorID<BenchClass>(run_config["tactic_id2"], tactic_ids2);
|
|
}
|
|
else
|
|
{
|
|
parseTacticToVectorID<BenchClass>(run_config["tactic_id"], tactic_ids1);
|
|
has_tactic_ids2 = false;
|
|
tactic_ids2.resize(1); // Dummy value so we loop exactly once below
|
|
}
|
|
if (tactic_ids1.empty() || tactic_ids2.empty())
|
|
{
|
|
std::cerr << "Warning: Skipping benchmark, no valid tactic found" << std::endl;
|
|
static bool printed = false;
|
|
if (!printed)
|
|
{
|
|
printed = true;
|
|
std::cerr << __PRETTY_FUNCTION__ << ": Valid Tactics are:\n";
|
|
auto confs = listAllTactics<BenchClass>();
|
|
for (auto c : confs)
|
|
std::cerr << c.toString();
|
|
}
|
|
|
|
continue;
|
|
}
|
|
|
|
auto get_or = [&](auto name, auto def)
|
|
{ return run_config.contains(name) ? run_config[name].template get<decltype(def)>() : def; };
|
|
int tp_size = get_or("tp_size", 1);
|
|
int ep_size = get_or("ep_size", 1);
|
|
int world_rank = get_or("world_rank", 0);
|
|
int bias = get_or("bias", 0);
|
|
int do_final_scale = get_or("do_final_scale", 1); // Default to scales on
|
|
TLLM_CHECK_WITH_INFO(world_rank < tp_size * ep_size, "Rank is out of bounds of tp*ep");
|
|
|
|
auto get_range = [&](std::string name, int min = 1, int max = INT32_MAX)
|
|
{
|
|
auto val = run_config.at(name).get<int>();
|
|
if (val < min || val > max)
|
|
{
|
|
throw std::invalid_argument(name + " must be a positive integer");
|
|
}
|
|
return val;
|
|
};
|
|
|
|
for (auto t1 : tactic_ids1)
|
|
{
|
|
// tactic_ids2 will have one dummy value if has_tactic_ids2 = false
|
|
for (auto t2 : tactic_ids2)
|
|
{
|
|
if (!has_tactic_ids2)
|
|
t2 = t1;
|
|
|
|
benchmark->Args({num_experts, //
|
|
get_range("k"), //
|
|
get_range("hidden_size"), //
|
|
get_range("inter_size"), //
|
|
tp_size, ep_size, world_rank, //
|
|
get_range("num_tokens"), //
|
|
bias, do_final_scale, //
|
|
get_range("act_fn", 0, (int) ActivationType::Identity), //
|
|
t1, //
|
|
t2, //
|
|
*routing_config});
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <class BenchClass>
|
|
void argGenHardcoded(benchmark::internal::Benchmark* benchmark)
|
|
{
|
|
auto num_experts = {1, 8, 9, 64, 65, 257}; // {1, 8, 64, 65, 1024};
|
|
auto top_k = {1, 2, 3, 16}; // {1, 2, 3, 42};
|
|
auto hidden_size = {4096};
|
|
auto inter_size_mul = {4.f}; // {7.f/2.f, 4.f};
|
|
auto num_tokens = {2048}; // {1, 20, 200, 2048};
|
|
auto use_bias = {0}; // {0, 1};
|
|
auto activation_type = {ActivationType::Gelu};
|
|
// {ActivationType::Relu, ActivationType::Gelu,
|
|
// ActivationType::Silu, ActivationType::Geglu,
|
|
// ActivationType::Swiglu};
|
|
auto cutlass_tactic = {-1}; // {0,..., listAllTactics<BenchClass>().size()};
|
|
auto routing_config = {LOAD_BALANCED_ROUTING_CONFIG}; // {0, 1, 2};
|
|
|
|
for (auto num_expert : num_experts)
|
|
for (auto k : top_k)
|
|
if (k <= num_expert)
|
|
for (auto size : hidden_size)
|
|
for (auto inter_mul : inter_size_mul)
|
|
{
|
|
auto inter_size = static_cast<int>(size * inter_mul);
|
|
for (auto tokens : num_tokens)
|
|
for (auto bias : use_bias)
|
|
for (auto act : activation_type)
|
|
for (auto tactic1 : cutlass_tactic)
|
|
for (auto tactic2 : cutlass_tactic)
|
|
for (auto routing : routing_config)
|
|
benchmark->Args({num_expert, k, size, inter_size, 1, 1, 0, tokens, bias,
|
|
1, (int) act, tactic1, tactic2, routing});
|
|
}
|
|
}
|
|
|
|
template <class BenchClass>
|
|
void argGen(benchmark::internal::Benchmark* benchmark)
|
|
{
|
|
if (LOG_LEVEL >= VERBOSE)
|
|
{
|
|
std::cout << "List of all tactics for dtype " << (int) BenchClass::toDTypeID() << ":\n";
|
|
int i = 0;
|
|
for (auto& t : listAllTactics<BenchClass>())
|
|
{
|
|
std::cout << "Tactic " << i << ":\n";
|
|
std::cout << t.toString() << std::endl;
|
|
|
|
i++;
|
|
}
|
|
}
|
|
|
|
// Generic setup
|
|
benchmark->UseManualTime();
|
|
benchmark->ArgNames(
|
|
{"Num Experts", "K", "Hidden Size", "Inter Size", "TP Size", "EP Size", "World Rank", "Num Tokens", "Use Bias",
|
|
"Use Final Scale", "Activation Function", "Tactic ID 1", "Tactic ID 2", "Routing ID"});
|
|
|
|
if (workloadFile)
|
|
argGenLoadFile<BenchClass>(benchmark);
|
|
else
|
|
argGenHardcoded<BenchClass>(benchmark);
|
|
}
|
|
|
|
BENCHMARK_BASIC(float, float, float)
|
|
BENCHMARK_BASIC(half, half, half)
|
|
using uint8 = uint8_t;
|
|
BENCHMARK_BASIC(half, uint8, half)
|
|
using cutlass::uint4b_t;
|
|
BENCHMARK_BASIC(half, uint4b_t, half)
|
|
#ifdef ENABLE_BF16
|
|
BENCHMARK_BASIC(nv_bfloat16, nv_bfloat16, nv_bfloat16)
|
|
#endif
|
|
#ifdef ENABLE_FP8
|
|
BENCHMARK_BASIC(SafeFP8, SafeFP8, half)
|
|
#endif
|
|
#ifdef ENABLE_FP4
|
|
BENCHMARK_BASIC(SafeFP4, SafeFP4, half)
|
|
BENCHMARK_BASIC(SafeFP8, SafeFP4, half)
|
|
#endif
|
|
|
|
void delayedRegisterBenchmark()
|
|
{
|
|
BENCHMARK_BASIC_DO_REGISTER(half, half, half);
|
|
#ifdef ENABLE_FP8
|
|
BENCHMARK_BASIC_DO_REGISTER(SafeFP8, SafeFP8, half);
|
|
#endif
|
|
if (workloadFile)
|
|
{
|
|
// Extra ones we don't want for hardcoded runs
|
|
BENCHMARK_BASIC_DO_REGISTER(float, float, float);
|
|
BENCHMARK_BASIC_DO_REGISTER(half, uint8, half);
|
|
BENCHMARK_BASIC_DO_REGISTER(half, uint4b_t, half);
|
|
#ifdef ENABLE_BF16
|
|
BENCHMARK_BASIC_DO_REGISTER(nv_bfloat16, nv_bfloat16, nv_bfloat16);
|
|
#endif
|
|
#ifdef ENABLE_FP4
|
|
BENCHMARK_BASIC_DO_REGISTER(SafeFP4, SafeFP4, half);
|
|
BENCHMARK_BASIC_DO_REGISTER(SafeFP8, SafeFP4, half);
|
|
#endif
|
|
}
|
|
}
|
|
|
|
void doCleanup()
|
|
{
|
|
bufferManager.reset();
|
|
streamPtr.reset();
|
|
}
|
|
|
|
void help()
|
|
{
|
|
std::cout << "Usage: mixtureOfExpertsBackendBenchmark [--disable_cuda_graphs] [--input_file <file>] [benchmark "
|
|
"options]\n";
|
|
std::cout
|
|
<< "--disable_cuda_graphs\t\tPrevent the benchmark from using cuda graphs. Useful for getting the performance "
|
|
"of specific kernels\n"
|
|
<< "--input_file\t\tA JSON file describing the benchmark configurations\n\n"
|
|
<< "File schema\n"
|
|
"[\n"
|
|
" {\n"
|
|
" \"num_experts\": int,\n"
|
|
" \"k\": int,\n"
|
|
" \"hidden_size\": int,\n"
|
|
" \"inter_size\": int,\n"
|
|
" \"tp_size\": int, (optional)\n"
|
|
" \"ep_size\": int, (optional)\n"
|
|
" \"world_rank\": int, (optional)\n"
|
|
" \"num_tokens\": int,\n"
|
|
" \"bias\": int, (optional)\n"
|
|
" \"do_final_scale\": int, (optional)\n"
|
|
" \"act_fn\": int,\n"
|
|
" \"tactic_id\": tactic, (see below)\n"
|
|
" \"tactic_id1\": tactic, (see below)\n"
|
|
" \"tactic_id2\": tactic, (see below)\n"
|
|
" \"dtypes\": [string, ...], (optional)\n"
|
|
" \"routing_name\": string, (optional)\n"
|
|
" \"selected_experts\": [int, ...], or string, (optional, length is a multiple of k)\n"
|
|
" \"expert_distribution\": [float, ...], or string, (optional, length is num_experts)\n"
|
|
" },\n"
|
|
" ...\n"
|
|
"]\n"
|
|
"Explanation:\n"
|
|
"- \"num_experts\" - The number of experts\n"
|
|
"- \"k\" - The top k\n"
|
|
"- \"hidden_size\" - The hidden size\n"
|
|
"- \"inter_size\" - The inter size\n"
|
|
"- \"tp_size\" - The TP size to use\n"
|
|
"- \"ep_size\" - The EP size to use\n"
|
|
"- \"world_rank\" - The world rank = tp_rank * ep_size + ep_rank\n"
|
|
"- \"num_tokens\" - The total number of tokens to benchmark\n"
|
|
"- \"bias\" - If bias should be used, 0 = no bias, 1 = bias\n"
|
|
"- \"do_final_scale\" - If final scales should be applied, 0 = no scale, 1 = scale\n"
|
|
"- \"act_fn\" - The activation function to use, 0 = identity, 1 = relu, 2 = gelu, 3 = silu, 4 = geglu, 5 = "
|
|
"swiglu\n"
|
|
"- \"tactic_id, tactic_id1, tactic_id2\"\n"
|
|
"The config for the CUTLASS GEMM. tactic_id sets the same tactic for both to the same tactic (except in "
|
|
"auto mode)\n"
|
|
"Use tactic_idX to set the tactic for the corresponding GEMM"
|
|
"Valid tactics are:\n"
|
|
" - An object:\n"
|
|
" {\n"
|
|
" \"is_tma_warp_specialized\": bool,\n"
|
|
" \"tile_shape\": [int, int, int] or int,\n"
|
|
" \"cluster_shape\": [int, int, int] or int, (required for sm90, type must be an int if tile_shape "
|
|
"is "
|
|
"an int)\n"
|
|
" \"warp_shape\": [int, int, int], (required for non-sm90 if tile_shape is an array)\n"
|
|
" \"stages\": int, (required for non-sm90)\n"
|
|
" },\n"
|
|
" - An integer: corresponds to an index in the tactics array. WARNING this is not stable between test "
|
|
"configurations\n"
|
|
" - An array: of integers or objects, forms a list of tactics to sweep\n"
|
|
" - The string \"all\": This will sweep through all possible tactics\n"
|
|
" - The string \"auto\": This runs a short benchmark to pick the fastest tactic before each benchmark "
|
|
"case. "
|
|
"Useful for quick perf tests, prefer a full sweep and manually setting the tactic for more accurate "
|
|
"results"
|
|
"- dtypes - A list of dtypes to run this config through.\n"
|
|
"Allowed values are: fp8, fp4, wfp4afp8, int4, int8, float, half, bfloat16\n"
|
|
"If this argument is omitted all dtypes will be run. Note, not all tactics are supported for all "
|
|
"dtypes,\n"
|
|
"unsupported tactics will be skipped with a warning.\n"
|
|
"- \"routing_name\" - a name to help identify the routing pattern. This can be used by later "
|
|
"benchmarks to reuse the config\n"
|
|
"- \"selected_experts\" - a flat array of selected experts to define a new config,\n or a string "
|
|
"referencing the name of a previous config. Defaults to pre-defined config \"balanced\",\n"
|
|
"which is short-hand for a perfectly balanced expert distribution\n"
|
|
"These define the routing values used as input to the moe backend, and is intended to allow comparing "
|
|
"different routing behaviours.\n"
|
|
"When defining an array, it must have `T*k` floating point values. Each set of\n"
|
|
"`k` values defines the input for a single token. If `num_tokens` is greater than `T` it will "
|
|
"repeat from the beginning\n"
|
|
"- \"expert_distribution\" - instead of explicitly setting selected_experts, define a random distribution "
|
|
"that experts will be randomly sampled from."
|
|
"There is also pre-defined config \"uniform\", which is short-hand for a random uniform distribution\n"
|
|
"\n";
|
|
|
|
std::cout << "benchmark options:\n";
|
|
benchmark::PrintDefaultHelp();
|
|
}
|
|
|
|
void gbenchCustomHelp()
|
|
{
|
|
help();
|
|
// google-benchmark calls exit() so we need to cleanup manually
|
|
doCleanup();
|
|
}
|
|
|
|
int parseArgsAndRunBench(int argc, char** argv)
|
|
{
|
|
try
|
|
{
|
|
int shift = 0;
|
|
for (int i = 1; i < argc; i++)
|
|
{
|
|
argv[i - shift] = argv[i];
|
|
if (strcmp("--input_file", argv[i]) == 0)
|
|
{
|
|
i += 1;
|
|
if (i == argc)
|
|
{
|
|
std::cerr << "Missing file name for input_file\n";
|
|
return -1;
|
|
}
|
|
workloadFile = argv[i];
|
|
if (workloadFile[0] == '-')
|
|
{
|
|
std::cerr << "Workload file " << workloadFile << " not a valid file name\n";
|
|
return -2;
|
|
}
|
|
shift += 2;
|
|
}
|
|
else if (strcmp("--disable_cuda_graphs", argv[i]) == 0)
|
|
{
|
|
useCudaGraph = false;
|
|
shift++;
|
|
}
|
|
else if (strcmp("--help", argv[i]) == 0 || strcmp("-h", argv[i]) == 0)
|
|
{
|
|
help();
|
|
return 0;
|
|
}
|
|
}
|
|
argc -= shift;
|
|
|
|
// Delay after we know if the user passed a config file
|
|
delayedRegisterBenchmark();
|
|
|
|
benchmark::Initialize(&argc, argv, &gbenchCustomHelp);
|
|
|
|
if (argc > 1)
|
|
{
|
|
help();
|
|
std::cout << std::flush; // Force flush
|
|
// Print the error second, so it's easy to see
|
|
std::cerr << "\nUnrecognised argument: " << argv[1] << std::endl;
|
|
return -4;
|
|
}
|
|
|
|
benchmark::RunSpecifiedBenchmarks();
|
|
benchmark::Shutdown();
|
|
|
|
return 0;
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
std::cerr << "Exiting benchmarks with exception: " << e.what() << std::endl;
|
|
return -3;
|
|
}
|
|
}
|
|
|
|
int main(int argc, char** argv)
|
|
{
|
|
deviceCount = getDeviceCount();
|
|
if (deviceCount < 0)
|
|
return 0;
|
|
streamPtr = std::make_shared<CudaStream>();
|
|
bufferManager = std::make_unique<BufferManager>(streamPtr);
|
|
|
|
int res = -1;
|
|
try
|
|
{
|
|
res = parseArgsAndRunBench(argc, argv);
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
std::cout << "Benchmark exited with unhandled exception: " << e.what() << std::endl;
|
|
}
|
|
|
|
doCleanup();
|
|
return res;
|
|
}
|