[Public release 26/04] Introducing Mega MoE, FP4 Indexer and other features/fixes (#304)
* Merge with private repo * Update README * Update README * Update README * Add PyTorch requirements * Fix sync scopes for MQA logits (#256) * Update README
This commit is contained in:
Chenggang Zhao
@@ -1,339 +1,54 @@
|
||||
#pragma once
|
||||
|
||||
#include <deep_gemm/common/types.hpp>
|
||||
#include <unordered_set>
|
||||
#include <deep_gemm/common/types.cuh>
|
||||
|
||||
#include "../../utils/math.hpp"
|
||||
#include "config.hpp"
|
||||
#include "runtime.hpp"
|
||||
#include "../../utils/layout.hpp"
|
||||
#include "../../utils/system.hpp"
|
||||
|
||||
namespace deep_gemm {
|
||||
|
||||
struct MulticastConfig {
|
||||
int num_multicast;
|
||||
bool is_multicast_on_a;
|
||||
|
||||
MulticastConfig(const int& num_multicast, const bool& is_multicast_on_a):
|
||||
num_multicast(num_multicast), is_multicast_on_a(is_multicast_on_a) {
|
||||
DG_HOST_ASSERT(1 <= num_multicast and num_multicast <= 2);
|
||||
}
|
||||
};
|
||||
|
||||
struct SharedMemoryConfig {
|
||||
int smem_size;
|
||||
int swizzle_a_mode;
|
||||
int swizzle_b_mode;
|
||||
int swizzle_cd_mode;
|
||||
};
|
||||
|
||||
struct ThreadConfig {
|
||||
int num_threads;
|
||||
|
||||
// SM90
|
||||
int num_tma_threads;
|
||||
int num_math_threads;
|
||||
|
||||
// SM100
|
||||
int num_non_epilogue_threads;
|
||||
int num_epilogue_threads;
|
||||
|
||||
static ThreadConfig sm90(const int& num_tma_threads,
|
||||
const int& num_math_threads) {
|
||||
auto config = ThreadConfig();
|
||||
config.num_threads = num_tma_threads + num_math_threads;
|
||||
config.num_tma_threads = num_tma_threads;
|
||||
config.num_math_threads = num_math_threads;
|
||||
return config;
|
||||
}
|
||||
|
||||
static ThreadConfig sm100(const int& num_non_epilogue_threads,
|
||||
const int& num_epilogue_threads) {
|
||||
auto config = ThreadConfig();
|
||||
config.num_threads = num_non_epilogue_threads + num_epilogue_threads;
|
||||
config.num_non_epilogue_threads = num_non_epilogue_threads;
|
||||
config.num_epilogue_threads = num_epilogue_threads;
|
||||
return config;
|
||||
}
|
||||
};
|
||||
|
||||
struct GemmConfig {
|
||||
// Templated configs
|
||||
GemmType gemm_type;
|
||||
KernelType kernel_type;
|
||||
MmaKind mma_kind;
|
||||
at::ScalarType a_dtype, b_dtype, cd_dtype;
|
||||
cute::UMMA::Major major_a;
|
||||
cute::UMMA::Major major_b;
|
||||
bool with_accumulation;
|
||||
int block_m, block_n, block_k;
|
||||
int num_stages, num_last_stages;
|
||||
|
||||
// Templated device configs
|
||||
int num_sms;
|
||||
int tc_util;
|
||||
|
||||
// Structured configs
|
||||
MulticastConfig multicast_config;
|
||||
SharedMemoryConfig smem_config;
|
||||
ThreadConfig thread_config;
|
||||
};
|
||||
|
||||
static bool is_multicast_legal(const int& shape_dim, const int& block_dim,
|
||||
const int& num_multicast, const int& num_sms,
|
||||
const bool& require_divisible) {
|
||||
const bool& divisible = ceil_div(shape_dim, block_dim) % num_multicast == 0 or not require_divisible;
|
||||
return divisible and num_sms % num_multicast == 0;
|
||||
}
|
||||
|
||||
template <typename size_type_t>
|
||||
static int get_swizzle_mode(const int& block_size, const size_type_t& elem_size) {
|
||||
// `> 0` means interleaving
|
||||
// 16B actually means non-swizzling (but interleaving)
|
||||
for (const int& mode: {128, 64, 32, 16}) {
|
||||
if ((block_size * static_cast<int>(elem_size)) % mode == 0)
|
||||
return mode;
|
||||
}
|
||||
DG_HOST_UNREACHABLE("Unreachable");
|
||||
}
|
||||
|
||||
template <typename ArchSpec>
|
||||
static SharedMemoryConfig get_smem_config(const GemmType& gemm_type, const KernelType& kernel_type,
|
||||
const int& m, const int& n, const int& k,
|
||||
const int& block_m, const int& block_n, const int& block_k,
|
||||
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
|
||||
const MmaKind& mma_kind, const at::ScalarType& cd_dtype,
|
||||
const int& num_stages, const MulticastConfig& multicast_config) {
|
||||
const int& ab_elem_size = static_cast<int>(get_element_size(mma_kind));
|
||||
const int& cd_elem_size = static_cast<int>(c10::elementSize(cd_dtype));
|
||||
static GemmConfig get_best_config(const GemmDesc& desc) {
|
||||
desc.check_validity();
|
||||
|
||||
const int& load_block_m = ArchSpec::get_ab_load_block_m(multicast_config, block_m);
|
||||
const int& load_block_n = ArchSpec::get_ab_load_block_n(multicast_config, block_n);
|
||||
const int& swizzle_a_mode = get_swizzle_mode(major_a == cute::UMMA::Major::K ? block_k : load_block_m, ab_elem_size);
|
||||
const int& swizzle_b_mode = get_swizzle_mode(major_b == cute::UMMA::Major::K ? block_k : load_block_n, ab_elem_size);
|
||||
const int& swizzle_cd_mode = ArchSpec::enable_cd_swizzle(cd_dtype) ? get_swizzle_mode(block_n, cd_elem_size) : 0;
|
||||
|
||||
// Different archs have different epilogue pipelines
|
||||
const int& smem_cd = ArchSpec::get_smem_cd_size(kernel_type, block_m, block_n, swizzle_cd_mode, cd_dtype);
|
||||
|
||||
// A/B shared memory
|
||||
const int& smem_a_per_stage = load_block_m * block_k * ab_elem_size;
|
||||
const int& smem_b_per_stage = load_block_n * block_k * ab_elem_size;
|
||||
|
||||
// SF shared memory
|
||||
const auto& [smem_sfa_per_stage, smem_sfb_per_stage] =
|
||||
ArchSpec::get_sf_smem_size_per_stage(kernel_type, block_m, block_n, block_k, mma_kind, cd_dtype);
|
||||
const int& smem_extra_sfb = ArchSpec::get_extra_sfb_smem_size(m, n, k, block_m, block_n, block_k);
|
||||
|
||||
// M-barriers and tensor memory pointers
|
||||
const int& smem_barrier = ArchSpec::get_barrier_smem_size(num_stages);
|
||||
const int& smem_tmem_ptr = ArchSpec::get_tmem_ptr_smem_size();
|
||||
const int& smem_tensor_map = ArchSpec::get_tensormap_smem_size(gemm_type);
|
||||
|
||||
// Sum them up
|
||||
int smem_size = 0;
|
||||
smem_size += smem_tensor_map;
|
||||
smem_size += smem_cd;
|
||||
smem_size += num_stages * smem_a_per_stage;
|
||||
smem_size += num_stages * smem_b_per_stage;
|
||||
smem_size += num_stages * smem_sfa_per_stage;
|
||||
smem_size += num_stages * smem_sfb_per_stage;
|
||||
smem_size += smem_extra_sfb;
|
||||
smem_size += smem_barrier;
|
||||
smem_size += smem_tmem_ptr;
|
||||
|
||||
return SharedMemoryConfig {
|
||||
.smem_size = smem_size,
|
||||
.swizzle_a_mode = swizzle_a_mode,
|
||||
.swizzle_b_mode = swizzle_b_mode,
|
||||
.swizzle_cd_mode = swizzle_cd_mode,
|
||||
};
|
||||
}
|
||||
|
||||
template <typename ArchSpec>
|
||||
static GemmConfig get_best_config(const GemmType& gemm_type, const KernelType& kernel_type,
|
||||
const int& m, const int& n, const int& k, const int& num_groups,
|
||||
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
|
||||
const at::ScalarType& a_dtype, const at::ScalarType& b_dtype,
|
||||
const at::ScalarType& cd_dtype,
|
||||
const bool& with_accumulation, const int& num_sms) {
|
||||
const auto mma_kind = (a_dtype == torch::kBFloat16 ? MmaKind::BF16 : MmaKind::MXFP8FP4);
|
||||
if (mma_kind == MmaKind::BF16) {
|
||||
DG_HOST_ASSERT(a_dtype == torch::kBFloat16 and b_dtype == torch::kBFloat16);
|
||||
} else {
|
||||
DG_HOST_ASSERT(a_dtype == torch::kFloat8_e4m3fn or a_dtype == kPackedFP4);
|
||||
DG_HOST_ASSERT(b_dtype == torch::kFloat8_e4m3fn or b_dtype == kPackedFP4);
|
||||
}
|
||||
DG_HOST_ASSERT(cd_dtype == torch::kBFloat16 or cd_dtype == torch::kFloat);
|
||||
|
||||
// Select M/N block sizes
|
||||
auto block_ms = ArchSpec::get_block_m_candidates(kernel_type, major_a, m);
|
||||
if (gemm_type == GemmType::MGroupedContiguous)
|
||||
block_ms = std::vector{get_mk_alignment_for_contiguous_layout()};
|
||||
if (gemm_type == GemmType::MGroupedMasked or gemm_type == GemmType::MGroupedContiguousWithPsumLayout)
|
||||
block_ms = std::vector{64, 128}; // Exclude 256 for performance
|
||||
auto block_ns = ArchSpec::get_block_n_candidates(kernel_type, cd_dtype);
|
||||
|
||||
// NOTES: TMA copy .b4x16_p64 only supports Swizzle 128B
|
||||
// TODO: Optimize it
|
||||
if (a_dtype == kPackedFP4 and major_a == cute::UMMA::Major::MN)
|
||||
block_ms = std::vector{128};
|
||||
if (b_dtype == kPackedFP4 and major_b == cute::UMMA::Major::MN)
|
||||
block_ns = std::vector{128};
|
||||
|
||||
// K block size is selected in a fixed manner
|
||||
const auto& block_k = (mma_kind == MmaKind::BF16 ? 64 : 128);
|
||||
|
||||
// Some util functions
|
||||
const auto& get_num_blocks = [=](const int& block_m, const int& block_n) {
|
||||
return ceil_div(m, block_m) * ceil_div(n, block_n) * num_groups;
|
||||
};
|
||||
const auto& get_num_waves = [=](const int& block_m, const int& block_n) {
|
||||
return ceil_div(get_num_blocks(block_m, block_n), num_sms);
|
||||
};
|
||||
const auto& get_last_wave_util = [=](const int& block_m, const int& block_n) {
|
||||
const auto& num_last_blocks = get_num_blocks(block_m, block_n) % num_sms;
|
||||
return num_last_blocks == 0 ? num_sms : num_last_blocks;
|
||||
};
|
||||
|
||||
// Decide block sizes by waves
|
||||
int best_block_m = 0, best_block_n = 0;
|
||||
int best_num_waves = 0, best_last_util = 0;
|
||||
for (const auto& block_m: block_ms) {
|
||||
for (const auto& block_n: block_ns) {
|
||||
const int& num_waves = get_num_waves(block_m, block_n);
|
||||
const auto& last_util = get_last_wave_util(block_m, block_n);
|
||||
if (not ArchSpec::is_block_size_legal(kernel_type, major_a, major_b, mma_kind, cd_dtype, m, n, k, block_m, block_n, block_k))
|
||||
continue;
|
||||
|
||||
bool success = false;
|
||||
if (best_block_m == 0 or best_block_n == 0 or num_waves < best_num_waves) {
|
||||
success = true;
|
||||
} else if (num_waves == best_num_waves) {
|
||||
// Check last wave utilization
|
||||
success = last_util > best_last_util;
|
||||
if (last_util == best_last_util) {
|
||||
// Case 1: same `block_m`, smaller `block_n` (wasted)
|
||||
success |= block_m == best_block_m and block_n < best_block_n;
|
||||
// Case 2: same `block_n`, smaller `block_m` (wasted)
|
||||
success |= block_n == best_block_n and block_m < best_block_m;
|
||||
// Case 3: different for both `block_m` and `block_n`, larger `block_n` is better
|
||||
// NOTES: don't pick `block_m/block_n` larger than shape `m/n` in this case
|
||||
success |= block_m != best_block_m and block_n > best_block_n
|
||||
and block_n <= n and block_m <= m;
|
||||
}
|
||||
}
|
||||
|
||||
// Replace with the new config if successful
|
||||
if (success) {
|
||||
best_block_m = block_m, best_block_n = block_n;
|
||||
best_num_waves = num_waves, best_last_util = last_util;
|
||||
}
|
||||
}
|
||||
}
|
||||
DG_HOST_ASSERT(best_block_m > 0 and best_block_n > 0);
|
||||
|
||||
// Decide the number of TMA multicasts and whether broadcast on A
|
||||
MulticastConfig best_multicast_config = {1, false};
|
||||
auto [is_legal_on_a, is_legal_on_b] = ArchSpec::get_multicast_legality(
|
||||
gemm_type, num_groups, m, n, best_block_m, best_block_n, num_sms);
|
||||
|
||||
// NOTES: TMA copy .b4x16_p64 only supports Swizzle 128B
|
||||
// TODO: Optimize it
|
||||
if (a_dtype == kPackedFP4 and major_a == cute::UMMA::Major::MN)
|
||||
is_legal_on_a = false;
|
||||
if (b_dtype == kPackedFP4 and major_b == cute::UMMA::Major::MN)
|
||||
is_legal_on_b = false;
|
||||
|
||||
const bool is_legal[2] = {is_legal_on_b, is_legal_on_a};
|
||||
bool order[2] = {false, true};
|
||||
if (best_block_m > best_block_n)
|
||||
std::swap(order[0], order[1]);
|
||||
for (const bool& is_multicast_on_a: order) {
|
||||
if (m >= 512 and is_legal[static_cast<int>(is_multicast_on_a)]) {
|
||||
best_multicast_config = {2, is_multicast_on_a};
|
||||
break;
|
||||
}
|
||||
// Choose the best layout
|
||||
const auto layout_candidates = ArchSpec::get_layout_candidates(desc);
|
||||
DG_HOST_ASSERT(not layout_candidates.empty());
|
||||
auto layout = layout_candidates[0];
|
||||
auto layout_info = ArchSpec::get_layout_info(desc, layout);
|
||||
for (int i = 1; i < static_cast<int>(layout_candidates.size()); ++ i) {
|
||||
const auto candidate_info = ArchSpec::get_layout_info(desc, layout_candidates[i]);
|
||||
if (ArchSpec::compare(candidate_info, layout_info))
|
||||
layout = layout_candidates[i], layout_info = candidate_info;
|
||||
}
|
||||
|
||||
// Always pick the largest number of stage
|
||||
constexpr int smem_capacity = ArchSpec::smem_capacity;
|
||||
int best_num_stages = 0;
|
||||
SharedMemoryConfig best_smem_config;
|
||||
for (int num_stages = 32; num_stages > 0; -- num_stages) {
|
||||
if (not ArchSpec::is_num_stages_legal(mma_kind, cd_dtype, num_stages, best_block_m, best_block_n, block_k))
|
||||
continue;
|
||||
|
||||
best_smem_config = get_smem_config<ArchSpec>(gemm_type, kernel_type,
|
||||
m, n, k,
|
||||
best_block_m, best_block_n, block_k,
|
||||
major_a, major_b,
|
||||
mma_kind, cd_dtype,
|
||||
num_stages, best_multicast_config);
|
||||
if (best_smem_config.smem_size <= smem_capacity) {
|
||||
best_num_stages = num_stages;
|
||||
break;
|
||||
}
|
||||
}
|
||||
DG_HOST_ASSERT(best_num_stages != 0);
|
||||
|
||||
// Recompute the minimal number of SMs required
|
||||
// NOTES: less L2 cache usage and less GPU frequency drop
|
||||
int num_min_sms = num_sms;
|
||||
if (get_env<int>("DG_JIT_MINIMIZE_NUM_SMS", 0)) {
|
||||
num_min_sms = ceil_div(ceil_div(m, best_block_m) * ceil_div(n, best_block_n) * num_groups, best_num_waves);
|
||||
num_min_sms = align(num_min_sms, best_multicast_config.num_multicast);
|
||||
DG_HOST_ASSERT(num_min_sms <= num_sms);
|
||||
}
|
||||
|
||||
const auto& config = GemmConfig {
|
||||
.gemm_type = gemm_type,
|
||||
.kernel_type = kernel_type,
|
||||
.mma_kind = mma_kind,
|
||||
.a_dtype = a_dtype,
|
||||
.b_dtype = b_dtype,
|
||||
.cd_dtype = cd_dtype,
|
||||
.major_a = major_a,
|
||||
.major_b = major_b,
|
||||
.with_accumulation = with_accumulation,
|
||||
.block_m = best_block_m,
|
||||
.block_n = best_block_n,
|
||||
.block_k = block_k,
|
||||
.num_stages = best_num_stages,
|
||||
.num_last_stages = ceil_div(k, block_k) % best_num_stages,
|
||||
.num_sms = num_min_sms,
|
||||
.tc_util = device_runtime->get_tc_util(),
|
||||
.multicast_config = best_multicast_config,
|
||||
// ReSharper disable once CppLocalVariableMightNotBeInitialized
|
||||
.smem_config = best_smem_config,
|
||||
.thread_config = ArchSpec::get_thread_config(kernel_type, best_block_m, best_block_n)
|
||||
// Infer other configs
|
||||
const auto storage_config = ArchSpec::get_storage_config(desc, layout);
|
||||
const auto pipeline_config = ArchSpec::get_pipeline_config(desc, layout, storage_config);
|
||||
const auto launch_config = ArchSpec::get_launch_config(desc, layout);
|
||||
const auto gemm_config = GemmConfig {
|
||||
.layout = layout,
|
||||
.storage_config = storage_config,
|
||||
.pipeline_config = pipeline_config,
|
||||
.launch_config = launch_config
|
||||
};
|
||||
|
||||
// Only SM100 BF16 kernels support tensor core control
|
||||
if (config.tc_util < 100)
|
||||
DG_HOST_ASSERT(device_runtime->get_arch_major() == 10 and mma_kind == MmaKind::BF16);
|
||||
|
||||
// Print configs for the first time
|
||||
if (get_env<int>("DG_JIT_DEBUG") or get_env<int>("DG_PRINT_CONFIGS")) {
|
||||
auto key = std::make_tuple(gemm_type, kernel_type, m, n, k, num_groups, major_a, major_b,
|
||||
mma_kind, a_dtype, b_dtype, cd_dtype, with_accumulation, num_sms);
|
||||
static std::set<decltype(key)> printed;
|
||||
std::stringstream ss;
|
||||
ss << desc;
|
||||
const auto key = ss.str();
|
||||
|
||||
static std::unordered_set<std::string> printed;
|
||||
if (printed.count(key) == 0) {
|
||||
printf("GEMM type: %d, kernel type: %d, M: %d, N: %d, K: %d, groups: %d, "
|
||||
"A major: %d, B major: %d, MMA kind: %d, A dtype: %s, B dtype: %s, CD dtype: %s, accumulation: %d, "
|
||||
"SM limit: %d -> block M: %d, block N: %d, block K: %d, stages: %d, last stages: %d, "
|
||||
"SMs: %d, multicast: %d, multicast on A: %d, shared memory: %d bytes, swizzle A: %d, "
|
||||
"swizzle B: %d, swizzle CD: %d, SMs: %d, threads: %d, TC util: %d%%\n",
|
||||
static_cast<int>(gemm_type), static_cast<int>(kernel_type), m, n, k, num_groups,
|
||||
static_cast<int>(major_a), static_cast<int>(major_b), static_cast<int>(mma_kind),
|
||||
c10::toString(a_dtype), c10::toString(b_dtype), c10::toString(cd_dtype),
|
||||
static_cast<int>(with_accumulation), num_sms, best_block_m, best_block_n, block_k,
|
||||
best_num_stages, config.num_last_stages, num_min_sms, best_multicast_config.num_multicast,
|
||||
static_cast<int>(best_multicast_config.is_multicast_on_a),
|
||||
best_smem_config.smem_size, best_smem_config.swizzle_a_mode, best_smem_config.swizzle_b_mode,
|
||||
best_smem_config.swizzle_cd_mode, config.num_sms, config.thread_config.num_threads, config.tc_util);
|
||||
std::cout << desc << ": " << gemm_config << ", " << layout_info << std::endl;
|
||||
printed.insert(key);
|
||||
}
|
||||
}
|
||||
return config;
|
||||
return gemm_config;
|
||||
}
|
||||
|
||||
} // namespace deep_gemm
|
||||
|
||||
@@ -0,0 +1,171 @@
|
||||
#pragma once
|
||||
|
||||
#include <cute/arch/mma_sm100_desc.hpp>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <deep_gemm/common/types.cuh>
|
||||
|
||||
#include "../../utils/math.hpp"
|
||||
|
||||
namespace deep_gemm {
|
||||
|
||||
/// GEMM descriptors
|
||||
struct GemmDesc {
|
||||
GemmType gemm_type;
|
||||
KernelType kernel_type;
|
||||
int m, n, k, num_groups;
|
||||
at::ScalarType a_dtype, b_dtype, cd_dtype;
|
||||
cute::UMMA::Major major_a;
|
||||
cute::UMMA::Major major_b;
|
||||
bool with_accumulation;
|
||||
|
||||
// Requirements from users
|
||||
int num_sms, tc_util;
|
||||
std::string compiled_dims;
|
||||
|
||||
// Shape for heuristic generation
|
||||
int expected_m = 0, expected_n = 0, expected_k = 0, expected_num_groups = 0;
|
||||
int get_expected_m() const { return expected_m > 0 ? expected_m : m; }
|
||||
int get_expected_n() const { return expected_n > 0 ? expected_n : n; }
|
||||
int get_expected_k() const { return expected_k > 0 ? expected_k : k; }
|
||||
int get_expected_num_groups() const { return expected_num_groups > 0 ? expected_num_groups : num_groups; }
|
||||
|
||||
MmaKind get_mma_kind() const {
|
||||
return a_dtype == torch::kBFloat16 ? MmaKind::BF16 : MmaKind::MXFP8FP4;
|
||||
}
|
||||
|
||||
void check_validity() const {
|
||||
if (get_mma_kind() == MmaKind::BF16) {
|
||||
DG_HOST_ASSERT(a_dtype == torch::kBFloat16 and b_dtype == torch::kBFloat16);
|
||||
} else {
|
||||
DG_HOST_ASSERT(a_dtype == torch::kFloat8_e4m3fn or a_dtype == kPackedFP4);
|
||||
DG_HOST_ASSERT(b_dtype == torch::kFloat8_e4m3fn or b_dtype == kPackedFP4);
|
||||
}
|
||||
DG_HOST_ASSERT(cd_dtype == torch::kBFloat16 or cd_dtype == torch::kFloat);
|
||||
DG_HOST_ASSERT(num_sms % 2 == 0);
|
||||
}
|
||||
|
||||
friend std::ostream& operator << (std::ostream& os, const GemmDesc& desc) {
|
||||
MmaKind mma_kind = desc.get_mma_kind();
|
||||
os << "GemmDesc(gemm_type=" << static_cast<int>(desc.gemm_type)
|
||||
<< ", kernel_type=" << static_cast<int>(desc.kernel_type)
|
||||
<< ", m=" << desc.m << ", n=" << desc.n << ", k=" << desc.k
|
||||
<< ", num_groups=" << desc.num_groups
|
||||
<< ", major_a=" << static_cast<int>(desc.major_a)
|
||||
<< ", major_b=" << static_cast<int>(desc.major_b)
|
||||
<< ", mma_kind=" << static_cast<int>(mma_kind)
|
||||
<< ", a_dtype=" << c10::toString(desc.a_dtype)
|
||||
<< ", b_dtype=" << c10::toString(desc.b_dtype)
|
||||
<< ", cd_dtype=" << c10::toString(desc.cd_dtype)
|
||||
<< ", with_accumulation=" << static_cast<int>(desc.with_accumulation)
|
||||
<< ", num_sms=" << desc.num_sms
|
||||
<< ", tc_util=" << desc.tc_util
|
||||
<< ", compiled_dims=" << desc.compiled_dims
|
||||
<< ", expected_m=" << desc.expected_m
|
||||
<< ", expected_n=" << desc.expected_n
|
||||
<< ", expected_k=" << desc.expected_k
|
||||
<< ", expected_num_groups=" << desc.expected_num_groups << ")";
|
||||
return os;
|
||||
}
|
||||
};
|
||||
|
||||
/// GEMM configs
|
||||
struct Layout {
|
||||
int swap_ab;
|
||||
int block_m, block_n, block_k;
|
||||
int cluster_m, cluster_n;
|
||||
|
||||
int get_cluster_size() const {
|
||||
return cluster_m * cluster_n;
|
||||
}
|
||||
|
||||
friend std::ostream& operator << (std::ostream& os, const Layout& layout) {
|
||||
os << "Layout(swap_ab=" << layout.swap_ab
|
||||
<< ", block_m=" << layout.block_m << ", block_n=" << layout.block_n << ", block_k=" << layout.block_k
|
||||
<< ", cluster_m=" << layout.cluster_m << ", cluster_n=" << layout.cluster_n << ")";
|
||||
return os;
|
||||
}
|
||||
};
|
||||
|
||||
struct StorageConfig {
|
||||
int load_block_m, load_block_n;
|
||||
int store_block_m, store_block_n;
|
||||
|
||||
int swizzle_a_mode, swizzle_b_mode;
|
||||
int swizzle_cd_mode;
|
||||
|
||||
friend std::ostream& operator << (std::ostream& os, const StorageConfig& config) {
|
||||
os << "StorageConfig("
|
||||
<< "load_block_m=" << config.load_block_m << ", load_block_n=" << config.load_block_n
|
||||
<< ", store_block_m=" << config.store_block_m << ", store_block_n=" << config.store_block_n
|
||||
<< ", swizzle_a_mode=" << config.swizzle_a_mode << ", swizzle_b_mode=" << config.swizzle_b_mode
|
||||
<< ", swizzle_cd_mode=" << config.swizzle_cd_mode << ")";
|
||||
return os;
|
||||
}
|
||||
};
|
||||
|
||||
struct PipelineConfig {
|
||||
int smem_size;
|
||||
int num_stages;
|
||||
|
||||
friend std::ostream& operator << (std::ostream& os, const PipelineConfig& config) {
|
||||
os << "PipelineConfig("
|
||||
<< "smem_size=" << config.smem_size
|
||||
<< ", num_stages=" << config.num_stages << ")";
|
||||
return os;
|
||||
}
|
||||
};
|
||||
|
||||
struct LaunchConfig {
|
||||
int num_sms;
|
||||
int num_sms_per_cluster;
|
||||
int num_threads;
|
||||
|
||||
int num_tma_threads;
|
||||
int num_math_threads;
|
||||
int num_non_epilogue_threads;
|
||||
int num_epilogue_threads;
|
||||
|
||||
friend std::ostream& operator << (std::ostream& os, const LaunchConfig& config) {
|
||||
os << "LaunchConfig("
|
||||
<< "num_sms=" << config.num_sms << ", num_sms_per_cluster=" << config.num_sms_per_cluster
|
||||
<< ", num_threads=" << config.num_threads
|
||||
<< ", num_tma_threads=" << config.num_tma_threads << ", num_math_threads=" << config.num_math_threads
|
||||
<< ", num_non_epilogue_threads=" << config.num_non_epilogue_threads
|
||||
<< ", num_epilogue_threads=" << config.num_epilogue_threads << ")";
|
||||
return os;
|
||||
}
|
||||
};
|
||||
|
||||
struct GemmConfig {
|
||||
Layout layout;
|
||||
StorageConfig storage_config;
|
||||
PipelineConfig pipeline_config;
|
||||
LaunchConfig launch_config;
|
||||
|
||||
friend std::ostream& operator << (std::ostream& os, const GemmConfig& config) {
|
||||
os << "GemmConfig("
|
||||
<< "layout=" << config.layout
|
||||
<< ", storage_config=" << config.storage_config
|
||||
<< ", pipeline_config=" << config.pipeline_config
|
||||
<< ", launch_config=" << config.launch_config << ")";
|
||||
return os;
|
||||
}
|
||||
};
|
||||
|
||||
/// Config comparators
|
||||
struct LayoutInfo {
|
||||
int num_waves;
|
||||
int last_wave_util;
|
||||
int64_t num_cycles;
|
||||
Layout layout;
|
||||
|
||||
friend std::ostream& operator << (std::ostream& os, const LayoutInfo& config) {
|
||||
os << "LayoutInfo("
|
||||
<< "num_waves=" << config.num_waves
|
||||
<< ", last_wave_util=" << config.last_wave_util
|
||||
<< ", num_cycles=" << config.num_cycles << ")";
|
||||
return os;
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace deep_gemm
|
||||
@@ -0,0 +1,211 @@
|
||||
#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"
|
||||
#include "sm100.hpp"
|
||||
|
||||
namespace deep_gemm {
|
||||
|
||||
struct MegaMoEConfig {
|
||||
// Block tiling
|
||||
int block_m, block_n, block_k;
|
||||
int load_block_m, load_block_n;
|
||||
int store_block_m;
|
||||
|
||||
// SF block sizes (UTCCP 128-aligned)
|
||||
int sf_block_m, sf_block_n;
|
||||
|
||||
// Pool capacity and SF-padded token count
|
||||
int num_max_pool_tokens;
|
||||
int num_padded_sf_pool_tokens;
|
||||
|
||||
// Swizzle modes for TMA descriptors
|
||||
int swizzle_acts_mode, swizzle_weights_mode;
|
||||
|
||||
// Number of experts to process per wave
|
||||
int num_experts_per_wave;
|
||||
|
||||
// Pipeline stages and shared memory
|
||||
int num_stages, smem_size;
|
||||
|
||||
// Thread layout
|
||||
int num_dispatch_threads, num_non_epilogue_threads, num_epilogue_threads;
|
||||
|
||||
friend std::ostream& operator << (std::ostream& os, const MegaMoEConfig& config) {
|
||||
os << "MegaMoEConfig("
|
||||
<< "block_m=" << config.block_m << ", block_n=" << config.block_n << ", block_k=" << config.block_k
|
||||
<< ", load_block_m=" << config.load_block_m << ", load_block_n=" << config.load_block_n
|
||||
<< ", store_block_m=" << config.store_block_m
|
||||
<< ", sf_block_m=" << config.sf_block_m << ", sf_block_n=" << config.sf_block_n
|
||||
<< ", num_max_pool_tokens=" << config.num_max_pool_tokens
|
||||
<< ", num_padded_sf_pool_tokens=" << config.num_padded_sf_pool_tokens
|
||||
<< ", swizzle_acts_mode=" << config.swizzle_acts_mode << ", swizzle_weights_mode=" << config.swizzle_weights_mode
|
||||
<< ", 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_non_epilogue_threads=" << config.num_non_epilogue_threads
|
||||
<< ", num_epilogue_threads=" << config.num_epilogue_threads << ")";
|
||||
return os;
|
||||
}
|
||||
};
|
||||
|
||||
static int get_block_m_for_mega_moe(const int& num_ranks, const int& num_experts,
|
||||
const int& num_max_tokens_per_rank, const int& num_topk) {
|
||||
// TODO: compute based on configs
|
||||
return 192;
|
||||
}
|
||||
|
||||
static int get_num_experts_per_wave_for_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) {
|
||||
// Reduce per-expert block count by this factor since uneven routing leaves some experts with fewer tokens
|
||||
constexpr int kImbalanceFactor = 2;
|
||||
|
||||
// TODO: support num_experts_per_rank > 32
|
||||
// Find the largest divisor of num_experts_per_rank that fits in 32 as the upper bound
|
||||
int max_num_experts_per_wave = std::min(32, num_experts_per_rank);
|
||||
while (max_num_experts_per_wave > 1 and num_experts_per_rank % max_num_experts_per_wave != 0)
|
||||
-- max_num_experts_per_wave;
|
||||
|
||||
// Count L1 blocks per expert assuming tokens are evenly spread across experts
|
||||
const int expected_tokens_per_expert =
|
||||
num_tokens * num_topk / num_experts_per_rank + 1;
|
||||
const int num_m_blocks = ceil_div(expected_tokens_per_expert, block_m);
|
||||
const int num_n_blocks = intermediate_hidden / block_n;
|
||||
const int num_l1_blocks_per_expert = num_m_blocks * num_n_blocks;
|
||||
|
||||
// Pick the smallest value whose total blocks (after imbalance reduction) can keep all SMs busy
|
||||
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, max_num_experts_per_wave);
|
||||
|
||||
// Round up to the nearest divisor of num_experts_per_rank so every wave processes the same count
|
||||
while (num_experts_per_wave < max_num_experts_per_wave 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_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& sf_block_m, const int& sf_block_n,
|
||||
const int& num_dispatch_warps, const int& num_epilogue_warps) {
|
||||
constexpr int kSmemAlignment = 1024;
|
||||
constexpr int kNumEpilogueStages = 2;
|
||||
constexpr int kNumTMAStoreStages = 2;
|
||||
|
||||
// Always multicast on A
|
||||
const int load_block_m = block_m / 2;
|
||||
|
||||
// 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, BLOCK_N/2 post-SwiGLU) and L2 BF16 (1 stage)
|
||||
const auto num_epilogue_warpgroups = num_epilogue_warps / 4;
|
||||
const int smem_cd_l1 = num_epilogue_warpgroups * store_block_m * (block_n / 2) * kNumTMAStoreStages;
|
||||
const int smem_cd_l2 = num_epilogue_warpgroups * store_block_m * block_n * static_cast<int>(sizeof(nv_bfloat16));
|
||||
const int smem_cd = std::max(smem_cd_l1, smem_cd_l2);
|
||||
|
||||
// Barriers (stage-independent): dispatch + tensor memory full/empty + combine (2 per epilogue warp)
|
||||
const int smem_barriers = (num_dispatch_warps + kNumEpilogueStages * 2 + num_epilogue_warps * 2) * 8;
|
||||
|
||||
// Amax reduction
|
||||
const int smem_amax_reduction = store_block_m * num_epilogue_warps * static_cast<int>(sizeof(float));
|
||||
|
||||
// Tensor memory pointer
|
||||
const int smem_tmem_ptr = 4;
|
||||
|
||||
// SF is aligned to UTCCP 128-element granularity
|
||||
const int smem_sfa_per_stage = sf_block_m * 4;
|
||||
const int smem_sfb_per_stage = sf_block_n * 4;
|
||||
|
||||
// Per-stage: A tile + B tile + SFA tile + SFB tile + full/empty barriers
|
||||
const int smem_per_stage = load_block_m * block_k + block_n * block_k + smem_sfa_per_stage + smem_sfb_per_stage + 2 * 8;
|
||||
|
||||
// Fixed total
|
||||
const int smem_fixed = smem_dispatch_size + smem_cd + smem_amax_reduction + smem_barriers + smem_tmem_ptr;
|
||||
|
||||
// Select maximum num_stages
|
||||
const int num_stages = (smem_capacity - smem_fixed) / smem_per_stage;
|
||||
DG_HOST_ASSERT(num_stages >= 2);
|
||||
|
||||
return {num_stages, smem_fixed + num_stages * smem_per_stage};
|
||||
}
|
||||
|
||||
static MegaMoEConfig get_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) {
|
||||
// Block tiling
|
||||
const int block_m = get_block_m_for_mega_moe(num_ranks, num_experts, num_max_tokens_per_rank, num_topk);
|
||||
const int block_n = 128;
|
||||
const int block_k = 128;
|
||||
const int load_block_m = block_m / 2;
|
||||
const int load_block_n = block_n;
|
||||
const int store_block_m = 32;
|
||||
const auto [sf_block_m, sf_block_n] = SM100ArchSpec::get_sf_uttcp_aligned_block_sizes(block_m, block_n, MmaKind::MXFP8FP4);
|
||||
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, block_m);
|
||||
const int num_padded_sf_pool_tokens = layout::get_num_padded_sf_pool_tokens(num_max_pool_tokens, block_m);
|
||||
// NOTES: FP8 activations and FP4 weights (unpacked to 8-bit in smem) both use 128B swizzle
|
||||
const int swizzle_acts_mode = 128;
|
||||
const int swizzle_weights_mode = 128;
|
||||
|
||||
// Waves
|
||||
const int num_sms = device_runtime->get_num_sms();
|
||||
const int num_experts_per_wave = get_num_experts_per_wave_for_mega_moe(
|
||||
num_experts_per_rank, num_tokens, num_topk,
|
||||
intermediate_hidden, block_m, block_n, num_sms);
|
||||
|
||||
// Thread layout
|
||||
const int num_dispatch_threads = 128;
|
||||
const int num_non_epilogue_threads = 128;
|
||||
const int num_epilogue_threads = 256;
|
||||
|
||||
// Pipeline
|
||||
const auto [num_stages, smem_size] = get_pipeline_config_for_mega_moe(
|
||||
SM100ArchSpec::smem_capacity,
|
||||
num_experts, hidden,
|
||||
block_m, block_n, block_k, store_block_m,
|
||||
sf_block_m, sf_block_n,
|
||||
num_dispatch_threads / 32, num_epilogue_threads / 32);
|
||||
|
||||
const auto config = MegaMoEConfig {
|
||||
block_m, block_n, block_k,
|
||||
load_block_m, load_block_n, store_block_m,
|
||||
sf_block_m, sf_block_n,
|
||||
num_max_pool_tokens, num_padded_sf_pool_tokens,
|
||||
swizzle_acts_mode, swizzle_weights_mode,
|
||||
num_experts_per_wave,
|
||||
num_stages, smem_size,
|
||||
num_dispatch_threads, num_non_epilogue_threads, num_epilogue_threads
|
||||
};
|
||||
|
||||
// Print configs for the first time
|
||||
if (get_env<int>("DG_JIT_DEBUG") or get_env<int>("DG_PRINT_CONFIGS")) {
|
||||
const auto key = fmt::format(
|
||||
"MegaMoEConfig(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,62 @@
|
||||
#pragma once
|
||||
|
||||
#include "../../jit/device_runtime.hpp"
|
||||
#include "../../utils/exception.hpp"
|
||||
#include "../../utils/lazy_init.hpp"
|
||||
|
||||
namespace deep_gemm {
|
||||
|
||||
class HeuristicsRuntime {
|
||||
static constexpr int kLegacyMKAlignmentForContiguousLayout = 128;
|
||||
|
||||
bool ignore_compile_dims = false;
|
||||
int block_m_multiple_of = 1;
|
||||
int block_n_multiple_of = 1;
|
||||
int mk_alignment_for_contiguous_layout = kLegacyMKAlignmentForContiguousLayout;
|
||||
|
||||
public:
|
||||
void set_ignore_compile_dims(const bool& new_value) {
|
||||
ignore_compile_dims = new_value;
|
||||
}
|
||||
|
||||
bool get_ignore_compile_dims() const {
|
||||
return ignore_compile_dims;
|
||||
}
|
||||
|
||||
void set_block_size_multiple_of(const int& new_block_m_multiple_of, const int& new_block_n_multiple_of) {
|
||||
block_m_multiple_of = new_block_m_multiple_of;
|
||||
block_n_multiple_of = new_block_n_multiple_of;
|
||||
}
|
||||
|
||||
int get_block_m_multiple_of() const {
|
||||
return block_m_multiple_of;
|
||||
}
|
||||
|
||||
int get_block_n_multiple_of() const {
|
||||
return block_n_multiple_of;
|
||||
}
|
||||
|
||||
void set_mk_alignment_for_contiguous_layout(const int& new_value) {
|
||||
mk_alignment_for_contiguous_layout = new_value;
|
||||
}
|
||||
|
||||
int get_mk_alignment_for_contiguous_layout() const {
|
||||
return mk_alignment_for_contiguous_layout;
|
||||
}
|
||||
|
||||
static int get_theoretical_mk_alignment_for_contiguous_layout(const std::optional<int>& expected_m) {
|
||||
if (device_runtime->get_arch_major() != 10)
|
||||
return kLegacyMKAlignmentForContiguousLayout;
|
||||
|
||||
int block_m = 240, mma_step = 16;
|
||||
if (expected_m.has_value()) {
|
||||
// Reduce `block_m` while ensuring it covers `m`
|
||||
for (; block_m > 32 and block_m - mma_step >= expected_m.value(); block_m -= mma_step);
|
||||
}
|
||||
return block_m;
|
||||
}
|
||||
};
|
||||
|
||||
static auto heuristics_runtime = LazyInit<HeuristicsRuntime>([](){ return std::make_shared<HeuristicsRuntime>(); });
|
||||
|
||||
} // namespace deep_gemm
|
||||
@@ -2,9 +2,11 @@
|
||||
|
||||
#include <cute/arch/mma_sm100_desc.hpp>
|
||||
// Reuse some types in the JIT modules
|
||||
#include <deep_gemm/common/types.hpp>
|
||||
#include <deep_gemm/common/types.cuh>
|
||||
|
||||
#include "common.hpp"
|
||||
#include "runtime.hpp"
|
||||
#include "utils.hpp"
|
||||
#include "../../utils/exception.hpp"
|
||||
|
||||
namespace deep_gemm {
|
||||
@@ -12,155 +14,255 @@ namespace deep_gemm {
|
||||
struct SM100ArchSpec {
|
||||
static constexpr int smem_capacity = 232448;
|
||||
|
||||
static std::vector<int> get_block_m_candidates(const KernelType& kernel_type, const cute::UMMA::Major& major_a, const int& m) {
|
||||
std::vector<int> candidates{128, 256};
|
||||
if ((kernel_type == KernelType::Kernel1D1D or kernel_type == KernelType::KernelNoSF) and major_a == cute::UMMA::Major::K) {
|
||||
// NOTES: `block_m = 32/64` is smaller than `LAYOUT_AD_M`, should be careful in handling this
|
||||
if (m <= 32) candidates.push_back(32);
|
||||
if (m <= 64) candidates.push_back(64);
|
||||
}
|
||||
return candidates;
|
||||
}
|
||||
|
||||
static std::vector<int> get_block_n_candidates(const KernelType& kernel_type, const at::ScalarType& cd_dtype) {
|
||||
// 16 is for better SM usage
|
||||
// Stride 32 is due to low-performance swizzle-16/32B
|
||||
std::vector<int> candidates = {16};
|
||||
for (int i = 32; i <= 256; i += 32)
|
||||
candidates.push_back(i);
|
||||
return candidates;
|
||||
}
|
||||
|
||||
static int get_ab_load_block_m(const MulticastConfig& config, const int& block_m) {
|
||||
return block_m / (config.is_multicast_on_a ? config.num_multicast : 1);
|
||||
}
|
||||
|
||||
static int get_ab_load_block_n(const MulticastConfig& config, const int& block_n) {
|
||||
return block_n / (config.is_multicast_on_a ? 1 : config.num_multicast);
|
||||
}
|
||||
|
||||
static int get_cd_store_block_m(const int& block_m) {
|
||||
constexpr int layout_ad_m = 128;
|
||||
return std::min(block_m, layout_ad_m);
|
||||
}
|
||||
|
||||
static int get_cd_store_block_n(const int& block_n) {
|
||||
return block_n;
|
||||
}
|
||||
|
||||
static bool enable_cd_swizzle(const at::ScalarType& cd_dtype) {
|
||||
return true;
|
||||
}
|
||||
|
||||
static std::pair<int, int> get_sf_uttcp_aligned_block_sizes(
|
||||
const int& block_m, const int& block_n, const MmaKind& mma_kind) {
|
||||
constexpr int num_utccp_aligned_elems = 128;
|
||||
switch (mma_kind) {
|
||||
case MmaKind::BF16: return {0, 0};
|
||||
case MmaKind::BF16: return {0, 0};
|
||||
case MmaKind::MXFP8FP4: return {align(block_m, num_utccp_aligned_elems), align(block_n, num_utccp_aligned_elems)};
|
||||
default: DG_HOST_UNREACHABLE("Unknown dtype");
|
||||
}
|
||||
}
|
||||
|
||||
static bool is_block_size_legal(const KernelType& kernel_type,
|
||||
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
|
||||
const MmaKind& mma_kind, const at::ScalarType& cd_dtype,
|
||||
const int& m, const int& n, const int& k,
|
||||
const int& block_m, const int& block_n, const int& block_k) {
|
||||
// Layout A/D does not support `block_n % 16 != 0`
|
||||
if (block_n % 16 != 0)
|
||||
return false;
|
||||
static std::vector<Layout> get_layout_candidates(const GemmDesc& desc) {
|
||||
// Block K is always in a fixed manner
|
||||
const int block_k = 128 / get_element_size(desc.get_mma_kind());
|
||||
|
||||
// Performance is lower with 1D1D and `block_m == 256`
|
||||
if (kernel_type == KernelType::Kernel1D1D and major_b == cute::UMMA::Major::K and block_m > 128)
|
||||
return false;
|
||||
|
||||
// For small K, fewer store blocks improve store/compute overlap and reduce epilogue bottleneck
|
||||
if (k <= 256 and (block_n > 128 or block_m > 128))
|
||||
return false;
|
||||
|
||||
// Check tensor memory validity
|
||||
int sf_block_m = 0, sf_block_n = 0;
|
||||
if (kernel_type == KernelType::Kernel1D1D) {
|
||||
const auto& [sf_block_m_, sf_block_n_] = get_sf_uttcp_aligned_block_sizes(block_m, block_n, mma_kind);
|
||||
sf_block_m = sf_block_m_, sf_block_n = sf_block_n_;
|
||||
// Always enable swap A/B (and multicasting if possible) for m-grouped GEMMs
|
||||
if (desc.gemm_type == GemmType::MGroupedContiguous or
|
||||
desc.gemm_type == GemmType::MGroupedContiguousWithPsumLayout or
|
||||
desc.gemm_type == GemmType::MGroupedMasked) {
|
||||
const bool swap_ab = true;
|
||||
const auto block_n = 128;
|
||||
const auto block_m = heuristics_runtime->get_mk_alignment_for_contiguous_layout();
|
||||
const auto cluster_m = 1;
|
||||
const auto cluster_n = ceil_div(desc.n, block_n) % 2 == 0 and desc.num_sms % 2 == 0 ? 2 : 1;
|
||||
const auto layout = Layout{swap_ab, block_m, block_n, block_k, cluster_m, cluster_n};
|
||||
std::vector<Layout> candidates = {layout};
|
||||
return candidates;
|
||||
}
|
||||
if (((2 * block_n) + (sf_block_m / 32) + (sf_block_n / 32)) > 512)
|
||||
return false;
|
||||
|
||||
// NOTES: when B is MN-major, we restrict `block_n` to multiples of 64,
|
||||
// since TMA performance degrades when `swizzle_b <= 32B` (i.e., when `block_ns % 64 != 0`), even with 3D TMA
|
||||
return major_b == cute::UMMA::Major::K or (block_n * get_element_size(mma_kind)) % 64 == 0;
|
||||
// Enumerate all candidates
|
||||
std::vector<Layout> candidates;
|
||||
for (int swap_ab = 0; swap_ab < 2; ++ swap_ab) {
|
||||
// Block M/N candidates
|
||||
std::vector<int> block_m_candidates;
|
||||
std::vector<int> block_n_candidates;
|
||||
if (swap_ab) {
|
||||
int step = std::lcm(16, heuristics_runtime->get_block_m_multiple_of());
|
||||
int end = 256;
|
||||
for (int i = step; i <= end; i += step)
|
||||
block_m_candidates.push_back(i);
|
||||
|
||||
// TODO: consider other block N
|
||||
block_n_candidates = {128};
|
||||
} else {
|
||||
// NOTES: smaller block M can avoid TMA L2 OOB bound
|
||||
// TODO: consider block M = 256
|
||||
if (desc.m <= 32) block_m_candidates = {32};
|
||||
else if (desc.m <= 64) block_m_candidates = {64};
|
||||
else block_m_candidates = {128};
|
||||
|
||||
// Small block size for small shape
|
||||
if (16 % heuristics_runtime->get_block_n_multiple_of() == 0)
|
||||
block_n_candidates.push_back(16);
|
||||
int step = std::lcm(32, heuristics_runtime->get_block_n_multiple_of());
|
||||
// For small K, fewer store blocks improve store/compute overlap and reduce epilogue bottleneck
|
||||
int end = desc.k <= 256 ? 128 : 256;
|
||||
for (int i = step; i <= end; i += step)
|
||||
block_n_candidates.push_back(i);
|
||||
}
|
||||
|
||||
for (int cluster_m = 1; cluster_m <= 2; ++ cluster_m) {
|
||||
// After swapping, layout A/D can only do on cluster N
|
||||
if (swap_ab == 1 and cluster_m > 1)
|
||||
continue;
|
||||
|
||||
for (int cluster_n = 1; cluster_n <= 2; ++ cluster_n) {
|
||||
// We only support cluster 2
|
||||
if (cluster_m * cluster_n > 2)
|
||||
continue;
|
||||
|
||||
// Only support layout A/D
|
||||
if (swap_ab == 0 and cluster_n > 1)
|
||||
continue;
|
||||
|
||||
// SM count must be divisible
|
||||
if (desc.num_sms % (cluster_m * cluster_n) != 0)
|
||||
continue;
|
||||
|
||||
for (int block_m: block_m_candidates) {
|
||||
// Ensure large swizzle sizes (32B swizzle yields poor performance)
|
||||
const auto swizzle_a_requirement = desc.a_dtype == kPackedFP4 ? 128 : 64;
|
||||
// Enforce swizzle alignment for MN major; otherwise check base MMA shape
|
||||
const auto load_block_m_requirement = desc.major_a == cute::UMMA::Major::MN ? swizzle_a_requirement : 8;
|
||||
if ((block_m / cluster_n) % load_block_m_requirement != 0)
|
||||
continue;
|
||||
|
||||
// Shape must be divisible for multicast
|
||||
if (ceil_div(desc.m, block_m) % cluster_m != 0)
|
||||
continue;
|
||||
|
||||
for (int block_n: block_n_candidates) {
|
||||
// Ensure large swizzle sizes (32B swizzle yields poor performance)
|
||||
const auto swizzle_b_requirement = desc.b_dtype == kPackedFP4 ? 128 : 64;
|
||||
// Enforce swizzle alignment for MN major; otherwise check base MMA shape
|
||||
const auto load_block_n_requirement = desc.major_b == cute::UMMA::Major::MN ? swizzle_b_requirement : 8;
|
||||
if ((block_n / cluster_m) % load_block_n_requirement != 0)
|
||||
continue;
|
||||
|
||||
// Shape must be divisible for multicast
|
||||
if (ceil_div(desc.n, block_n) % cluster_n != 0)
|
||||
continue;
|
||||
|
||||
// SwapAB requires block N is layout A/D' UMMA M
|
||||
constexpr int layout_ad_m = 128;
|
||||
if (swap_ab and block_n != layout_ad_m)
|
||||
continue;
|
||||
|
||||
// Check tensor memory capacity
|
||||
const auto [sf_block_m, sf_block_n] = get_sf_uttcp_aligned_block_sizes(block_m, block_n, desc.get_mma_kind());
|
||||
const auto tmem_sf_cols = desc.get_mma_kind() == MmaKind::MXFP8FP4 ? sf_block_m / 32 + sf_block_n / 32 : 0;
|
||||
const auto umma_n = swap_ab ? block_m : block_n;
|
||||
if (2 * umma_n + tmem_sf_cols > 512)
|
||||
continue;
|
||||
|
||||
const auto layout = Layout{swap_ab, block_m, block_n, block_k, cluster_m, cluster_n};
|
||||
|
||||
// When neither A nor B is MN major, 128B swizzle is always feasible
|
||||
if (desc.major_a == cute::UMMA::Major::K or desc.major_b == cute::UMMA::Major::K) {
|
||||
const auto storage_config = get_storage_config(desc, layout);
|
||||
if (storage_config.swizzle_a_mode != 128 or storage_config.swizzle_b_mode != 128)
|
||||
continue;
|
||||
}
|
||||
|
||||
candidates.push_back(layout);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
DG_HOST_ASSERT(not candidates.empty());
|
||||
return candidates;
|
||||
}
|
||||
|
||||
static bool is_num_stages_legal(const MmaKind& mma_kind, const at::ScalarType& cd_dtype,
|
||||
const int& num_stages,
|
||||
const int& block_m, const int& block_n, const int& block_k) {
|
||||
return true;
|
||||
}
|
||||
static StorageConfig get_storage_config(const GemmDesc& desc, const Layout& layout) {
|
||||
constexpr int layout_ad_m = 128;
|
||||
constexpr int umma_step_n = 16;
|
||||
|
||||
// Load/store block sizes (w/o consideration of swizzling atoms, w/ consideration of loop atoms)
|
||||
const auto load_block_m = layout.block_m / layout.cluster_n;
|
||||
const auto load_block_n = layout.block_n / layout.cluster_m;
|
||||
const auto store_block_m = layout.swap_ab ? umma_step_n : std::min(layout_ad_m, layout.block_m);
|
||||
const auto store_block_n = layout.block_n;
|
||||
|
||||
// Decide swizzling by the inner dim
|
||||
// TODO: support FP4 sub-byte
|
||||
const auto swizzle_mode_a = get_swizzle_mode(
|
||||
desc.major_a == cute::UMMA::Major::K ? layout.block_k : load_block_m, c10::elementSize(desc.a_dtype));
|
||||
const auto swizzle_mode_b = get_swizzle_mode(
|
||||
desc.major_b == cute::UMMA::Major::K ? layout.block_k : load_block_n, c10::elementSize(desc.b_dtype));
|
||||
const auto swizzle_mode_cd = get_swizzle_mode(
|
||||
store_block_n, c10::elementSize(desc.cd_dtype));
|
||||
|
||||
static std::pair<bool, bool> get_multicast_legality(const GemmType& gemm_type, const int& num_groups,
|
||||
const int& m, const int& n, const int& block_m, const int& block_n,
|
||||
const int& num_sms) {
|
||||
// TODO: support other layouts
|
||||
return {
|
||||
false,
|
||||
is_multicast_legal(m, block_m, 2, num_sms, true) and (gemm_type == GemmType::Normal or gemm_type == GemmType::KGroupedContiguous
|
||||
or (gemm_type == GemmType::Batched and num_groups <= 32)),
|
||||
load_block_m, load_block_n,
|
||||
store_block_m, store_block_n,
|
||||
swizzle_mode_a, swizzle_mode_b, swizzle_mode_cd
|
||||
};
|
||||
}
|
||||
|
||||
static ThreadConfig get_thread_config(const KernelType& kernel_type,
|
||||
const int& block_m, const int& block_n) {
|
||||
return ThreadConfig::sm100(128, 128);
|
||||
}
|
||||
static PipelineConfig get_pipeline_config(const GemmDesc& desc, const Layout& layout, const StorageConfig& storage_config) {
|
||||
constexpr int kNumMaxStages = 32;
|
||||
|
||||
static int get_smem_cd_size(const KernelType& kernel_type,
|
||||
const int& block_m, const int& block_n,
|
||||
const int& swizzle_cd_mode,
|
||||
const at::ScalarType& cd_dtype) {
|
||||
constexpr static int layout_ad_m = 128;
|
||||
return std::min(block_m, layout_ad_m) * swizzle_cd_mode * 2;
|
||||
}
|
||||
// C/D for TMA stores
|
||||
const int smem_cd = layout.swap_ab ? storage_config.store_block_m * storage_config.store_block_n * c10::elementSize(desc.cd_dtype) * 2
|
||||
: storage_config.store_block_m * storage_config.swizzle_cd_mode * 2;
|
||||
|
||||
static std::pair<int, int> get_sf_smem_size_per_stage(const KernelType& kernel_type,
|
||||
const int& block_m, const int& block_n, const int& block_k,
|
||||
const MmaKind& mma_kind, const at::ScalarType& cd_dtype) {
|
||||
if (mma_kind == MmaKind::BF16)
|
||||
return {0, 0};
|
||||
|
||||
int smem_sfa_per_stage = 0;
|
||||
int smem_sfb_per_stage = 0;
|
||||
if (kernel_type == KernelType::Kernel1D1D) {
|
||||
const auto [sf_block_m, sf_block_n] = get_sf_uttcp_aligned_block_sizes(block_m, block_n, mma_kind);
|
||||
smem_sfa_per_stage = sf_block_m * 4;
|
||||
smem_sfb_per_stage = sf_block_n * 4;
|
||||
} else {
|
||||
smem_sfa_per_stage = block_m * 4;
|
||||
smem_sfb_per_stage = 0;
|
||||
}
|
||||
return {smem_sfa_per_stage, smem_sfb_per_stage};
|
||||
}
|
||||
|
||||
static int get_extra_sfb_smem_size(const int& m, const int& n, const int& k,
|
||||
const int& block_m, const int& block_n, const int& block_k) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
static int get_barrier_smem_size(const int& num_stages) {
|
||||
// TODO: remove SF barriers for BF16 GEMMs
|
||||
// TMA full/empty barriers, with-SF full barriers, tensor memory full/empty barriers
|
||||
// NOTES: some shapes may only have 1 epilogue stage, but we still allocate space for 2 stages
|
||||
// NOTES: the last barrier is for tensor core utilization control
|
||||
return num_stages * 8 * 3 + 2 * 8 * 2 + 8;
|
||||
const int smem_barriers = kNumMaxStages * 8 * 3 + 2 * 8 * 2 + 8;
|
||||
|
||||
// Tensor memory pointer
|
||||
const int smem_tmem_ptr = 4;
|
||||
|
||||
// Calculate A/B per stages
|
||||
// TODO: consider FP4
|
||||
const int smem_a_per_stage = storage_config.load_block_m * layout.block_k * c10::elementSize(desc.a_dtype);
|
||||
const int smem_b_per_stage = storage_config.load_block_n * layout.block_k * c10::elementSize(desc.b_dtype);
|
||||
|
||||
// Calculate SF A/B per stages
|
||||
int smem_sfa_per_stage = 0;
|
||||
int smem_sfb_per_stage = 0;
|
||||
if (desc.kernel_type == KernelType::Kernel1D1D) {
|
||||
const auto [sf_block_m, sf_block_n] = get_sf_uttcp_aligned_block_sizes(
|
||||
layout.block_m, layout.block_n, desc.get_mma_kind());
|
||||
smem_sfa_per_stage = sf_block_m * 4;
|
||||
smem_sfb_per_stage = sf_block_n * 4;
|
||||
}
|
||||
|
||||
// Calculate stages
|
||||
int smem_extra = smem_cd + smem_barriers + smem_tmem_ptr;
|
||||
int smem_per_stage = smem_a_per_stage + smem_b_per_stage + smem_sfa_per_stage + smem_sfb_per_stage;
|
||||
int num_stages = std::min(
|
||||
(smem_capacity - smem_extra) / smem_per_stage,
|
||||
kNumMaxStages);
|
||||
return {
|
||||
smem_extra + num_stages * smem_per_stage,
|
||||
num_stages
|
||||
};
|
||||
}
|
||||
|
||||
static int get_tmem_ptr_smem_size() {
|
||||
return 4;
|
||||
static LaunchConfig get_launch_config(const GemmDesc& desc, const Layout& layout) {
|
||||
return {
|
||||
desc.num_sms,
|
||||
layout.get_cluster_size(),
|
||||
256,
|
||||
32, 128, 128, 128
|
||||
};
|
||||
}
|
||||
|
||||
static int get_tensormap_smem_size(const GemmType& gemm_type) {
|
||||
return 0;
|
||||
static LayoutInfo get_layout_info(const GemmDesc& desc, const Layout& layout) {
|
||||
const auto num_blocks =
|
||||
ceil_div(desc.get_expected_m(), layout.block_m) *
|
||||
ceil_div(desc.get_expected_n(), layout.block_n) *
|
||||
desc.get_expected_num_groups();
|
||||
const auto num_waves = ceil_div(num_blocks, desc.num_sms);
|
||||
const auto num_last_blocks = num_blocks % desc.num_sms;
|
||||
const auto last_wave_util = num_last_blocks == 0 ? desc.num_sms : num_last_blocks;
|
||||
// TODO: calculate expected cycles
|
||||
return {num_waves, last_wave_util, 0, layout};
|
||||
}
|
||||
|
||||
// A regular comparator
|
||||
static bool compare(const LayoutInfo& a, const LayoutInfo& b) {
|
||||
// Single wave is always better
|
||||
if ((a.num_waves == 1 or b.num_waves == 1) and a.num_waves != b.num_waves)
|
||||
return a.num_waves < b.num_waves;
|
||||
|
||||
// Doing multicast is better
|
||||
if (a.layout.get_cluster_size() != b.layout.get_cluster_size())
|
||||
return a.layout.get_cluster_size() > b.layout.get_cluster_size();
|
||||
|
||||
// Smaller number of waves is better
|
||||
if (a.num_waves != b.num_waves)
|
||||
return a.num_waves < b.num_waves;
|
||||
|
||||
// Larger last wave utilization is better
|
||||
if (a.last_wave_util != b.last_wave_util)
|
||||
return a.last_wave_util > b.last_wave_util;
|
||||
|
||||
// More stages is better
|
||||
// Same block M, smaller block N is better
|
||||
// Same block N, smaller block M is better
|
||||
if (a.layout.block_m + a.layout.block_n != b.layout.block_m + b.layout.block_n)
|
||||
return a.layout.block_m + a.layout.block_n < b.layout.block_m + b.layout.block_n;
|
||||
|
||||
// Less shared memory C/D, more stages is better
|
||||
return a.layout.block_m * a.layout.block_n < b.layout.block_m * b.layout.block_n;
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@@ -2,162 +2,244 @@
|
||||
|
||||
#include <cute/arch/mma_sm100_desc.hpp>
|
||||
// Reuse some types in the JIT modules
|
||||
#include <deep_gemm/common/types.hpp>
|
||||
#include <deep_gemm/common/types.cuh>
|
||||
|
||||
#include "common.hpp"
|
||||
#include "utils.hpp"
|
||||
#include "../../utils/exception.hpp"
|
||||
|
||||
namespace deep_gemm {
|
||||
|
||||
struct SM90ArchSpec {
|
||||
static constexpr int smem_capacity = 232448;
|
||||
|
||||
static std::vector<int> get_block_m_candidates(const KernelType& kernel_type, const cute::UMMA::Major& major_a, const int& m) {
|
||||
std::vector<int> candidates{64, 128, 256};
|
||||
if ((kernel_type == KernelType::Kernel1D2D or kernel_type == KernelType::KernelNoSF) and major_a == cute::UMMA::Major::K) {
|
||||
// NOTES: `block_m = 16/32` is smaller than MMA M size, should be careful in handling this
|
||||
if (m <= 16) candidates.push_back(16);
|
||||
if (m <= 32) candidates.push_back(32);
|
||||
|
||||
static std::vector<Layout> get_layout_candidates(const GemmDesc& desc) {
|
||||
// Block M candidates
|
||||
std::vector<int> block_m_candidates;
|
||||
if (desc.gemm_type == GemmType::Normal or
|
||||
desc.gemm_type == GemmType::Batched or
|
||||
desc.gemm_type == GemmType::KGroupedContiguous) {
|
||||
// TODO: check 256's performance
|
||||
block_m_candidates = {64, 128};
|
||||
// NOTES: smaller block M can avoid TMA L2 OOB bound
|
||||
if (desc.m <= 16) block_m_candidates.push_back(16);
|
||||
if (desc.m <= 32) block_m_candidates.push_back(32);
|
||||
|
||||
// BF16 output GEMM supports 256
|
||||
if (desc.cd_dtype != torch::kFloat)
|
||||
block_m_candidates.push_back(256);
|
||||
} else if (desc.gemm_type == GemmType::MGroupedContiguous or
|
||||
desc.gemm_type == GemmType::MGroupedContiguousWithPsumLayout) {
|
||||
block_m_candidates = std::vector{heuristics_runtime->get_mk_alignment_for_contiguous_layout()};
|
||||
} else if (desc.gemm_type == GemmType::MGroupedMasked) {
|
||||
block_m_candidates = {64, 128};
|
||||
}
|
||||
return candidates;
|
||||
}
|
||||
|
||||
static std::vector<int> get_block_n_candidates(const KernelType& kernel_type, const at::ScalarType& cd_dtype) {
|
||||
int start = 16;
|
||||
|
||||
// Block N candidates
|
||||
std::vector<int> block_n_candidates;
|
||||
int step = std::lcm(16, heuristics_runtime->get_block_n_multiple_of());
|
||||
int start = step;
|
||||
// Avoid bank conflicts for 1D1D kernel FP32 output
|
||||
std::vector<int> candidates;
|
||||
if (kernel_type == KernelType::Kernel1D1D and cd_dtype == torch::kFloat) {
|
||||
candidates.push_back(16);
|
||||
if (desc.kernel_type == KernelType::Kernel1D1D and desc.cd_dtype == torch::kFloat) {
|
||||
DG_HOST_ASSERT(desc.major_a == cute::UMMA::Major::K);
|
||||
DG_HOST_ASSERT(desc.major_b == cute::UMMA::Major::K);
|
||||
start = 24;
|
||||
block_n_candidates.push_back(16);
|
||||
}
|
||||
// Register spills
|
||||
int end = 256;
|
||||
if (desc.kernel_type == KernelType::Kernel1D2D)
|
||||
end = 192;
|
||||
if (desc.kernel_type == KernelType::Kernel1D1D)
|
||||
end = 160;
|
||||
// Enumerate
|
||||
for (int i = start; i <= end; i += step)
|
||||
block_n_candidates.push_back(i);
|
||||
|
||||
// Block K is always in a fixed manner
|
||||
const int block_k = 128 / get_element_size(desc.get_mma_kind());
|
||||
|
||||
// Disable multicast for performance
|
||||
const bool disable_multicast =
|
||||
// The number of k-groups is large (a heuristic)
|
||||
(desc.gemm_type == GemmType::KGroupedContiguous and desc.num_groups > 4) or
|
||||
// Not supported
|
||||
(desc.gemm_type == GemmType::Batched);
|
||||
|
||||
// Enumerate all candidates
|
||||
std::vector<Layout> candidates;
|
||||
for (int cluster_m = 1; cluster_m <= (disable_multicast ? 1 : 2); ++ cluster_m) {
|
||||
for (int cluster_n = 1; cluster_n <= (disable_multicast ? 1 : 2); ++ cluster_n) {
|
||||
// We only support cluster 2
|
||||
if (cluster_m * cluster_n > 2)
|
||||
continue;
|
||||
|
||||
// SM count must be divisible
|
||||
if (desc.num_sms % (cluster_m * cluster_n) != 0)
|
||||
continue;
|
||||
|
||||
for (int block_m: block_m_candidates) {
|
||||
for (int block_n: block_n_candidates) {
|
||||
// 1D2D kernel unroll requirement
|
||||
if (desc.kernel_type == KernelType::Kernel1D2D and block_n > block_k and (block_n % (block_n - block_k) != 0 and block_k % (block_n - block_k) != 0))
|
||||
continue;
|
||||
|
||||
// Multicast legality for masked layout
|
||||
// TODO: add some comments about it
|
||||
if ((desc.gemm_type == GemmType::MGroupedMasked or desc.gemm_type == GemmType::MGroupedContiguousWithPsumLayout) and
|
||||
ceil_div(desc.n, block_n) % (cluster_m * cluster_n) != 0)
|
||||
continue;
|
||||
|
||||
// The block sizes cannot be too large (for enough registers), so at least one dim less than 128
|
||||
if (block_m > 128 and block_n > 128)
|
||||
continue;
|
||||
|
||||
// Calculate swizzling
|
||||
const auto layout = Layout{0, block_m, block_n, block_k, cluster_m, cluster_n};
|
||||
const auto storage_config = get_storage_config(desc, layout);
|
||||
|
||||
// Make sure swizzling is large enough (32B's performance is low)
|
||||
if (storage_config.swizzle_a_mode % 64 != 0 or storage_config.swizzle_b_mode % 64 != 0)
|
||||
continue;
|
||||
|
||||
// To hide TMA latency, the stage count should be at least 3; for small matrices, at least 4
|
||||
int num_stages = get_pipeline_config(desc, layout, storage_config).num_stages;
|
||||
if (num_stages < 3 or (block_m * block_n < 128 * 192 and num_stages < 4))
|
||||
continue;
|
||||
|
||||
candidates.push_back(layout);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Push the strided options
|
||||
for (int i = start; i <= 256; i += 16)
|
||||
candidates.push_back(i);
|
||||
DG_HOST_ASSERT(not candidates.empty());
|
||||
return candidates;
|
||||
}
|
||||
|
||||
static int get_ab_load_block_m(const MulticastConfig& multicast_config, const int& block_m) {
|
||||
return block_m;
|
||||
}
|
||||
|
||||
static int get_ab_load_block_n(const MulticastConfig& multicast_config, const int& block_n) {
|
||||
return block_n;
|
||||
}
|
||||
|
||||
static int get_cd_store_block_m(const int& block_m, const bool& single_warpgroup_sync = false) {
|
||||
static StorageConfig get_storage_config(const GemmDesc& desc, const Layout& layout) {
|
||||
constexpr int wgmma_m = 64;
|
||||
return single_warpgroup_sync ? wgmma_m : block_m;
|
||||
}
|
||||
|
||||
static int get_cd_store_block_n(const int& block_n) {
|
||||
return block_n;
|
||||
}
|
||||
// Load/store block sizes (w/o consideration of swizzling atoms, w/ consideration of loop atoms)
|
||||
// TODO: support swap AB
|
||||
DG_HOST_ASSERT(layout.swap_ab == 0);
|
||||
const auto load_block_m = layout.block_m;
|
||||
const auto load_block_n = layout.block_n;
|
||||
// 1D1D kernel will do single warp-group stores
|
||||
const auto store_block_m = desc.kernel_type == KernelType::Kernel1D1D ? wgmma_m : layout.block_m;
|
||||
const auto store_block_n = layout.block_n;
|
||||
|
||||
static bool enable_cd_swizzle(const at::ScalarType& cd_dtype) {
|
||||
return cd_dtype != torch::kFloat;
|
||||
}
|
||||
|
||||
static bool is_block_size_legal(const KernelType& kernel_type,
|
||||
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
|
||||
const MmaKind& mma_kind, const at::ScalarType& cd_dtype,
|
||||
const int& m, const int& n, const int& k,
|
||||
const int& block_m, const int& block_n, const int& block_k) {
|
||||
// SM90 FP32 output does not support `block_m == 256`
|
||||
if (cd_dtype == at::kFloat and block_m == 256)
|
||||
return false;
|
||||
|
||||
// Avoid large C/D shared memory for FP32 output
|
||||
// Ensure `num_stages >= 4` (for 1D1D Kernel), `num_stages >= 3` (for No SF kernel)
|
||||
if (block_n > 128 and cd_dtype == torch::kFloat) {
|
||||
if (kernel_type == KernelType::Kernel1D1D and block_n > 152)
|
||||
return false;
|
||||
if (kernel_type == KernelType::KernelNoSF and block_n > 200)
|
||||
return false;
|
||||
}
|
||||
|
||||
// When B is N Major, use swizzle 128B for better performance; only affects SM90 BF16 GEMM
|
||||
if (major_b == cute::UMMA::Major::MN and block_n >= 128 and block_n % 64 != 0)
|
||||
return false;
|
||||
|
||||
// Too many scaling factors in a single block: `block_n > block_k and std::gcd(block_n, block_k) != block_n - block_k`
|
||||
// Or too many register spills
|
||||
if (block_n > 128 and kernel_type == KernelType::Kernel1D2D and (block_n != 144 and block_n != 160 and block_n != 192))
|
||||
return false;
|
||||
|
||||
// The block sizes cannot be too large (for enough registers), so at least one dim less than 128
|
||||
return block_m <= 128 or block_n <= 128;
|
||||
}
|
||||
|
||||
static bool is_num_stages_legal(const MmaKind& mma_kind, const at::ScalarType& cd_dtype,
|
||||
const int& num_stages,
|
||||
const int& block_m, const int& block_n, const int& block_k) {
|
||||
// Unrolling both stages and `num_former_iters` will cause large code size
|
||||
if (mma_kind == MmaKind::MXFP8FP4 and block_k % block_n != 0 and block_k / std::gcd(block_n, block_k) <= 4)
|
||||
return num_stages <= 4;
|
||||
return true;
|
||||
}
|
||||
|
||||
static std::pair<bool, bool> get_multicast_legality(const GemmType& gemm_type, const int& num_groups,
|
||||
const int& m, const int& n, const int& block_m, const int& block_n,
|
||||
const int& num_sms) {
|
||||
// Disable multicast when the number of k-groups is large (a heuristic)
|
||||
if (gemm_type == GemmType::KGroupedContiguous and num_groups > 4)
|
||||
return {false, false};
|
||||
|
||||
if (gemm_type == GemmType::Batched)
|
||||
return {false, false};
|
||||
// Decide swizzling by the inner dim
|
||||
const auto swizzle_mode_a = get_swizzle_mode(
|
||||
desc.major_a == cute::UMMA::Major::K ? layout.block_k : load_block_m, c10::elementSize(desc.a_dtype));
|
||||
const auto swizzle_mode_b = get_swizzle_mode(
|
||||
desc.major_b == cute::UMMA::Major::K ? layout.block_k : load_block_n, c10::elementSize(desc.b_dtype));
|
||||
// We only enable swizzling for non-FP32 outputs
|
||||
const auto swizzle_mode_cd = desc.cd_dtype != torch::kFloat ?
|
||||
get_swizzle_mode(store_block_n, c10::elementSize(desc.cd_dtype)) : 0;
|
||||
|
||||
return {
|
||||
is_multicast_legal(n, block_n, 2, num_sms, gemm_type == GemmType::MGroupedMasked),
|
||||
// For masked GEMM layout, divisibility on N is also required as we must ensure the total number of blocks is even
|
||||
is_multicast_legal(m, block_m, 2, num_sms, false)
|
||||
and (gemm_type != GemmType::MGroupedMasked or is_multicast_legal(n, block_n, 2, num_sms, true))
|
||||
load_block_m, load_block_n,
|
||||
store_block_m, store_block_n,
|
||||
swizzle_mode_a, swizzle_mode_b, swizzle_mode_cd
|
||||
};
|
||||
}
|
||||
|
||||
static ThreadConfig get_thread_config(const KernelType& kernel_type,
|
||||
const int& block_m, const int& block_n) {
|
||||
return ThreadConfig::sm90(128, (block_m <= 64 ? 1 : 2) * 128);
|
||||
}
|
||||
static PipelineConfig get_pipeline_config(const GemmDesc& desc, const Layout& layout, const StorageConfig& storage_config) {
|
||||
constexpr int kNumMaxStages = 16;
|
||||
|
||||
static int get_smem_cd_size(const KernelType& kernel_type,
|
||||
const int& block_m, const int& block_n,
|
||||
const int& swizzle_cd_mode, const at::ScalarType& cd_dtype) {
|
||||
// TODO: consider swap AB
|
||||
// C/D for TMA stores
|
||||
// NOTES: 1024 is for TMA swizzling alignment requirement
|
||||
return align(block_m * block_n * static_cast<int>(c10::elementSize(cd_dtype)), 1024);
|
||||
const int smem_cd =
|
||||
align(layout.block_m * layout.block_n * static_cast<int>(c10::elementSize(desc.cd_dtype)), 1024);
|
||||
const int smem_barriers = kNumMaxStages * 8 * 2;
|
||||
|
||||
// Calculate A/B per stages
|
||||
const int smem_a_per_stage = storage_config.load_block_m * layout.block_k * c10::elementSize(desc.a_dtype);
|
||||
const int smem_b_per_stage = storage_config.load_block_n * layout.block_k * c10::elementSize(desc.b_dtype);
|
||||
|
||||
// Calculate SF A/B per stages
|
||||
const int smem_sfa_per_stage = desc.kernel_type == KernelType::KernelNoSF ?
|
||||
0 : align(layout.block_m * static_cast<int>(sizeof(float)), 128);
|
||||
const int smem_sfb_per_stage = desc.kernel_type != KernelType::Kernel1D1D ?
|
||||
0 : align(layout.block_n * static_cast<int>(sizeof(float)), 128);
|
||||
|
||||
// Extra SFB sizes for 1D2D kernels
|
||||
const int use_uniform_sfb = layout.block_k % layout.block_n == 0 ? 1 : 2;
|
||||
const int smem_extra_sfb = desc.kernel_type != KernelType::Kernel1D2D ?
|
||||
0 : align<int>(ceil_div(desc.k, layout.block_k) * static_cast<int>(sizeof(float)) * use_uniform_sfb, 8);
|
||||
|
||||
// Extra tensormap for 1D1D kernels
|
||||
const int smem_tensormap =
|
||||
desc.gemm_type == GemmType::KGroupedContiguous ? 4 * static_cast<int>(sizeof(CUtensorMap)) : 0;
|
||||
|
||||
// Calculate stages
|
||||
const int smem_extra = smem_cd + smem_barriers + smem_extra_sfb + smem_tensormap;
|
||||
const int smem_per_stage = smem_a_per_stage + smem_b_per_stage + smem_sfa_per_stage + smem_sfb_per_stage;
|
||||
const int num_stages = std::min(
|
||||
(smem_capacity - smem_extra) / smem_per_stage,
|
||||
kNumMaxStages);
|
||||
return {
|
||||
smem_extra + num_stages * smem_per_stage,
|
||||
num_stages
|
||||
};
|
||||
}
|
||||
|
||||
static std::pair<int, int> get_sf_smem_size_per_stage(const KernelType& kernel_type,
|
||||
const int& block_m, const int& block_n, const int& block_k,
|
||||
const MmaKind& mma_kind, const at::ScalarType& cd_dtype) {
|
||||
if (mma_kind == MmaKind::BF16)
|
||||
return {0, 0};
|
||||
|
||||
// NOTES: 128 is for 2D TMA alignment requirement
|
||||
int smem_sfa_per_stage = align(block_m * static_cast<int>(sizeof(float)), 128);
|
||||
int smem_sfb_per_stage = 0;
|
||||
if (kernel_type == KernelType::Kernel1D1D)
|
||||
smem_sfb_per_stage = align(block_n * 4, 128);
|
||||
return {smem_sfa_per_stage, smem_sfb_per_stage};
|
||||
static LaunchConfig get_launch_config(const GemmDesc& desc, const Layout& layout) {
|
||||
const int num_tma_threads = 128;
|
||||
const int num_math_threads = layout.block_m <= 64 ? 128 : 256;
|
||||
return {
|
||||
desc.num_sms,
|
||||
layout.get_cluster_size(),
|
||||
num_tma_threads + num_math_threads,
|
||||
num_tma_threads, num_math_threads,
|
||||
0, 0 // Meaningless for SM90
|
||||
};
|
||||
}
|
||||
|
||||
static int get_extra_sfb_smem_size(const int& m, const int& n, const int& k,
|
||||
const int& block_m, const int& block_n, const int& block_k) {
|
||||
const auto& use_uniform_sfb = block_k % block_n == 0 ? 1 : 2;
|
||||
return align<int>(ceil_div(k, block_k) * static_cast<int>(sizeof(float)) * use_uniform_sfb, 8);
|
||||
static LayoutInfo get_layout_info(const GemmDesc& desc, const Layout& layout) {
|
||||
const auto num_blocks =
|
||||
ceil_div(desc.get_expected_m(), layout.block_m) *
|
||||
ceil_div(desc.get_expected_n(), layout.block_n) *
|
||||
desc.get_expected_num_groups();
|
||||
const auto num_waves = ceil_div(num_blocks, desc.num_sms);
|
||||
const auto num_last_blocks = num_blocks % desc.num_sms;
|
||||
const auto last_wave_util = num_last_blocks == 0 ? desc.num_sms : num_last_blocks;
|
||||
|
||||
// Utils
|
||||
const int l2_bandwidth_per_cycle = std::min(64. * desc.num_sms, 8e6 / (1.3e3)); // B/cycle
|
||||
const int l1_bandwidth_per_cycle = 128 * desc.num_sms; // B/cycle
|
||||
const int wgmma_m = 64;
|
||||
const int elem_size_ab = c10::elementSize(desc.a_dtype);
|
||||
const int elem_size_cd = c10::elementSize(desc.cd_dtype);
|
||||
DG_HOST_ASSERT(desc.a_dtype == desc.b_dtype);
|
||||
|
||||
// Data movement per block
|
||||
int64_t expected_k = desc.get_expected_k();
|
||||
int64_t num_bytes_l2_ab = expected_k * (layout.block_m / layout.cluster_n + layout.block_n / layout.cluster_m) * elem_size_ab;
|
||||
int64_t num_bytes_l1_ab = expected_k * (layout.block_m + layout.block_n) * elem_size_ab;
|
||||
int64_t num_bytes_l1_tc = expected_k * (std::max(wgmma_m, layout.block_m) + layout.block_n) * elem_size_ab
|
||||
+ layout.block_m * layout.block_n * elem_size_cd;
|
||||
int64_t num_bytes_l1_l2_cd = layout.block_m * layout.block_n * elem_size_cd * (desc.with_accumulation ? 2 : 1);
|
||||
|
||||
// HBM bandwidth and total compute (Tensor/CUDA cores) are constant across configs
|
||||
// We only model L1/L2 cycles as they are the primary variables between configs
|
||||
int64_t num_l2_cycles = (num_bytes_l2_ab + num_bytes_l1_l2_cd) * num_blocks / l2_bandwidth_per_cycle;
|
||||
int64_t num_l1_cycles = (num_bytes_l1_ab + num_bytes_l1_tc + num_bytes_l1_l2_cd) * num_blocks / l1_bandwidth_per_cycle;
|
||||
float wave_efficiency = static_cast<float>(num_blocks) / (num_waves * desc.num_sms);
|
||||
int64_t num_cycles = std::max(num_l1_cycles, num_l2_cycles) / wave_efficiency;
|
||||
|
||||
// Disable multicasting if only one wave exists
|
||||
if (layout.cluster_n * layout.cluster_m > 1 and num_waves <= 1)
|
||||
num_cycles = std::numeric_limits<int64_t>::max();
|
||||
|
||||
return {num_waves, last_wave_util, num_cycles, layout};
|
||||
}
|
||||
|
||||
static int get_barrier_smem_size(const int& num_stages) {
|
||||
return num_stages * 8 * 2;
|
||||
}
|
||||
|
||||
static int get_tmem_ptr_smem_size() {
|
||||
return 0;
|
||||
}
|
||||
|
||||
static int get_tensormap_smem_size(const GemmType& gemm_type) {
|
||||
return gemm_type == GemmType::KGroupedContiguous ? 4 * static_cast<int>(sizeof(CUtensorMap)) : 0;
|
||||
// A regular comparator
|
||||
static bool compare(const LayoutInfo& a, const LayoutInfo& b) {
|
||||
return a.num_cycles < b.num_cycles;
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@@ -0,0 +1,23 @@
|
||||
#pragma once
|
||||
|
||||
#include <cute/arch/mma_sm100_desc.hpp>
|
||||
// Reuse some types in the JIT modules
|
||||
#include <deep_gemm/common/types.cuh>
|
||||
|
||||
#include "common.hpp"
|
||||
#include "../../utils/exception.hpp"
|
||||
|
||||
namespace deep_gemm {
|
||||
|
||||
template <typename size_type_t>
|
||||
static int get_swizzle_mode(const int& block_size, const size_type_t& elem_size) {
|
||||
// `> 0` means interleaving
|
||||
// 16B actually means non-swizzling (but interleaving)
|
||||
for (const int& mode: {128, 64, 32, 16}) {
|
||||
if ((block_size * static_cast<int>(elem_size)) % mode == 0)
|
||||
return mode;
|
||||
}
|
||||
DG_HOST_UNREACHABLE("Unreachable");
|
||||
}
|
||||
|
||||
} // namespace deep_gemm
|
||||
Reference in New Issue
Block a user