biondizzle/DeepGEMM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[Public release 26/04] Introducing Mega MoE, FP4 Indexer and other features/fixes (#304)

Browse Source 
* 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
2026-04-17 09:45:14 +08:00
committed by GitHub
parent d30fc36c8f
commit 7f2a703ed5
109 changed files with 12101 additions and 3219 deletions
Download Patch File Download Diff File Expand all files Collapse all files

347
csrc/jit_kernels/heuristics/common.hpp
View File

@@ -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

171
csrc/jit_kernels/heuristics/config.hpp Normal file
View File

@@ -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

211
csrc/jit_kernels/heuristics/mega_moe.hpp Normal file
View File

@@ -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

62
csrc/jit_kernels/heuristics/runtime.hpp Normal file
View File

@@ -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

344
csrc/jit_kernels/heuristics/sm100.hpp
View File

@@ -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;
}
};

330
csrc/jit_kernels/heuristics/sm90.hpp
View File

@@ -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;
}
};

23
csrc/jit_kernels/heuristics/utils.hpp Normal file
View File

@@ -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

2
csrc/jit_kernels/impls/epilogue.hpp
View File

@@ -6,7 +6,7 @@
namespace deep_gemm {
static std::string get_default_epilogue_type(const std::optional<std::string>& epilogue_type) {
return epilogue_type.value_or("EpilogueIdentity");
return epilogue_type.value_or("epilogue::transform::EpilogueIdentity");
}
} // namespace deep_gemm

72
csrc/jit_kernels/impls/runtime_utils.hpp
View File

@@ -20,6 +20,9 @@ static int get_non_contiguous_dim(const cute::UMMA::Major& major) {
}
static int get_compiled_dim(const int& dim, const char& name, const std::string& compiled_dims) {
if (heuristics_runtime->get_ignore_compile_dims())
return 0;
for (const char& c: compiled_dims) {
if (name == c)
return dim;
@@ -58,8 +61,19 @@ static std::string to_string(const at::ScalarType& dtype) {
}
}
static std::string to_string(const float& v) {
if (std::isfinite(v)) {
return fmt::format(R"({:a}f)", v);
} else if (std::isinf(v)) {
return v > 0 ? "cute::numeric_limits<float>::infinity()"
: "-cute::numeric_limits<float>::infinity()";
}
DG_HOST_UNREACHABLE("NaN input is not supported");
}
static CUtensorMapDataType aten_dtype_to_tensor_map_dtype(const at::ScalarType& dtype,
const bool& allow_tf32) {
const bool& allow_tf32,
const bool& fp4_unpacked_smem) {
if (allow_tf32 and dtype == torch::kFloat)
return CU_TENSOR_MAP_DATA_TYPE_TFLOAT32;
@@ -68,13 +82,16 @@ static CUtensorMapDataType aten_dtype_to_tensor_map_dtype(const at::ScalarType&
case torch::kFloat: return CU_TENSOR_MAP_DATA_TYPE_FLOAT32;
case torch::kBFloat16: return CU_TENSOR_MAP_DATA_TYPE_BFLOAT16;
case torch::kFloat8_e4m3fn: return CU_TENSOR_MAP_DATA_TYPE_UINT8;
case kPackedFP4: return CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B;
#if CUDA_VERSION >= 12080
case kPackedFP4: return fp4_unpacked_smem ? CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B
: CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B;
#endif
default: DG_HOST_UNREACHABLE("Unsupported dtype");
}
}
static CUtensorMapSwizzle mode_into_tensor_map_swizzle(const int& mode, const int& base) {
#if CUDART_VERSION >= 12080
#if CUDA_VERSION >= 12080
if (base != 0) {
DG_HOST_ASSERT(base == 32 and mode == 128);
return CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B;
@@ -97,14 +114,20 @@ static CUtensorMap make_tma_2d_desc(const torch::Tensor& t,
int smem_inner_dim, int smem_outer_dim,
const int& gmem_outer_stride,
const int& swizzle_mode, const int& swizzle_base = 0,
const bool& allow_tf32 = false) {
const auto& elem_size = static_cast<int>(t.element_size());
const bool& allow_tf32 = false,
const bool& fp4_unpacked_smem = true) {
const auto elem_size = static_cast<int>(t.element_size());
if (swizzle_mode != 0)
smem_inner_dim = swizzle_mode / elem_size;
// Inner dim must be a multiple of 64B for .b4x16_p64
if (t.scalar_type() == kPackedFP4)
DG_HOST_ASSERT(gmem_inner_dim % 128 == 0);
if (t.scalar_type() == kPackedFP4) {
// Inner dim must be a multiple of 64B for .b4x16_p64
DG_HOST_ASSERT(not fp4_unpacked_smem or gmem_inner_dim % 128 == 0);
// Fix FP4 packed smem
if (not fp4_unpacked_smem and swizzle_mode != 0)
smem_inner_dim = swizzle_mode * 2;
}
CUtensorMap tensor_map;
const cuuint64_t gmem_dims[2] = {static_cast<cuuint64_t>(gmem_inner_dim), static_cast<cuuint64_t>(gmem_outer_dim)};
@@ -112,12 +135,13 @@ static CUtensorMap make_tma_2d_desc(const torch::Tensor& t,
const cuuint64_t gmem_strides[1] = {static_cast<cuuint64_t>(gmem_outer_stride * elem_size), };
const cuuint32_t elem_strides[2] = {1, 1};
if (get_env<int>("DG_JIT_DEBUG")) {
printf("Making TMA desc: global memory: %d %d, shared memory: %d %d, outer stride: %d, swizzle: %d (base: %d), elem size: %d\n",
printf("Making TMA desc: global memory: %d %d, shared memory: %d %d, outer stride: %d, swizzle: %d (base: %d), elem size: %d, pointer: %llu\n",
gmem_inner_dim, gmem_outer_dim, smem_inner_dim, smem_outer_dim,
gmem_outer_stride, swizzle_mode, swizzle_base, elem_size);
gmem_outer_stride, swizzle_mode, swizzle_base, elem_size,
reinterpret_cast<unsigned long long>(t.data_ptr()));
}
DG_CUDA_DRIVER_CHECK(lazy_cuTensorMapEncodeTiled(
&tensor_map, aten_dtype_to_tensor_map_dtype(t.scalar_type(), allow_tf32),
&tensor_map, aten_dtype_to_tensor_map_dtype(t.scalar_type(), allow_tf32, fp4_unpacked_smem),
2, t.data_ptr(), gmem_dims, gmem_strides, smem_dims, elem_strides,
CU_TENSOR_MAP_INTERLEAVE_NONE, mode_into_tensor_map_swizzle(swizzle_mode, swizzle_base),
CU_TENSOR_MAP_L2_PROMOTION_L2_256B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE));
@@ -129,14 +153,20 @@ static CUtensorMap make_tma_3d_desc(const torch::Tensor& t,
int smem_dim_0, int smem_dim_1, int smem_dim_2,
const int& gmem_stride_0, const int& gmem_stride_1,
const int& swizzle_mode, const int& swizzle_base = 0,
const bool& allow_tf32 = false) {
const auto& elem_size = static_cast<int>(t.element_size());
const bool& allow_tf32 = false,
const bool& fp4_unpacked_smem = true) {
const auto elem_size = static_cast<int>(t.element_size());
if (swizzle_mode != 0)
smem_dim_0 = swizzle_mode / elem_size;
// Inner dim must be a multiple of 64B for .b4x16_p64
if (t.scalar_type() == kPackedFP4)
DG_HOST_ASSERT(gmem_dim_0 % 128 == 0);
if (t.scalar_type() == kPackedFP4) {
// Inner dim must be a multiple of 64B for .b4x16_p64
DG_HOST_ASSERT(not fp4_unpacked_smem or gmem_dim_0 % 128 == 0);
// Fix fp4 packed smem
if (not fp4_unpacked_smem and swizzle_mode != 0)
smem_dim_0 = swizzle_mode * 2;
}
CUtensorMap tensor_map;
const cuuint64_t gmem_dims[3] = {static_cast<cuuint64_t>(gmem_dim_0), static_cast<cuuint64_t>(gmem_dim_1), static_cast<cuuint64_t>(gmem_dim_2),};
@@ -149,7 +179,7 @@ static CUtensorMap make_tma_3d_desc(const torch::Tensor& t,
gmem_stride_0, gmem_stride_1, swizzle_mode, elem_size);
}
DG_CUDA_DRIVER_CHECK(lazy_cuTensorMapEncodeTiled(
&tensor_map, aten_dtype_to_tensor_map_dtype(t.scalar_type(), allow_tf32),
&tensor_map, aten_dtype_to_tensor_map_dtype(t.scalar_type(), allow_tf32, fp4_unpacked_smem),
3, t.data_ptr(), gmem_dims, gmem_strides, smem_dims, elem_strides,
CU_TENSOR_MAP_INTERLEAVE_NONE, mode_into_tensor_map_swizzle(swizzle_mode, swizzle_base),
CU_TENSOR_MAP_L2_PROMOTION_L2_256B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE));
@@ -166,8 +196,8 @@ static CUtensorMap make_tma_a_desc(const cute::UMMA::Major& major,
const bool& allow_tf32 = false) {
if (num_groups > 1)
DG_HOST_ASSERT(major == cute::UMMA::Major::K);
const auto& [gmem_inner_dim, gmem_outer_dim] = get_inner_outer_dims(major, shape_k, shape_m * num_groups);
const auto& [smem_inner_dim, smem_outer_dim] = get_inner_outer_dims(major, block_k, block_m);
const auto [gmem_inner_dim, gmem_outer_dim] = get_inner_outer_dims(major, shape_k, shape_m * num_groups);
const auto [smem_inner_dim, smem_outer_dim] = get_inner_outer_dims(major, block_k, block_m);
return make_tma_2d_desc(t,
gmem_inner_dim, gmem_outer_dim,
smem_inner_dim, smem_outer_dim,
@@ -184,8 +214,8 @@ static CUtensorMap make_tma_b_desc(const cute::UMMA::Major& major,
const int& num_groups,
const int& swizzle_mode, const int& swizzle_base = 0,
const bool& allow_tf32 = false) {
const auto& [gmem_inner_dim, gmem_outer_dim] = get_inner_outer_dims(major, shape_k, shape_n);
const auto& [smem_inner_dim, smem_outer_dim] = get_inner_outer_dims(major, block_k, block_n);
const auto [gmem_inner_dim, gmem_outer_dim] = get_inner_outer_dims(major, shape_k, shape_n);
const auto [smem_inner_dim, smem_outer_dim] = get_inner_outer_dims(major, block_k, block_n);
// `num_groups` is always applied into the outer dimensions
return make_tma_2d_desc(t,

426
csrc/jit_kernels/impls/sm100_bf16_gemm.hpp
View File

@@ -16,9 +16,7 @@ namespace deep_gemm {
class SM100BF16GemmRuntime final: public LaunchRuntime<SM100BF16GemmRuntime> {
public:
struct Args {
int m, n, k, num_groups;
const std::string& compiled_dims;
GemmDesc gemm_desc;
GemmConfig gemm_config;
LaunchArgs launch_args;
@@ -45,28 +43,32 @@ static void __instantiate_kernel() {{
{}, {},
{}, {},
{},
{},
{}, {}, {},
{}
>);
}};
)",
to_string(args.gemm_config.major_a), to_string(args.gemm_config.major_b),
get_compiled_dim(args.m, 'm', args.compiled_dims), get_compiled_dim(args.n, 'n', args.compiled_dims), get_compiled_dim(args.k, 'k', args.compiled_dims),
args.gemm_config.block_m, args.gemm_config.block_n, args.gemm_config.block_k,
args.num_groups,
args.gemm_config.smem_config.swizzle_a_mode, args.gemm_config.smem_config.swizzle_b_mode, args.gemm_config.smem_config.swizzle_cd_mode,
args.gemm_config.num_stages,
args.gemm_config.thread_config.num_non_epilogue_threads, args.gemm_config.thread_config.num_epilogue_threads,
args.gemm_config.multicast_config.num_multicast, args.gemm_config.multicast_config.is_multicast_on_a,
args.gemm_config.num_sms,
to_string(args.gemm_config.gemm_type), args.gemm_config.with_accumulation, to_string(args.gemm_config.cd_dtype),
args.gemm_config.tc_util);
to_string(args.gemm_desc.major_a), to_string(args.gemm_desc.major_b),
get_compiled_dim(args.gemm_desc.m, 'm', args.gemm_desc.compiled_dims),
get_compiled_dim(args.gemm_desc.n, 'n', args.gemm_desc.compiled_dims),
get_compiled_dim(args.gemm_desc.k, 'k', args.gemm_desc.compiled_dims),
args.gemm_config.layout.block_m, args.gemm_config.layout.block_n, args.gemm_config.layout.block_k,
args.gemm_desc.num_groups,
args.gemm_config.storage_config.swizzle_a_mode, args.gemm_config.storage_config.swizzle_b_mode, args.gemm_config.storage_config.swizzle_cd_mode,
args.gemm_config.pipeline_config.num_stages,
args.gemm_config.launch_config.num_non_epilogue_threads, args.gemm_config.launch_config.num_epilogue_threads,
args.gemm_config.layout.get_cluster_size(), args.gemm_config.layout.cluster_n > 1,
args.gemm_config.launch_config.num_sms,
args.gemm_config.layout.swap_ab,
to_string(args.gemm_desc.gemm_type), args.gemm_desc.with_accumulation, to_string(args.gemm_desc.cd_dtype),
args.gemm_desc.tc_util);
}
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
// TODO: optimize `args` copy
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.grouped_layout, args.m, args.n, args.k,
args.grouped_layout, args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k,
args.tensor_map_a, args.tensor_map_b,
args.tensor_map_cd));
}
@@ -79,45 +81,49 @@ static void sm100_bf16_gemm(const torch::Tensor& a,
const int& m, const int& n, const int& k,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims) {
const auto& config = get_best_config<SM100ArchSpec>(
GemmType::Normal, KernelType::KernelNoSF,
m, n, k, 1, major_a, major_b,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), c.has_value(),
device_runtime->get_num_sms());
const auto desc = GemmDesc {
.gemm_type = GemmType::Normal,
.kernel_type = KernelType::KernelNoSF,
.m = m, .n = n, .k = k, .num_groups = 1,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = c.has_value(),
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims
};
const auto config = get_best_config<SM100ArchSpec>(desc);
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM100ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(major_b, b, n, k,
SM100ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), 1,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), 1,
config.smem_config.swizzle_cd_mode);
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), 1,
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), 1,
config.storage_config.swizzle_cd_mode);
// Launch
const SM100BF16GemmRuntime::Args& args = {
.m = m, .n = n, .k = k,
.num_groups = 1,
.compiled_dims = compiled_dims,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd
};
const auto& code = SM100BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm100_bf16_gemm", code);
const auto code = SM100BF16GemmRuntime::generate(args);
const auto runtime = compiler->build("sm100_bf16_gemm", code);
SM100BF16GemmRuntime::launch(runtime, args);
}
@@ -130,53 +136,61 @@ static void sm100_m_grouped_bf16_gemm_contiguous(const torch::Tensor& a,
const std::string& compiled_dims,
const bool& use_psum_layout,
const std::optional<int>& expected_m_for_psum_layout) {
const auto& gemm_type = use_psum_layout ? GemmType::MGroupedContiguousWithPsumLayout : GemmType::MGroupedContiguous;
const auto gemm_type = use_psum_layout ?
GemmType::MGroupedContiguousWithPsumLayout : GemmType::MGroupedContiguous;
// Only psum layout can use expected m
if (expected_m_for_psum_layout)
DG_HOST_ASSERT(use_psum_layout);
// NOTES: If actual M is dynamic, estimate config via `num_groups` and `expected_m`.
// Otherwise, treat the contiguous layout as a whole.
const auto& m_for_config = expected_m_for_psum_layout.has_value() ? expected_m_for_psum_layout.value() : m;
const auto& num_groups_for_config = expected_m_for_psum_layout.has_value() ? num_groups : 1;
const auto desc = GemmDesc {
.gemm_type = gemm_type,
.kernel_type = KernelType::KernelNoSF,
.m = m, .n = n, .k = k, .num_groups = num_groups,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = false,
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims,
.expected_m = expected_m_for_psum_layout.value_or(m),
.expected_n = n, .expected_k = k,
.expected_num_groups = expected_m_for_psum_layout.has_value() ? num_groups : 1
};
const auto config = get_best_config<SM100ArchSpec>(desc);
const auto& config = get_best_config<SM100ArchSpec>(
gemm_type, KernelType::KernelNoSF,
// NOTES: `num_groups` is 1, since the contiguous layout is seen as a whole
m_for_config, n, k, num_groups_for_config, major_a, major_b,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), false,
device_runtime->get_num_sms());
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM100ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(major_b, b, n, k,
SM100ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), 1,
config.smem_config.swizzle_cd_mode);
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), 1,
config.storage_config.swizzle_cd_mode);
// Launch
const SM100BF16GemmRuntime::Args& args = {
.m = m, .n = n, .k = k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
const SM100BF16GemmRuntime::Args args = {
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.grouped_layout = grouped_layout.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd
};
const auto& code = SM100BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm100_bf16_m_grouped_gemm_contiguous", code);
const auto code = SM100BF16GemmRuntime::generate(args);
const auto runtime = compiler->build("sm100_bf16_m_grouped_gemm_contiguous", code);
SM100BF16GemmRuntime::launch(runtime, args);
}
@@ -188,45 +202,50 @@ static void sm100_m_grouped_bf16_gemm_masked(const torch::Tensor& a,
const int& expected_m,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims) {
const auto& config = get_best_config<SM100ArchSpec>(
GemmType::MGroupedMasked, KernelType::KernelNoSF,
expected_m, n, k, num_groups, major_a, major_b,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), false,
device_runtime->get_num_sms());
const auto desc = GemmDesc {
.gemm_type = GemmType::MGroupedMasked,
.kernel_type = KernelType::KernelNoSF,
.m = m, .n = n, .k = k, .num_groups = num_groups,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = false,
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims,
.expected_m = expected_m, .expected_n = n, .expected_k = k, .expected_num_groups = num_groups
};
const auto config = get_best_config<SM100ArchSpec>(desc);
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM100ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), num_groups,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(major_b, b, n, k,
SM100ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), num_groups,
config.smem_config.swizzle_cd_mode);
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), num_groups,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), num_groups,
config.storage_config.swizzle_cd_mode);
// Launch
const SM100BF16GemmRuntime::Args& args = {
.m = m, .n = n, .k = k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
const SM100BF16GemmRuntime::Args args = {
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.grouped_layout = masked_m.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd
};
const auto& code = SM100BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm100_bf16_m_grouped_gemm_masked", code);
const auto code = SM100BF16GemmRuntime::generate(args);
const auto runtime = compiler->build("sm100_bf16_m_grouped_gemm_masked", code);
SM100BF16GemmRuntime::launch(runtime, args);
}
@@ -241,54 +260,59 @@ static void sm100_bf16_k_grouped_gemm(const torch::Tensor& a,
DG_HOST_ASSERT(major_a == cute::UMMA::Major::MN and major_b == cute::UMMA::Major::MN);
int sum_k = 0;
for (const auto& k: ks) {
for (const auto k: ks) {
sum_k += k;
DG_HOST_ASSERT(k % 128 == 0);
}
const auto& num_groups = static_cast<int>(ks.size());
const auto num_groups = static_cast<int>(ks.size());
// Get config using max K for better performance
const auto& max_k = *std::max_element(ks.begin(), ks.end());
const auto& config = get_best_config<SM100ArchSpec>(
GemmType::KGroupedContiguous, KernelType::KernelNoSF,
m, n, max_k, num_groups, cute::UMMA::Major::MN, cute::UMMA::Major::MN,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), c.has_value(),
device_runtime->get_num_sms());
const auto max_k = *std::max_element(ks.begin(), ks.end());
const auto desc = GemmDesc {
.gemm_type = GemmType::KGroupedContiguous,
.kernel_type = KernelType::KernelNoSF,
.m = m, .n = n, .k = sum_k, .num_groups = num_groups,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = c.has_value(),
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims,
.expected_m = m, .expected_n = n, .expected_k = max_k, .expected_num_groups = num_groups
};
const auto config = get_best_config<SM100ArchSpec>(desc);
// Create tensor descriptors
const auto& tensor_map_a = make_tma_a_desc(cute::UMMA::Major::MN, a, m, sum_k,
SM100ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(0)), 1,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(cute::UMMA::Major::MN, b, n, sum_k,
SM100ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(0)), 1,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(1)), num_groups,
config.smem_config.swizzle_cd_mode);
const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::MN, a, m, sum_k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(0)), 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(cute::UMMA::Major::MN, b, n, sum_k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(0)), 1,
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(1)), num_groups,
config.storage_config.swizzle_cd_mode);
// Launch kernel
const SM100BF16GemmRuntime::Args& args = {
.m = m, .n = n, .k = sum_k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.grouped_layout = ks_tensor.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd
};
const auto& code = SM100BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm100_bf16_k_grouped_gemm", code);
const auto code = SM100BF16GemmRuntime::generate(args);
const auto runtime = compiler->build("sm100_bf16_k_grouped_gemm", code);
SM100BF16GemmRuntime::launch(runtime, args);
}
@@ -297,46 +321,46 @@ static void sm100_bf16_bhr_hdr_bhd(const torch::Tensor& tensor_a,
const torch::Tensor& tensor_d,
const int& b, const int& h, const int& r, const int& d,
const std::string& compiled_dims = "nk") {
const auto& config = get_best_config<SM100ArchSpec>(
GemmType::Batched, KernelType::KernelNoSF,
b, d, r, h, cute::UMMA::Major::K, cute::UMMA::Major::K,
tensor_a.scalar_type(), tensor_b.scalar_type(),
tensor_d.scalar_type(), false,
device_runtime->get_num_sms());
const auto desc = GemmDesc {
.gemm_type = GemmType::Batched,
.kernel_type = KernelType::KernelNoSF,
.m = b, .n = d, .k = r, .num_groups = h,
.a_dtype = tensor_a.scalar_type(), .b_dtype = tensor_b.scalar_type(),
.cd_dtype = tensor_d.scalar_type(),
.major_a = cute::UMMA::Major::K, .major_b = cute::UMMA::Major::K,
.with_accumulation = false,
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims
};
const auto config = get_best_config<SM100ArchSpec>(desc);
const int& load_block_m = SM100ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m);
const auto& tensor_map_a = make_tma_3d_desc(tensor_a, r, b, h,
config.block_k, load_block_m, 1,
tensor_a.stride(0), tensor_a.stride(1),
config.smem_config.swizzle_a_mode);
const int& load_block_n = SM100ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n);
const auto& tensor_map_b = make_tma_3d_desc(tensor_b, r, d, h,
config.block_k, load_block_n, 1,
tensor_b.stride(1), tensor_b.stride(0),
config.smem_config.swizzle_b_mode);
const int& store_block_m = SM100ArchSpec::get_cd_store_block_m(config.block_m);
const int& store_block_n = SM100ArchSpec::get_cd_store_block_n(config.block_n);
const auto& tensor_map_cd = make_tma_3d_desc(tensor_d, d, b, h,
store_block_n, store_block_m, 1,
tensor_d.stride(0), tensor_d.stride(1),
config.smem_config.swizzle_cd_mode);
const auto tensor_map_a = make_tma_3d_desc(tensor_a, r, b, h,
config.layout.block_k, config.storage_config.load_block_m, 1,
tensor_a.stride(0), tensor_a.stride(1),
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_3d_desc(tensor_b, r, d, h,
config.layout.block_k, config.storage_config.load_block_n, 1,
tensor_b.stride(1), tensor_b.stride(0),
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_3d_desc(tensor_d, d, b, h,
config.storage_config.store_block_n, config.storage_config.store_block_m, 1,
tensor_d.stride(0), tensor_d.stride(1),
config.storage_config.swizzle_cd_mode);
// Launch
const SM100BF16GemmRuntime::Args& args = {
.m = b, .n = d, .k = r,
.num_groups = h,
.compiled_dims = compiled_dims,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd
};
const auto& code = SM100BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm100_bf16_bhr_hdr_bhd", code);
const auto code = SM100BF16GemmRuntime::generate(args);
const auto runtime = compiler->build("sm100_bf16_bhr_hdr_bhd", code);
SM100BF16GemmRuntime::launch(runtime, args);
}
@@ -345,46 +369,46 @@ static void sm100_bf16_bhd_hdr_bhr(const torch::Tensor& tensor_a,
const torch::Tensor& tensor_d,
const int& b, const int& h, const int& r, const int& d,
const std::string& compiled_dims = "nk") {
const auto& config = get_best_config<SM100ArchSpec>(
GemmType::Batched, KernelType::KernelNoSF,
b, r, d, h, cute::UMMA::Major::K, cute::UMMA::Major::MN,
tensor_a.scalar_type(), tensor_b.scalar_type(),
tensor_d.scalar_type(), false,
device_runtime->get_num_sms());
const auto desc = GemmDesc {
.gemm_type = GemmType::Batched,
.kernel_type = KernelType::KernelNoSF,
.m = b, .n = r, .k = d, .num_groups = h,
.a_dtype = tensor_a.scalar_type(), .b_dtype = tensor_b.scalar_type(),
.cd_dtype = tensor_d.scalar_type(),
.major_a = cute::UMMA::Major::K, .major_b = cute::UMMA::Major::MN,
.with_accumulation = false,
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims
};
const auto config = get_best_config<SM100ArchSpec>(desc);
const int& load_block_m = SM100ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m);
const auto& tensor_map_a = make_tma_3d_desc(tensor_a, d, b, h,
config.block_k, load_block_m, 1,
tensor_a.stride(0), tensor_a.stride(1),
config.smem_config.swizzle_a_mode);
const int& load_block_n = SM100ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n);
const auto& tensor_map_b = make_tma_3d_desc(tensor_b, r, d, h,
load_block_n, config.block_k, 1,
tensor_b.stride(1), tensor_b.stride(0),
config.smem_config.swizzle_b_mode);
const int& store_block_m = SM100ArchSpec::get_cd_store_block_m(config.block_m);
const int& store_block_n = SM100ArchSpec::get_cd_store_block_n(config.block_n);
const auto& tensor_map_cd = make_tma_3d_desc(tensor_d, r, b, h,
store_block_n, store_block_m, 1,
tensor_d.stride(0), tensor_d.stride(1),
config.smem_config.swizzle_cd_mode);
const auto tensor_map_a = make_tma_3d_desc(tensor_a, d, b, h,
config.layout.block_k, config.storage_config.load_block_m, 1,
tensor_a.stride(0), tensor_a.stride(1),
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_3d_desc(tensor_b, r, d, h,
config.storage_config.load_block_n, config.layout.block_k, 1,
tensor_b.stride(1), tensor_b.stride(0),
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_3d_desc(tensor_d, r, b, h,
config.storage_config.store_block_n, config.storage_config.store_block_m, 1,
tensor_d.stride(0), tensor_d.stride(1),
config.storage_config.swizzle_cd_mode);
// Launch
const SM100BF16GemmRuntime::Args& args = {
.m = b, .n = r, .k = d,
.num_groups = h,
.compiled_dims = compiled_dims,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd
};
const auto& code = SM100BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm100_bf16_bhd_hdr_bhr", code);
const auto code = SM100BF16GemmRuntime::generate(args);
const auto runtime = compiler->build("sm100_bf16_bhd_hdr_bhr", code);
SM100BF16GemmRuntime::launch(runtime, args);
}

22
csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp
View File

@@ -85,11 +85,11 @@ static void sm100_bmn_bnk_mn_gemm(const torch::Tensor &a,
// NOTES: we select 4 as start, as it is tested to be faster than values > 4
int num_stages = 4, smem_size = 0;
while (true) {
const int& smem_cd = block_m * swizzle_cd_mode * 2;
const int& smem_a_per_stage = block_m * block_k * sizeof(cutlass::bfloat16_t);
const int& smem_b_per_stage = block_n * block_k * sizeof(cutlass::bfloat16_t);
const int& smem_barrier = SM100ArchSpec::get_barrier_smem_size(num_stages);
const int& smem_tmem_ptr = SM100ArchSpec::get_tmem_ptr_smem_size();
const int smem_cd = block_m * swizzle_cd_mode * 2;
const int smem_a_per_stage = block_m * block_k * sizeof(cutlass::bfloat16_t);
const int smem_b_per_stage = block_n * block_k * sizeof(cutlass::bfloat16_t);
const int smem_barrier = num_stages * 8 * 3 + 2 * 8 * 2 + 8;
const int smem_tmem_ptr = 4;
smem_size = 0;
smem_size += smem_cd;
@@ -112,11 +112,11 @@ static void sm100_bmn_bnk_mn_gemm(const torch::Tensor &a,
num_stages, smem_size, swizzle_ab_mode, swizzle_cd_mode);
}
const auto& tensor_map_a = make_tma_2d_desc(a, k, s * m, block_k, block_m, k, swizzle_ab_mode);
const auto& tensor_map_b = make_tma_2d_desc(b, k, s * n, block_k, block_n, k, swizzle_ab_mode);
const auto& tensor_map_d = make_tma_2d_desc(d, n, m, block_n, block_m, n, swizzle_cd_mode);
const auto tensor_map_a = make_tma_2d_desc(a, k, s * m, block_k, block_m, k, swizzle_ab_mode);
const auto tensor_map_b = make_tma_2d_desc(b, k, s * n, block_k, block_n, k, swizzle_ab_mode);
const auto tensor_map_d = make_tma_2d_desc(d, n, m, block_n, block_m, n, swizzle_cd_mode);
const SM100BmkBnkMnRuntime::Args& args = {
const SM100BmkBnkMnRuntime::Args args = {
.s = s, .m = m, .n = n, .k = k,
.block_m = block_m, .block_n = block_n, .block_k = block_k,
.split_factor = split_factor,
@@ -129,8 +129,8 @@ static void sm100_bmn_bnk_mn_gemm(const torch::Tensor &a,
.tensor_map_b = tensor_map_b,
.tensor_map_d = tensor_map_d
};
const auto& code = SM100BmkBnkMnRuntime::generate(args);
const auto& runtime = compiler->build("sm100_bmn_bnk_mn_gemm", code);
const auto code = SM100BmkBnkMnRuntime::generate(args);
const auto runtime = compiler->build("sm100_bmn_bnk_mn_gemm", code);
SM100BmkBnkMnRuntime::launch(runtime, args);
}

459
csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp Normal file
View File

@@ -0,0 +1,459 @@
#pragma once
#include <torch/python.h>
#include "../../jit/compiler.hpp"
#include "../../jit/device_runtime.hpp"
#include "../../jit/kernel_runtime.hpp"
#include "../../utils/exception.hpp"
#include "../../utils/format.hpp"
#include "../../utils/math.hpp"
#include "../heuristics/sm100.hpp"
#include "epilogue.hpp"
#include "runtime_utils.hpp"
namespace deep_gemm {
class SM100FP8FP4Gemm1D1DRuntime final: public LaunchRuntime<SM100FP8FP4Gemm1D1DRuntime> {
public:
struct Args {
GemmDesc gemm_desc;
GemmConfig gemm_config;
LaunchArgs launch_args;
// TODO: move into descriptor
const std::optional<std::string> epilogue_type;
// TODO: move into descriptor
int gran_k_a, gran_k_b;
void* grouped_layout;
CUtensorMap tensor_map_a;
CUtensorMap tensor_map_b;
CUtensorMap tensor_map_sfa;
CUtensorMap tensor_map_sfb;
CUtensorMap tensor_map_cd;
};
static std::string generate_impl(const Args& args) {
// TODO: rename files
return fmt::format(R"(
#include <deep_gemm/impls/sm100_fp8_fp4_gemm_1d1d.cuh>
using namespace deep_gemm;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&sm100_fp8_fp4_gemm_1d1d_impl<
{}, {},
{}, {},
{}, {}, {},
{}, {}, {},
{},
{}, {}, {},
{},
{}, {},
{}, {},
{},
{},
{}, {},
{}, {}, {},
{}
>);
}};
)",
to_string(args.gemm_desc.major_a), to_string(args.gemm_desc.major_b),
args.gran_k_a, args.gran_k_b,
get_compiled_dim(args.gemm_desc.m, 'm', args.gemm_desc.compiled_dims),
get_compiled_dim(args.gemm_desc.n, 'n', args.gemm_desc.compiled_dims),
get_compiled_dim(args.gemm_desc.k, 'k', args.gemm_desc.compiled_dims),
args.gemm_config.layout.block_m, args.gemm_config.layout.block_n, args.gemm_config.layout.block_k,
args.gemm_desc.num_groups,
args.gemm_config.storage_config.swizzle_a_mode, args.gemm_config.storage_config.swizzle_b_mode, args.gemm_config.storage_config.swizzle_cd_mode,
args.gemm_config.pipeline_config.num_stages,
args.gemm_config.launch_config.num_non_epilogue_threads, args.gemm_config.launch_config.num_epilogue_threads,
args.gemm_config.layout.get_cluster_size(), args.gemm_config.layout.cluster_n > 1,
args.gemm_config.launch_config.num_sms,
args.gemm_config.layout.swap_ab,
to_string(args.gemm_desc.gemm_type), args.gemm_desc.with_accumulation,
to_string(args.gemm_desc.a_dtype), to_string(args.gemm_desc.b_dtype), to_string(args.gemm_desc.cd_dtype),
get_default_epilogue_type(args.epilogue_type));
}
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
// TODO: optimize `args` copy
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.grouped_layout, args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k,
args.tensor_map_a, args.tensor_map_b,
args.tensor_map_sfa, args.tensor_map_sfb,
args.tensor_map_cd));
}
};
static void sm100_fp8_fp4_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa,
const torch::Tensor& b, const torch::Tensor& sfb,
const std::optional<torch::Tensor>& c,
const torch::Tensor& d,
const int& m, const int& n, const int& k,
const int& gran_k_a, const int& gran_k_b,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims,
const std::optional<std::string>& epilogue_type = std::nullopt) {
const auto desc = GemmDesc {
.gemm_type = GemmType::Normal,
.kernel_type = KernelType::Kernel1D1D,
.m = m, .n = n, .k = k, .num_groups = 1,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = c.has_value(),
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(),
.compiled_dims = compiled_dims
};
const auto config = get_best_config<SM100ArchSpec>(desc);
const auto cd = c.value_or(d);
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), 1,
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_cd_desc(d, m, static_cast<int>(d.size(-1)),
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), 1,
config.storage_config.swizzle_cd_mode);
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.layout.block_m, gran_k_a, 1, 0);
const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k,
config.layout.block_n, gran_k_b, 1, 0);
// Launch
const SM100FP8FP4Gemm1D1DRuntime::Args args = {
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.epilogue_type = epilogue_type,
.gran_k_a = gran_k_a,
.gran_k_b = gran_k_b,
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_sfa = tensor_map_sfa,
.tensor_map_sfb = tensor_map_sfb,
.tensor_map_cd = tensor_map_cd
};
const auto code = SM100FP8FP4Gemm1D1DRuntime::generate(args);
const auto runtime = compiler->build("sm100_fp8_fp4_gemm_1d1d", code);
SM100FP8FP4Gemm1D1DRuntime::launch(runtime, args);
}
static void sm100_m_grouped_fp8_fp4_gemm_contiguous_1d1d(const torch::Tensor& a, const torch::Tensor& sfa,
const torch::Tensor& b, const torch::Tensor& sfb,
const torch::Tensor& d,
const torch::Tensor& grouped_layout,
const int& num_groups, const int& m, const int& n, const int& k,
const int& gran_k_a, const int& gran_k_b,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims,
const bool& use_psum_layout,
const std::optional<int>& expected_m_for_psum_layout) {
const auto gemm_type = use_psum_layout ?
GemmType::MGroupedContiguousWithPsumLayout : GemmType::MGroupedContiguous;
// Only psum layout can use expected m
if (expected_m_for_psum_layout)
DG_HOST_ASSERT(use_psum_layout);
// NOTES: If actual M is dynamic, estimate config via `num_groups` and `expected_m`.
// Otherwise, treat the contiguous layout as a whole.
const auto desc = GemmDesc {
.gemm_type = gemm_type,
.kernel_type = KernelType::Kernel1D1D,
.m = m, .n = n, .k = k, .num_groups = num_groups,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = false,
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(),
.compiled_dims = compiled_dims,
.expected_m = expected_m_for_psum_layout.value_or(m),
.expected_n = n, .expected_k = k,
.expected_num_groups = expected_m_for_psum_layout.has_value() ? num_groups : 1
};
const auto config = get_best_config<SM100ArchSpec>(desc);
// Create tensor descriptors
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), 1,
config.storage_config.swizzle_cd_mode);
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.layout.block_m, gran_k_a, 1, 0);
const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k,
config.layout.block_n, gran_k_b, num_groups, 0);
// Launch kernel
const SM100FP8FP4Gemm1D1DRuntime::Args args = {
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.epilogue_type = std::nullopt,
.gran_k_a = gran_k_a,
.gran_k_b = gran_k_b,
.grouped_layout = grouped_layout.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_sfa = tensor_map_sfa,
.tensor_map_sfb = tensor_map_sfb,
.tensor_map_cd = tensor_map_cd
};
const auto code = SM100FP8FP4Gemm1D1DRuntime::generate(args);
const auto runtime = compiler->build("sm100_m_grouped_fp8_fp4_gemm_contiguous_1d1d", code);
SM100FP8FP4Gemm1D1DRuntime::launch(runtime, args);
}
static void sm100_m_grouped_fp8_fp4_gemm_masked_1d1d(const torch::Tensor& a, const torch::Tensor& sfa,
const torch::Tensor& b, const torch::Tensor& sfb,
const torch::Tensor& d,
const torch::Tensor& masked_m,
const int& num_groups, const int& m, const int& n, const int& k,
const int& expected_m,
const int& gran_k_a, const int& gran_k_b,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims) {
const auto desc = GemmDesc {
.gemm_type = GemmType::MGroupedMasked,
.kernel_type = KernelType::Kernel1D1D,
.m = m, .n = n, .k = k, .num_groups = num_groups,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = false,
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(),
.compiled_dims = compiled_dims,
.expected_m = expected_m, .expected_n = n, .expected_k = k, .expected_num_groups = num_groups
};
const auto config = get_best_config<SM100ArchSpec>(desc);
// Create tensor descriptors
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), num_groups,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), num_groups,
config.storage_config.swizzle_cd_mode);
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.layout.block_m, gran_k_a, num_groups, 0);
const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k,
config.layout.block_n, gran_k_b, num_groups, 0);
// Launch kernel
const SM100FP8FP4Gemm1D1DRuntime::Args args = {
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.epilogue_type = std::nullopt,
.gran_k_a = gran_k_a,
.gran_k_b = gran_k_b,
.grouped_layout = masked_m.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_sfa = tensor_map_sfa,
.tensor_map_sfb = tensor_map_sfb,
.tensor_map_cd = tensor_map_cd
};
const auto code = SM100FP8FP4Gemm1D1DRuntime::generate(args);
const auto runtime = compiler->build("sm100_m_grouped_fp8_fp4_gemm_masked_1d1d", code);
SM100FP8FP4Gemm1D1DRuntime::launch(runtime, args);
}
static void sm100_k_grouped_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa,
const torch::Tensor& b, const torch::Tensor& sfb,
const std::optional<torch::Tensor>& c,
const torch::Tensor& d,
const int& m, const int& n,
const std::vector<int>& ks, const torch::Tensor& ks_tensor,
const int& gran_k,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims) {
DG_HOST_ASSERT(major_a == cute::UMMA::Major::MN and major_b == cute::UMMA::Major::MN);
DG_HOST_ASSERT(gran_k == 32 or gran_k == 128);
const int gran_k_a = gran_k;
const int gran_k_b = gran_k;
int sum_k = 0, sum_sf_k = 0;
for (const auto k: ks) {
sum_k += k, sum_sf_k += ceil_div(k, gran_k * 4);
DG_HOST_ASSERT(k % gran_k == 0);
}
const auto num_groups = static_cast<int>(ks.size());
// Get config using max K for better performance
const auto max_k = *std::max_element(ks.begin(), ks.end());
const auto desc = GemmDesc {
.gemm_type = GemmType::KGroupedContiguous,
.kernel_type = KernelType::Kernel1D1D,
.m = m, .n = n, .k = sum_k, .num_groups = num_groups,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = c.has_value(),
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(),
.compiled_dims = compiled_dims,
.expected_m = m, .expected_n = n, .expected_k = max_k, .expected_num_groups = num_groups
};
const auto config = get_best_config<SM100ArchSpec>(desc);
// Create tensor descriptors
const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::MN, a, m, sum_k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(0)), 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(cute::UMMA::Major::MN, b, n, sum_k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(0)), 1,
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(1)), num_groups,
config.storage_config.swizzle_cd_mode);
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, sum_sf_k * gran_k_a * 4,
config.layout.block_m, gran_k_a, 1, 0);
const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, sum_sf_k * gran_k_b * 4,
config.layout.block_n, gran_k_b, 1, 0);
// Launch kernel
const SM100FP8FP4Gemm1D1DRuntime::Args args = {
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.epilogue_type = std::nullopt,
.gran_k_a = gran_k_a,
.gran_k_b = gran_k_b,
.grouped_layout = ks_tensor.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_sfa = tensor_map_sfa,
.tensor_map_sfb = tensor_map_sfb,
.tensor_map_cd = tensor_map_cd
};
const auto code = SM100FP8FP4Gemm1D1DRuntime::generate(args);
const auto runtime = compiler->build("sm100_k_grouped_fp8_gemm_1d1d", code);
SM100FP8FP4Gemm1D1DRuntime::launch(runtime, args);
}
static void sm100_fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa,
const torch::Tensor& b, const torch::Tensor& sfb,
const std::optional<torch::Tensor>& c,
const torch::Tensor& d,
const int& batch_size, const int& m, const int& n, const int& k,
const int& gran_k_a, const int& gran_k_b,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims) {
const auto desc = GemmDesc {
.gemm_type = GemmType::Batched,
.kernel_type = KernelType::Kernel1D1D,
.m = m, .n = n, .k = k, .num_groups = batch_size,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = c.has_value(),
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(),
.compiled_dims = compiled_dims
};
const auto config = get_best_config<SM100ArchSpec>(desc);
const int load_block_m = config.storage_config.load_block_m;
const auto [inner_dim_a, outer_dim_a] = get_inner_outer_dims(major_a, k, m);
const auto [inner_block_a, outer_block_a] = get_inner_outer_dims(major_a, config.layout.block_k, load_block_m);
const auto tensor_map_a = make_tma_3d_desc(a, inner_dim_a, outer_dim_a, batch_size,
inner_block_a, outer_block_a, 1,
a.stride(major_a == cute::UMMA::Major::K ? 1 : 2),
a.stride(0),
config.storage_config.swizzle_a_mode);
const int load_block_n = config.storage_config.load_block_n;
const auto [inner_dim_b, outer_dim_b] = get_inner_outer_dims(major_b, k, n);
const auto [inner_block_b, outer_block_b] = get_inner_outer_dims(major_b, config.layout.block_k, load_block_n);
const auto tensor_map_b = make_tma_3d_desc(b, inner_dim_b, outer_dim_b, batch_size,
inner_block_b, outer_block_b, 1,
b.stride(major_b == cute::UMMA::Major::K ? 1 : 2),
b.stride(0),
config.storage_config.swizzle_b_mode);
const int store_block_m = config.storage_config.store_block_m;
const int store_block_n = config.storage_config.store_block_n;
const auto tensor_map_cd = make_tma_3d_desc(d, n, m, batch_size,
store_block_n, store_block_m, 1,
d.stride(1), d.stride(0),
config.storage_config.swizzle_cd_mode);
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.layout.block_m, gran_k_a, batch_size, 0);
const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k,
config.layout.block_n, gran_k_b, batch_size, 0);
// Launch
const SM100FP8FP4Gemm1D1DRuntime::Args args = {
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.epilogue_type = std::nullopt,
.gran_k_a = gran_k_a,
.gran_k_b = gran_k_b,
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_sfa = tensor_map_sfa,
.tensor_map_sfb = tensor_map_sfb,
.tensor_map_cd = tensor_map_cd
};
const auto code = SM100FP8FP4Gemm1D1DRuntime::generate(args);
const auto runtime = compiler->build("sm100_fp8_gemm_1d1d", code);
SM100FP8FP4Gemm1D1DRuntime::launch(runtime, args);
}
} // namespace deep_gemm

210
csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp Normal file
View File

@@ -0,0 +1,210 @@
#pragma once
#include <torch/python.h>
#include "../../jit/compiler.hpp"
#include "../../jit/kernel_runtime.hpp"
#include "../../utils/exception.hpp"
#include "../../utils/format.hpp"
#include "runtime_utils.hpp"
#include <deep_gemm/layout/mega_moe.cuh>
#include <deep_gemm/layout/sym_buffer.cuh>
#include "../heuristics/mega_moe.hpp"
namespace deep_gemm {
class SM100FP8FP4MegaMoERuntime final : public LaunchRuntime<SM100FP8FP4MegaMoERuntime> {
public:
struct Args {
// Templated arguments
int num_max_tokens_per_rank;
int hidden, intermediate_hidden;
int num_experts, num_topk;
int num_ranks;
float activation_clamp;
bool fast_math;
MegaMoEConfig config;
// Runtime arguments
void* y;
int num_tokens;
layout::SymBuffer<> sym_buffer_ptrs;
// Tensormap
CUtensorMap tensor_map_l1_acts;
CUtensorMap tensor_map_l1_acts_sf;
CUtensorMap tensor_map_l1_weights;
CUtensorMap tensor_map_l1_weights_sf;
CUtensorMap tensor_map_l1_output;
CUtensorMap tensor_map_l2_acts;
CUtensorMap tensor_map_l2_acts_sf;
CUtensorMap tensor_map_l2_weights;
CUtensorMap tensor_map_l2_weights_sf;
// Launch configs
LaunchArgs launch_args;
};
static std::string generate_impl(const Args& args) {
return fmt::format(R"(
#include <deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh>
using namespace deep_gemm;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&sm100_fp8_fp4_mega_moe_impl<
{},
{}, {},
{}, {},
{},
{}, {}, {},
{},
{}, {},
{},
{},
{},
{}, {}, {},
{}, {},
{},
{}
>);
}};
)", args.num_max_tokens_per_rank,
args.hidden, args.intermediate_hidden,
args.num_experts, args.num_topk,
args.config.num_experts_per_wave,
args.config.block_m, args.config.block_n, args.config.block_k,
args.config.store_block_m,
args.config.sf_block_m, args.config.sf_block_n,
args.config.num_max_pool_tokens,
args.config.num_padded_sf_pool_tokens,
args.config.num_stages,
args.config.num_dispatch_threads, args.config.num_non_epilogue_threads, args.config.num_epilogue_threads,
args.launch_args.grid_dim.first, args.num_ranks,
to_string(args.activation_clamp),
args.fast_math ? "true" : "false");
}
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
// TODO: optimize `args` copy
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.y,
args.num_tokens,
args.sym_buffer_ptrs,
args.tensor_map_l1_acts,
args.tensor_map_l1_acts_sf,
args.tensor_map_l1_weights,
args.tensor_map_l1_weights_sf,
args.tensor_map_l1_output,
args.tensor_map_l2_acts,
args.tensor_map_l2_acts_sf,
args.tensor_map_l2_weights,
args.tensor_map_l2_weights_sf
));
}
};
static void sm100_fp8_fp4_mega_moe(
const torch::Tensor& y,
const torch::Tensor& l1_acts, const torch::Tensor& l1_acts_sf,
const torch::Tensor& l2_acts, const torch::Tensor& l2_acts_sf,
const torch::Tensor& l1_weights, const torch::Tensor& l2_weights,
const torch::Tensor& l1_weights_sf, const torch::Tensor& l2_weights_sf,
const std::vector<int64_t>& sym_buffer_ptrs,
const int& rank_idx, const int& num_max_tokens_per_rank,
const int& num_experts_per_rank,
const int& num_tokens, const int& num_topk,
const int& hidden, const int& intermediate_hidden,
const float& activation_clamp,
const bool& fast_math
) {
const auto num_ranks = static_cast<int>(sym_buffer_ptrs.size());
const auto num_experts = num_experts_per_rank * num_ranks;
// Heuristics
const auto config = get_mega_moe_config(
num_ranks, num_experts, num_experts_per_rank,
num_max_tokens_per_rank, num_tokens, num_topk, hidden, intermediate_hidden);
// Make tensormap
constexpr int kGranK = 32;
const auto tensor_map_l1_acts = make_tma_2d_desc(l1_acts,
hidden, config.num_max_pool_tokens,
config.block_k, config.load_block_m,
static_cast<int>(l1_acts.stride(-2)),
config.swizzle_acts_mode);
const auto tensor_map_l1_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l1_acts_sf,
config.num_padded_sf_pool_tokens, hidden,
config.sf_block_m, kGranK,
1, 0);
const auto tensor_map_l1_weights = make_tma_2d_desc(l1_weights,
hidden, num_experts_per_rank * intermediate_hidden * 2,
config.block_k, config.load_block_n,
static_cast<int>(l1_weights.stride(-2)),
config.swizzle_weights_mode);
const auto tensor_map_l1_weights_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l1_weights_sf,
intermediate_hidden * 2, hidden,
config.block_n, kGranK,
num_experts_per_rank, 0);
// NOTES: L1 output and L2 activations are essentially the same tensor.
// Post-SwiGLU output has half the N width (`BLOCK_N / 2` per input tile),
// so the swizzle mode is also halved (128 -> 64).
const auto tensor_map_l1_output = make_tma_2d_desc(l2_acts,
intermediate_hidden, config.num_max_pool_tokens,
config.block_n / 2, config.store_block_m,
static_cast<int>(l2_acts.stride(-2)),
config.swizzle_acts_mode / 2);
const auto tensor_map_l2_acts = make_tma_2d_desc(l2_acts,
intermediate_hidden, config.num_max_pool_tokens,
config.block_k, config.load_block_m,
static_cast<int>(l2_acts.stride(-2)),
config.swizzle_acts_mode);
const auto tensor_map_l2_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l2_acts_sf,
config.num_padded_sf_pool_tokens, intermediate_hidden,
config.sf_block_m, kGranK,
1, 0);
const auto tensor_map_l2_weights = make_tma_2d_desc(l2_weights,
intermediate_hidden, num_experts_per_rank * hidden,
config.block_k, config.load_block_n,
static_cast<int>(l2_weights.stride(-2)),
config.swizzle_weights_mode);
const auto tensor_map_l2_weights_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l2_weights_sf,
hidden, intermediate_hidden,
config.block_n, kGranK,
num_experts_per_rank, 0);
// Launch
const auto num_sms = device_runtime->get_num_sms();
const SM100FP8FP4MegaMoERuntime::Args args = {
.num_max_tokens_per_rank = num_max_tokens_per_rank,
.hidden = hidden, .intermediate_hidden = intermediate_hidden,
.num_experts = num_experts, .num_topk = num_topk,
.num_ranks = num_ranks,
.activation_clamp = activation_clamp,
.fast_math = fast_math,
.config = config,
.y = y.data_ptr(),
.num_tokens = num_tokens,
.sym_buffer_ptrs = layout::SymBuffer<>(sym_buffer_ptrs, rank_idx),
.tensor_map_l1_acts = tensor_map_l1_acts,
.tensor_map_l1_acts_sf = tensor_map_l1_acts_sf,
.tensor_map_l1_weights = tensor_map_l1_weights,
.tensor_map_l1_weights_sf = tensor_map_l1_weights_sf,
.tensor_map_l1_output = tensor_map_l1_output,
.tensor_map_l2_acts = tensor_map_l2_acts,
.tensor_map_l2_acts_sf = tensor_map_l2_acts_sf,
.tensor_map_l2_weights = tensor_map_l2_weights,
.tensor_map_l2_weights_sf = tensor_map_l2_weights_sf,
.launch_args = LaunchArgs(num_sms,
config.num_dispatch_threads + config.num_non_epilogue_threads + config.num_epilogue_threads,
config.smem_size, 2)
};
const auto code = SM100FP8FP4MegaMoERuntime::generate(args);
const auto runtime = compiler->build("sm100_fp8_fp4_mega_moe", code);
SM100FP8FP4MegaMoERuntime::launch(runtime, args);
}
} // namespace deep_gemm

36
csrc/jit_kernels/impls/sm100_tf32_hc_prenorm_gemm.hpp
View File

@@ -79,21 +79,21 @@ static void sm100_tf32_hc_prenorm_gemm(const torch::Tensor& a,
DG_HOST_ASSERT(n <= 128 and n % 8 == 0);
DG_HOST_ASSERT(k % block_k == 0);
const auto& swizzle_cd_mode = get_swizzle_mode(block_n, sizeof(float));
const auto& tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a, m, k,
block_m, block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(cute::UMMA::Major::K))), 1,
get_swizzle_mode(block_k, a.element_size()), 0,
true);
const auto& tensor_map_b = make_tma_b_desc(cute::UMMA::Major::K, b, n, k,
block_n, block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(cute::UMMA::Major::K))), 1,
get_swizzle_mode(block_k, b.element_size()), 0,
true);
const auto& tensor_map_d = num_splits == 1 ? make_tma_cd_desc(d, m, n,
block_m, block_n,
static_cast<int>(d.stride(-2)), 1,
swizzle_cd_mode)
const auto swizzle_cd_mode = get_swizzle_mode(block_n, sizeof(float));
const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a, m, k,
block_m, block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(cute::UMMA::Major::K))), 1,
get_swizzle_mode(block_k, a.element_size()), 0,
true);
const auto tensor_map_b = make_tma_b_desc(cute::UMMA::Major::K, b, n, k,
block_n, block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(cute::UMMA::Major::K))), 1,
get_swizzle_mode(block_k, b.element_size()), 0,
true);
const auto tensor_map_d = num_splits == 1 ? make_tma_cd_desc(d, m, n,
block_m, block_n,
static_cast<int>(d.stride(-2)), 1,
swizzle_cd_mode)
: make_tma_3d_desc(d, n, m, num_splits,
block_n, block_m, 1,
static_cast<int>(d.stride(-2)),
@@ -135,14 +135,14 @@ static void sm100_tf32_hc_prenorm_gemm(const torch::Tensor& a,
.num_stages = num_stages,
.num_mma_threads = num_mma_threads,
.num_cast_and_reduce_threads = num_cast_and_reduce_threads,
.launch_args = LaunchArgs(num_splits * ceil_div(m, block_m), num_mma_threads + num_cast_and_reduce_threads, smem_size, 1),
.launch_args = LaunchArgs(num_splits * ceil_div(m, block_m), num_mma_threads + num_cast_and_reduce_threads, smem_size),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_d = tensor_map_d,
.sqr_sum = sqr_sum.data_ptr<float>()
};
const auto& code = SM100BF16HCPrenormGemmRuntime::generate(args);
const auto& runtime = compiler->build("sm100_tf32_hc_prenorm_gemm", code);
const auto code = SM100BF16HCPrenormGemmRuntime::generate(args);
const auto runtime = compiler->build("sm100_tf32_hc_prenorm_gemm", code);
SM100BF16HCPrenormGemmRuntime::launch(runtime, args);
}

418
csrc/jit_kernels/impls/sm90_bf16_gemm.hpp
View File

@@ -14,9 +14,7 @@ namespace deep_gemm {
class SM90BF16GemmRuntime final: public LaunchRuntime<SM90BF16GemmRuntime> {
public:
struct Args {
int m, n, k, num_groups;
const std::string& compiled_dims;
GemmDesc gemm_desc;
GemmConfig gemm_config;
LaunchArgs launch_args;
@@ -49,24 +47,29 @@ static void __instantiate_kernel() {{
}};
)",
// TODO: add CD dtype
to_string(args.gemm_config.major_a), to_string(args.gemm_config.major_b),
get_compiled_dim(args.m, 'm', args.compiled_dims), get_compiled_dim(args.n, 'n', args.compiled_dims), get_compiled_dim(args.k, 'k', args.compiled_dims),
args.num_groups,
args.gemm_config.block_m, args.gemm_config.block_n, args.gemm_config.block_k,
args.gemm_config.smem_config.swizzle_a_mode, args.gemm_config.smem_config.swizzle_b_mode, args.gemm_config.smem_config.swizzle_cd_mode,
args.gemm_config.num_stages,
args.gemm_config.thread_config.num_tma_threads, args.gemm_config.thread_config.num_math_threads,
args.gemm_config.multicast_config.num_multicast, args.gemm_config.multicast_config.is_multicast_on_a,
args.gemm_config.num_sms,
to_string(args.gemm_config.gemm_type), args.gemm_config.with_accumulation,
to_string(args.gemm_config.cd_dtype));
to_string(args.gemm_desc.major_a), to_string(args.gemm_desc.major_b),
get_compiled_dim(args.gemm_desc.m, 'm', args.gemm_desc.compiled_dims),
get_compiled_dim(args.gemm_desc.n, 'n', args.gemm_desc.compiled_dims),
get_compiled_dim(args.gemm_desc.k, 'k', args.gemm_desc.compiled_dims),
args.gemm_desc.num_groups,
args.gemm_config.layout.block_m, args.gemm_config.layout.block_n, args.gemm_config.layout.block_k,
args.gemm_config.storage_config.swizzle_a_mode,
args.gemm_config.storage_config.swizzle_b_mode,
args.gemm_config.storage_config.swizzle_cd_mode,
args.gemm_config.pipeline_config.num_stages,
args.gemm_config.launch_config.num_tma_threads, args.gemm_config.launch_config.num_math_threads,
// TODO: refactor with cluster M/N
args.gemm_config.layout.get_cluster_size(), args.gemm_config.layout.cluster_n > 1,
args.gemm_config.launch_config.num_sms,
to_string(args.gemm_desc.gemm_type), args.gemm_desc.with_accumulation,
to_string(args.gemm_desc.cd_dtype));
}
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
// TODO: optimize `args` copy
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.grouped_layout,
args.m, args.n, args.k,
args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k,
args.tensor_map_a, args.tensor_map_b,
args.tensor_map_cd));
}
@@ -79,46 +82,50 @@ static void sm90_bf16_gemm(const torch::Tensor& a,
const int& m, const int& n, const int& k,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims) {
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::Normal, KernelType::KernelNoSF,
m, n, k, 1, major_a, major_b,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), c.has_value(),
device_runtime->get_num_sms());
const auto desc = GemmDesc {
.gemm_type = GemmType::Normal,
.kernel_type = KernelType::KernelNoSF,
.m = m, .n = n, .k = k, .num_groups = 1,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = c.has_value(),
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims
};
const auto config = get_best_config<SM90ArchSpec>(desc);
// Requires no TMA splits
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(major_b, b, n, k,
SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), 1,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM90ArchSpec::get_cd_store_block_m(config.block_m),
SM90ArchSpec::get_cd_store_block_n(config.block_n),
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), 1,
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), 1,
config.smem_config.swizzle_cd_mode);
config.storage_config.swizzle_cd_mode);
// Launch
const SM90BF16GemmRuntime::Args& args = {
.m = m, .n = n, .k = k,
.num_groups = 1,
.compiled_dims = compiled_dims,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd,
};
const auto& code = SM90BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm90_bf16_gemm", code);
const auto code = SM90BF16GemmRuntime::generate(args);
const auto runtime = compiler->build("sm90_bf16_gemm", code);
SM90BF16GemmRuntime::launch(runtime, args);
}
@@ -128,51 +135,67 @@ static void sm90_m_grouped_bf16_gemm_contiguous(const torch::Tensor& a,
const torch::Tensor& m_indices,
const int& num_groups, const int& m, const int& n, const int& k,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims) {
const std::string& compiled_dims,
const bool& use_psum_layout,
const std::optional<int>& expected_m_for_psum_layout) {
DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16);
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K);
DG_HOST_ASSERT(k % 64 == 0);
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::MGroupedContiguous, KernelType::KernelNoSF,
m, n, k, 1, major_a, major_b,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), false,
device_runtime->get_num_sms());
const auto gemm_type = use_psum_layout ?
GemmType::MGroupedContiguousWithPsumLayout : GemmType::MGroupedContiguous;
// Only psum layout can use expected m
if (expected_m_for_psum_layout)
DG_HOST_ASSERT(use_psum_layout);
const auto desc = GemmDesc {
.gemm_type = gemm_type,
.kernel_type = KernelType::KernelNoSF,
.m = m, .n = n, .k = k, .num_groups = num_groups,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = false,
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims,
.expected_m = expected_m_for_psum_layout.value_or(m),
.expected_n = n, .expected_k = k,
.expected_num_groups = expected_m_for_psum_layout.has_value() ? num_groups : 1
};
const auto config = get_best_config<SM90ArchSpec>(desc);
// Requires no TMA splits
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(major_b, b, n, k,
SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM90ArchSpec::get_cd_store_block_m(config.block_m),
SM90ArchSpec::get_cd_store_block_n(config.block_n),
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), 1,
config.smem_config.swizzle_cd_mode);
config.storage_config.swizzle_cd_mode);
// Launch
const SM90BF16GemmRuntime::Args& args = {
.m = m, .n = n, .k = k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.grouped_layout = m_indices.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd,
};
const auto& code = SM90BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm90_m_grouped_bf16_gemm_contiguous", code);
const auto code = SM90BF16GemmRuntime::generate(args);
const auto runtime = compiler->build("sm90_m_grouped_bf16_gemm_contiguous", code);
SM90BF16GemmRuntime::launch(runtime, args);
}
@@ -188,46 +211,51 @@ static void sm90_bf16_m_grouped_gemm_masked(const torch::Tensor& a,
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K);
DG_HOST_ASSERT(k % 64 == 0);
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::MGroupedMasked, KernelType::KernelNoSF,
expected_m, n, k, num_groups, major_a, major_b,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), false,
device_runtime->get_num_sms());
const auto desc = GemmDesc {
.gemm_type = GemmType::MGroupedMasked,
.kernel_type = KernelType::KernelNoSF,
.m = m, .n = n, .k = k, .num_groups = num_groups,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = false,
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims,
.expected_m = expected_m, .expected_n = 0, .expected_k = 0, .expected_num_groups = num_groups
};
const auto config = get_best_config<SM90ArchSpec>(desc);
// Requires no TMA splits
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), num_groups,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(major_b, b, n, k,
SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM90ArchSpec::get_cd_store_block_m(config.block_m),
SM90ArchSpec::get_cd_store_block_n(config.block_n),
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), num_groups,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), num_groups,
config.smem_config.swizzle_cd_mode);
config.storage_config.swizzle_cd_mode);
// Launch
const SM90BF16GemmRuntime::Args& args = {
.m = m, .n = n, .k = k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.grouped_layout = masked_m.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd,
};
const auto& code = SM90BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm90_bf16_m_grouped_gemm_masked", code);
const auto code = SM90BF16GemmRuntime::generate(args);
const auto runtime = compiler->build("sm90_bf16_m_grouped_gemm_masked", code);
SM90BF16GemmRuntime::launch(runtime, args);
}
@@ -242,54 +270,59 @@ static void sm90_bf16_k_grouped_gemm(const torch::Tensor& a,
DG_HOST_ASSERT(major_a == cute::UMMA::Major::MN and major_b == cute::UMMA::Major::MN);
int sum_k = 0;
for (const auto& k: ks) {
for (const auto k: ks) {
sum_k += k;
DG_HOST_ASSERT(k % 128 == 0);
}
const auto& num_groups = static_cast<int>(ks.size());
const auto num_groups = static_cast<int>(ks.size());
// Get config using max K for better performance
const auto& max_k = *std::max_element(ks.begin(), ks.end());
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::KGroupedContiguous, KernelType::KernelNoSF,
m, n, max_k, num_groups, cute::UMMA::Major::MN, cute::UMMA::Major::MN,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), c.has_value(),
device_runtime->get_num_sms());
const auto max_k = *std::max_element(ks.begin(), ks.end());
const auto desc = GemmDesc {
.gemm_type = GemmType::KGroupedContiguous,
.kernel_type = KernelType::KernelNoSF,
.m = m, .n = n, .k = sum_k, .num_groups = num_groups,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = c.has_value(),
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims,
.expected_m = m, .expected_n = n, .expected_k = max_k, .expected_num_groups = num_groups
};
const auto config = get_best_config<SM90ArchSpec>(desc);
// Create tensor descriptors
const auto& tensor_map_a = make_tma_a_desc(cute::UMMA::Major::MN, a, m, sum_k,
SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(0)), 1,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(cute::UMMA::Major::MN, b, n, sum_k,
SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(0)), 1,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM90ArchSpec::get_cd_store_block_m(config.block_m),
SM90ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(1)), num_groups,
config.smem_config.swizzle_cd_mode);
const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::MN, a, m, sum_k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(0)), 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(cute::UMMA::Major::MN, b, n, sum_k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(0)), 1,
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(1)), num_groups,
config.storage_config.swizzle_cd_mode);
// Launch kernel
const SM90BF16GemmRuntime::Args& args = {
.m = m, .n = n, .k = sum_k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.grouped_layout = ks_tensor.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd,
};
const auto& code = SM90BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm90_bf16_k_grouped_gemm", code);
const auto code = SM90BF16GemmRuntime::generate(args);
const auto runtime = compiler->build("sm90_bf16_k_grouped_gemm", code);
SM90BF16GemmRuntime::launch(runtime, args);
}
@@ -298,45 +331,50 @@ static void sm90_bf16_bhr_hdr_bhd(const torch::Tensor& tensor_a,
const torch::Tensor& tensor_d,
const int& b, const int& h, const int& r, const int& d,
const std::string& compiled_dims = "nk") {
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::Batched, KernelType::KernelNoSF,
b, d, r, h, cute::UMMA::Major::K, cute::UMMA::Major::K,
tensor_a.scalar_type(), tensor_b.scalar_type(),
tensor_d.scalar_type(), false,
device_runtime->get_num_sms());
const auto desc = GemmDesc {
.gemm_type = GemmType::Batched,
.kernel_type = KernelType::KernelNoSF,
.m = b, .n = d, .k = r, .num_groups = h,
.a_dtype = tensor_a.scalar_type(), .b_dtype = tensor_b.scalar_type(),
.cd_dtype = tensor_d.scalar_type(),
.major_a = cute::UMMA::Major::K, .major_b = cute::UMMA::Major::K,
.with_accumulation = false,
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims
};
const auto config = get_best_config<SM90ArchSpec>(desc);
const int& load_block_m = SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m);
const auto& tensor_map_a = make_tma_3d_desc(tensor_a, r, b, h,
config.block_k, load_block_m, 1,
tensor_a.stride(0), tensor_a.stride(1),
config.smem_config.swizzle_a_mode);
const int& load_block_n = SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n);
const auto& tensor_map_b = make_tma_3d_desc(tensor_b, r, d, h,
config.block_k, load_block_n, 1,
tensor_b.stride(1), tensor_b.stride(0),
config.smem_config.swizzle_b_mode);
const int& store_block_m = SM90ArchSpec::get_cd_store_block_m(config.block_m);
const int& store_block_n = SM90ArchSpec::get_cd_store_block_n(config.block_n);
const auto& tensor_map_cd = make_tma_3d_desc(tensor_d, d, b, h,
const int load_block_m = config.storage_config.load_block_m;
const auto tensor_map_a = make_tma_3d_desc(tensor_a, r, b, h,
config.layout.block_k, load_block_m, 1,
tensor_a.stride(0), tensor_a.stride(1),
config.storage_config.swizzle_a_mode);
const int load_block_n = config.storage_config.load_block_n;
const auto tensor_map_b = make_tma_3d_desc(tensor_b, r, d, h,
config.layout.block_k, load_block_n, 1,
tensor_b.stride(1), tensor_b.stride(0),
config.storage_config.swizzle_b_mode);
const int store_block_m = config.storage_config.store_block_m;
const int store_block_n = config.storage_config.store_block_n;
const auto tensor_map_cd = make_tma_3d_desc(tensor_d, d, b, h,
store_block_n, store_block_m, 1,
tensor_d.stride(0), tensor_d.stride(1),
config.smem_config.swizzle_cd_mode);
config.storage_config.swizzle_cd_mode);
// Launch
const SM90BF16GemmRuntime::Args& args = {
.m = b, .n = d, .k = r,
.num_groups = h,
.compiled_dims = compiled_dims,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd,
};
const auto& code = SM90BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm90_bf16_bhr_hdr_bhd", code);
const auto code = SM90BF16GemmRuntime::generate(args);
const auto runtime = compiler->build("sm90_bf16_bhr_hdr_bhd", code);
SM90BF16GemmRuntime::launch(runtime, args);
}
@@ -345,45 +383,49 @@ static void sm90_bf16_bhd_hdr_bhr(const torch::Tensor& tensor_a,
const torch::Tensor& tensor_d,
const int& b, const int& h, const int& r, const int& d,
const std::string& compiled_dims = "nk") {
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::Batched, KernelType::KernelNoSF,
b, r, d, h, cute::UMMA::Major::K, cute::UMMA::Major::MN,
tensor_a.scalar_type(), tensor_b.scalar_type(),
tensor_d.scalar_type(), false,
device_runtime->get_num_sms());
const auto desc = GemmDesc {
.gemm_type = GemmType::Batched,
.kernel_type = KernelType::KernelNoSF,
.m = b, .n = r, .k = d, .num_groups = h,
.a_dtype = tensor_a.scalar_type(), .b_dtype = tensor_b.scalar_type(),
.cd_dtype = tensor_d.scalar_type(),
.major_a = cute::UMMA::Major::K, .major_b = cute::UMMA::Major::MN,
.with_accumulation = false,
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims
};
const auto config = get_best_config<SM90ArchSpec>(desc);
const int& load_block_m = SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m);
const auto& tensor_map_a = make_tma_3d_desc(tensor_a, d, b, h,
config.block_k, load_block_m, 1,
tensor_a.stride(0), tensor_a.stride(1),
config.smem_config.swizzle_a_mode);
const int& load_block_n = SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n);
const auto& tensor_map_b = make_tma_3d_desc(tensor_b, r, d, h,
load_block_n, config.block_k, 1,
tensor_b.stride(1), tensor_b.stride(0),
config.smem_config.swizzle_b_mode);
const int& store_block_m = SM90ArchSpec::get_cd_store_block_m(config.block_m);
const int& store_block_n = SM90ArchSpec::get_cd_store_block_n(config.block_n);
const auto& tensor_map_cd = make_tma_3d_desc(tensor_d, r, b, h,
const int load_block_m = config.storage_config.load_block_m;
const auto tensor_map_a = make_tma_3d_desc(tensor_a, d, b, h,
config.layout.block_k, load_block_m, 1,
tensor_a.stride(0), tensor_a.stride(1),
config.storage_config.swizzle_a_mode);
const int load_block_n = config.storage_config.load_block_n;
const auto tensor_map_b = make_tma_3d_desc(tensor_b, r, d, h,
load_block_n, config.layout.block_k, 1,
tensor_b.stride(1), tensor_b.stride(0),
config.storage_config.swizzle_b_mode);
const int store_block_m = config.storage_config.store_block_m;
const int store_block_n = config.storage_config.store_block_n;
const auto tensor_map_cd = make_tma_3d_desc(tensor_d, r, b, h,
store_block_n, store_block_m, 1,
tensor_d.stride(0), tensor_d.stride(1),
config.smem_config.swizzle_cd_mode);
config.storage_config.swizzle_cd_mode);
// Launch
const SM90BF16GemmRuntime::Args& args = {
.m = b, .n = r, .k = d,
.num_groups = h,
.compiled_dims = compiled_dims,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd,
};
const auto& code = SM90BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm90_bf16_bhd_hdr_bhr", code);
const auto code = SM90BF16GemmRuntime::generate(args);
const auto runtime = compiler->build("sm90_bf16_bhd_hdr_bhr", code);
SM90BF16GemmRuntime::launch(runtime, args);
}

14
csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp
View File

@@ -84,9 +84,9 @@ static void sm90_bmn_bnk_mn_gemm(const torch::Tensor &a,
// Select best number of stages
int num_stages = 4, smem_size = 0;
while (true) {
const int& smem_a_per_stage = block_m * block_k * sizeof(cutlass::bfloat16_t);
const int& smem_b_per_stage = block_n * block_k * sizeof(cutlass::bfloat16_t);
const int& smem_barrier = SM90ArchSpec::get_barrier_smem_size(num_stages);
const int smem_a_per_stage = block_m * block_k * sizeof(cutlass::bfloat16_t);
const int smem_b_per_stage = block_n * block_k * sizeof(cutlass::bfloat16_t);
const int smem_barrier = num_stages * 8 * 2;
smem_size = 0;
smem_size += (smem_a_per_stage + smem_b_per_stage) * num_stages;
@@ -108,8 +108,8 @@ static void sm90_bmn_bnk_mn_gemm(const torch::Tensor &a,
num_stages, smem_size, swizzle_ab_mode);
}
const auto& tensor_map_a = make_tma_2d_desc(a, k, s * m, block_k, block_m, k, swizzle_ab_mode);
const auto& tensor_map_b = make_tma_2d_desc(b, k, s * n, block_k, block_n, k, swizzle_ab_mode);
const auto tensor_map_a = make_tma_2d_desc(a, k, s * m, block_k, block_m, k, swizzle_ab_mode);
const auto tensor_map_b = make_tma_2d_desc(b, k, s * n, block_k, block_n, k, swizzle_ab_mode);
const SM90BmkBnkMnRuntime::Args& args = {
.s = s, .m = m, .n = n, .k = k,
@@ -123,8 +123,8 @@ static void sm90_bmn_bnk_mn_gemm(const torch::Tensor &a,
.tensor_map_b = tensor_map_b,
.d = d.data_ptr<float>()
};
const auto& code = SM90BmkBnkMnRuntime::generate(args);
const auto& runtime = compiler->build("sm90_bmn_bnk_mn_gemm", code);
const auto code = SM90BmkBnkMnRuntime::generate(args);
const auto runtime = compiler->build("sm90_bmn_bnk_mn_gemm", code);
SM90BmkBnkMnRuntime::launch(runtime, args);
}

183
csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp
View File

@@ -15,9 +15,7 @@ namespace deep_gemm {
class SM90FP8Gemm1D1DRuntime final: public LaunchRuntime<SM90FP8Gemm1D1DRuntime> {
public:
struct Args {
int m, n, k, num_groups;
const std::string& compiled_dims;
GemmDesc gemm_desc;
GemmConfig gemm_config;
LaunchArgs launch_args;
@@ -52,15 +50,17 @@ static void __instantiate_kernel() {{
>);
}};
)",
get_compiled_dim(args.m, 'm', args.compiled_dims), get_compiled_dim(args.n, 'n', args.compiled_dims), get_compiled_dim(args.k, 'k', args.compiled_dims),
args.num_groups,
args.gemm_config.block_m, args.gemm_config.block_n, args.gemm_config.block_k,
args.gemm_config.smem_config.swizzle_a_mode, args.gemm_config.smem_config.swizzle_b_mode,
args.gemm_config.num_stages,
args.gemm_config.thread_config.num_tma_threads, args.gemm_config.thread_config.num_math_threads,
args.gemm_config.multicast_config.num_multicast, args.gemm_config.multicast_config.is_multicast_on_a,
args.gemm_config.num_sms, to_string(args.gemm_config.gemm_type),
to_string(args.gemm_config.cd_dtype));
get_compiled_dim(args.gemm_desc.m, 'm', args.gemm_desc.compiled_dims),
get_compiled_dim(args.gemm_desc.n, 'n', args.gemm_desc.compiled_dims),
get_compiled_dim(args.gemm_desc.k, 'k', args.gemm_desc.compiled_dims),
args.gemm_desc.num_groups,
args.gemm_config.layout.block_m, args.gemm_config.layout.block_n, args.gemm_config.layout.block_k,
args.gemm_config.storage_config.swizzle_a_mode, args.gemm_config.storage_config.swizzle_b_mode,
args.gemm_config.pipeline_config.num_stages,
args.gemm_config.launch_config.num_tma_threads, args.gemm_config.launch_config.num_math_threads,
args.gemm_config.layout.get_cluster_size(), args.gemm_config.layout.cluster_n > 1,
args.gemm_config.launch_config.num_sms, to_string(args.gemm_desc.gemm_type),
to_string(args.gemm_desc.cd_dtype));
}
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
@@ -68,7 +68,7 @@ static void __instantiate_kernel() {{
args.gmem_a_ptr, args.gmem_b_ptr,
args.grouped_layout,
args.tensor_map_buffer,
args.m, args.n, args.k,
args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k,
args.tensor_map_a_base, args.tensor_map_b_base,
args.tensor_map_sfa, args.tensor_map_sfb,
args.tensor_map_cd));
@@ -85,44 +85,48 @@ static void sm90_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa,
DG_HOST_ASSERT(c.has_value() and d.scalar_type() == torch::kFloat);
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K);
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::Normal, KernelType::Kernel1D1D,
m, n, k, 1, major_a, major_b,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), c.has_value(),
device_runtime->get_num_sms());
const auto desc = GemmDesc {
.gemm_type = GemmType::Normal,
.kernel_type = KernelType::Kernel1D1D,
.m = m, .n = n, .k = k, .num_groups = 1,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = c.has_value(),
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims
};
const auto config = get_best_config<SM90ArchSpec>(desc);
// Requires no TMA splits
DG_HOST_ASSERT(config.smem_config.swizzle_a_mode == config.block_k);
DG_HOST_ASSERT(config.smem_config.swizzle_b_mode == config.block_k);
DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k);
DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k);
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k, k, 1,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(major_b, b, n, k,
SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k, k, 1,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.block_m, config.block_k, 1, 0);
const auto& tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k,
config.block_n, config.block_k, 1, 0);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM90ArchSpec::get_cd_store_block_m(config.block_m, true),
SM90ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), 1,
0);
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k, k, 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k, k, 1,
config.storage_config.swizzle_b_mode);
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.layout.block_m, config.layout.block_k, 1, 0);
const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k,
config.layout.block_n, config.layout.block_k, 1, 0);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), 1,
0);
// Launch
const SM90FP8Gemm1D1DRuntime::Args& args = {
.m = m, .n = n, .k = k,
.num_groups = 1,
.compiled_dims = compiled_dims,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.gmem_a_ptr = nullptr,
.gmem_b_ptr = nullptr,
.grouped_layout = nullptr,
@@ -133,8 +137,8 @@ static void sm90_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa,
.tensor_map_sfb = tensor_map_sfb,
.tensor_map_cd = tensor_map_cd,
};
const auto& code = SM90FP8Gemm1D1DRuntime::generate(args);
const auto& runtime = compiler->build("sm90_fp8_gemm_1d1d", code);
const auto code = SM90FP8Gemm1D1DRuntime::generate(args);
const auto runtime = compiler->build("sm90_fp8_gemm_1d1d", code);
SM90FP8Gemm1D1DRuntime::launch(runtime, args);
}
@@ -151,54 +155,61 @@ static void sm90_k_grouped_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Te
DG_HOST_ASSERT(c.has_value() and d.scalar_type() == torch::kFloat);
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K);
// Get config using max K for better performance
const auto& num_groups = static_cast<int>(ks.size());
const auto& max_k = *std::max_element(ks.begin(), ks.end());
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::KGroupedContiguous, KernelType::Kernel1D1D,
m, n, max_k, num_groups, major_a, major_b,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), c.has_value(),
device_runtime->get_num_sms());
// Requires no TMA splits
DG_HOST_ASSERT(config.smem_config.swizzle_a_mode == config.block_k);
DG_HOST_ASSERT(config.smem_config.swizzle_b_mode == config.block_k);
int first_k = 0, sum_k = 0, sum_sf_k = 0;
// TODO: refactor with the mk alignment function
const auto num_groups = static_cast<int>(ks.size());
int first_k = 0, sum_k = 0, sum_sf_k = 0, max_k = 0;
for (int i = 0; i < num_groups; ++ i) {
if (first_k == 0 and ks[i] != 0)
first_k = ks[i];
sum_k += ks[i], sum_sf_k += ceil_div(ks[i], 128);
max_k = std::max(max_k, ks[i]);
DG_HOST_ASSERT(ks[i] % 128 == 0);
}
const auto& tensor_map_a_base = make_tma_a_desc(major_a, a, m, first_k,
SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k, first_k, 1,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b_base = make_tma_b_desc(major_b, b, n, first_k,
SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k, first_k, 1,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, sum_sf_k * 128,
config.block_m, config.block_k, 1, 0);
const auto& tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, sum_sf_k * 128,
config.block_n, config.block_k, 1, 0);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM90ArchSpec::get_cd_store_block_m(config.block_m, true),
SM90ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), num_groups,
config.smem_config.swizzle_cd_mode);
// Get config using max K for better performance
const auto desc = GemmDesc {
.gemm_type = GemmType::KGroupedContiguous,
.kernel_type = KernelType::Kernel1D1D,
.m = m, .n = n, .k = sum_k, .num_groups = num_groups,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = c.has_value(),
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims,
.expected_m = m, .expected_n = n, .expected_k = max_k, .expected_num_groups = num_groups
};
const auto config = get_best_config<SM90ArchSpec>(desc);
// Requires no TMA splits
DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k);
DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k);
const auto tensor_map_a_base = make_tma_a_desc(major_a, a, m, first_k,
config.storage_config.load_block_m,
config.layout.block_k, first_k, 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b_base = make_tma_b_desc(major_b, b, n, first_k,
config.storage_config.load_block_n,
config.layout.block_k, first_k, 1,
config.storage_config.swizzle_b_mode);
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, sum_sf_k * 128,
config.layout.block_m, config.layout.block_k, 1, 0);
const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, sum_sf_k * 128,
config.layout.block_n, config.layout.block_k, 1, 0);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), num_groups,
config.storage_config.swizzle_cd_mode);
// Launch
const SM90FP8Gemm1D1DRuntime::Args& args = {
.m = m, .n = n, .k = sum_k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.gmem_a_ptr = a.data_ptr(),
.gmem_b_ptr = b.data_ptr(),
.grouped_layout = ks_tensor.data_ptr(),
@@ -209,8 +220,8 @@ static void sm90_k_grouped_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Te
.tensor_map_sfb = tensor_map_sfb,
.tensor_map_cd = tensor_map_cd,
};
const auto& code = SM90FP8Gemm1D1DRuntime::generate(args);
const auto& runtime = compiler->build("sm90_fp8_gemm_1d1d", code);
const auto code = SM90FP8Gemm1D1DRuntime::generate(args);
const auto runtime = compiler->build("sm90_fp8_gemm_1d1d", code);
SM90FP8Gemm1D1DRuntime::launch(runtime, args);
}

348
csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp
View File

@@ -17,14 +17,13 @@ namespace deep_gemm {
class SM90FP8Gemm1D2DRuntime final: public LaunchRuntime<SM90FP8Gemm1D2DRuntime> {
public:
struct Args {
cute::UMMA::Major major_sfb;
int m, n, k, num_groups;
const std::string& compiled_dims;
const std::optional<std::string>& epilogue_type;
GemmDesc gemm_desc;
GemmConfig gemm_config;
LaunchArgs launch_args;
// TODO: move this into `gemm_desc`
const std::optional<std::string>& epilogue_type;
cute::UMMA::Major major_sfb;
void *sfb, *grouped_layout;
CUtensorMap tensor_map_a;
CUtensorMap tensor_map_b;
@@ -45,7 +44,7 @@ static void __instantiate_kernel() {{
{},
{}, {}, {},
{}, {}, {},
{}, {},
{},
{}, {},
{}, {},
{}, {},
@@ -55,14 +54,16 @@ static void __instantiate_kernel() {{
)",
// TODO: add CD dtype
to_string(args.major_sfb),
get_compiled_dim(args.m, 'm', args.compiled_dims), get_compiled_dim(args.n, 'n', args.compiled_dims), get_compiled_dim(args.k, 'k', args.compiled_dims),
args.num_groups,
args.gemm_config.block_m, args.gemm_config.block_n, args.gemm_config.block_k,
args.gemm_config.smem_config.swizzle_a_mode, args.gemm_config.smem_config.swizzle_b_mode, args.gemm_config.smem_config.swizzle_cd_mode,
args.gemm_config.num_stages, args.gemm_config.num_last_stages,
args.gemm_config.thread_config.num_tma_threads, args.gemm_config.thread_config.num_math_threads,
args.gemm_config.multicast_config.num_multicast, args.gemm_config.multicast_config.is_multicast_on_a,
args.gemm_config.num_sms, to_string(args.gemm_config.gemm_type),
get_compiled_dim(args.gemm_desc.m, 'm', args.gemm_desc.compiled_dims),
get_compiled_dim(args.gemm_desc.n, 'n', args.gemm_desc.compiled_dims),
get_compiled_dim(args.gemm_desc.k, 'k', args.gemm_desc.compiled_dims),
args.gemm_desc.num_groups,
args.gemm_config.layout.block_m, args.gemm_config.layout.block_n, args.gemm_config.layout.block_k,
args.gemm_config.storage_config.swizzle_a_mode, args.gemm_config.storage_config.swizzle_b_mode, args.gemm_config.storage_config.swizzle_cd_mode,
args.gemm_config.pipeline_config.num_stages,
args.gemm_config.launch_config.num_tma_threads, args.gemm_config.launch_config.num_math_threads,
args.gemm_config.layout.get_cluster_size(), args.gemm_config.layout.cluster_n > 1,
args.gemm_config.launch_config.num_sms, to_string(args.gemm_desc.gemm_type),
get_default_epilogue_type(args.epilogue_type));
}
@@ -70,7 +71,7 @@ static void __instantiate_kernel() {{
// TODO: optimize `args` copy
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.sfb, args.grouped_layout,
args.m, args.n, args.k,
args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k,
args.tensor_map_a, args.tensor_map_b,
args.tensor_map_d, args.tensor_map_sfa));
}
@@ -87,45 +88,49 @@ static void sm90_fp8_gemm_1d2d(const torch::Tensor& a, const torch::Tensor& sfa,
DG_HOST_ASSERT(not c.has_value() and d.scalar_type() == torch::kBFloat16);
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K);
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::Normal, KernelType::Kernel1D2D,
m, n, k, 1, major_a, major_b,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), c.has_value(),
device_runtime->get_num_sms());
const auto desc = GemmDesc {
.gemm_type = GemmType::Normal,
.kernel_type = KernelType::Kernel1D2D,
.m = m, .n = n, .k = k, .num_groups = 1,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = c.has_value(),
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims
};
const auto config = get_best_config<SM90ArchSpec>(desc);
// Requires no TMA splits
DG_HOST_ASSERT(config.smem_config.swizzle_a_mode == config.block_k);
DG_HOST_ASSERT(config.smem_config.swizzle_b_mode == config.block_k);
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(major_b, b, n, k,
SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), 1,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_d = make_tma_cd_desc(d, m, static_cast<int>(d.size(-1)),
SM90ArchSpec::get_cd_store_block_m(config.block_m),
SM90ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), 1,
config.smem_config.swizzle_cd_mode);
const auto& tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.block_m, config.block_k, 1, 0);
DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k);
DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k);
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), 1,
config.storage_config.swizzle_b_mode);
const auto tensor_map_d = make_tma_cd_desc(d, m, static_cast<int>(d.size(-1)),
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), 1,
config.storage_config.swizzle_cd_mode);
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.layout.block_m, config.layout.block_k, 1, 0);
// Launch
const SM90FP8Gemm1D2DRuntime::Args& args = {
.major_sfb = major_sfb,
.m = m, .n = n, .k = k,
.num_groups = 1,
.compiled_dims = compiled_dims,
.epilogue_type = epilogue_type,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.epilogue_type = epilogue_type,
.major_sfb = major_sfb,
.sfb = sfb.data_ptr(),
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
@@ -133,8 +138,8 @@ static void sm90_fp8_gemm_1d2d(const torch::Tensor& a, const torch::Tensor& sfa,
.tensor_map_d = tensor_map_d,
.tensor_map_sfa = tensor_map_sfa,
};
const auto& code = SM90FP8Gemm1D2DRuntime::generate(args);
const auto& runtime = compiler->build("sm90_fp8_gemm_1d2d", code);
const auto code = SM90FP8Gemm1D2DRuntime::generate(args);
const auto runtime = compiler->build("sm90_fp8_gemm_1d2d", code);
SM90FP8Gemm1D2DRuntime::launch(runtime, args);
}
@@ -144,49 +149,65 @@ static void sm90_m_grouped_fp8_gemm_contiguous_1d2d(const torch::Tensor& a, cons
const torch::Tensor& m_indices,
const int& num_groups, const int& m, const int& n, const int& k,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b, const cute::UMMA::Major& major_sfb,
const std::string& compiled_dims) {
const std::string& compiled_dims,
const bool& use_psum_layout,
const std::optional<int>& expected_m_for_psum_layout) {
DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16);
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K);
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::MGroupedContiguous, KernelType::Kernel1D2D,
m, n, k, 1, major_a, major_b,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), false,
device_runtime->get_num_sms());
const auto gemm_type = use_psum_layout ?
GemmType::MGroupedContiguousWithPsumLayout : GemmType::MGroupedContiguous;
// Only psum layout can use expected m
if (expected_m_for_psum_layout)
DG_HOST_ASSERT(use_psum_layout);
const auto desc = GemmDesc {
.gemm_type = gemm_type,
.kernel_type = KernelType::Kernel1D2D,
.m = m, .n = n, .k = k, .num_groups = num_groups,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = false,
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims,
.expected_m = expected_m_for_psum_layout.value_or(m),
.expected_n = n, .expected_k = k,
.expected_num_groups = expected_m_for_psum_layout.has_value() ? num_groups : 1
};
const auto config = get_best_config<SM90ArchSpec>(desc);
// Requires no TMA splits
DG_HOST_ASSERT(config.smem_config.swizzle_a_mode == config.block_k);
DG_HOST_ASSERT(config.smem_config.swizzle_b_mode == config.block_k);
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(major_b, b, n, k,
SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_d = make_tma_cd_desc(d, m, n,
SM90ArchSpec::get_cd_store_block_m(config.block_m),
SM90ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), 1,
config.smem_config.swizzle_cd_mode);
const auto& tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.block_m, config.block_k, 1, 0);
DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k);
DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k);
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.storage_config.swizzle_b_mode);
const auto tensor_map_d = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), 1,
config.storage_config.swizzle_cd_mode);
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.layout.block_m, config.layout.block_k, 1, 0);
// Launch
const SM90FP8Gemm1D2DRuntime::Args& args = {
.major_sfb = major_sfb,
.m = m, .n = n, .k = k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
.epilogue_type = std::nullopt,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.epilogue_type = std::nullopt,
.major_sfb = major_sfb,
.sfb = sfb.data_ptr(),
.grouped_layout = m_indices.data_ptr(),
.tensor_map_a = tensor_map_a,
@@ -194,8 +215,8 @@ static void sm90_m_grouped_fp8_gemm_contiguous_1d2d(const torch::Tensor& a, cons
.tensor_map_d = tensor_map_d,
.tensor_map_sfa = tensor_map_sfa,
};
const auto& code = SM90FP8Gemm1D2DRuntime::generate(args);
const auto& runtime = compiler->build("sm90_m_grouped_fp8_gemm_contiguous_1d2d", code);
const auto code = SM90FP8Gemm1D2DRuntime::generate(args);
const auto runtime = compiler->build("sm90_m_grouped_fp8_gemm_contiguous_1d2d", code);
SM90FP8Gemm1D2DRuntime::launch(runtime, args);
}
@@ -210,45 +231,50 @@ static void sm90_m_grouped_fp8_gemm_masked_1d2d(const torch::Tensor& a, const to
DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16);
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K);
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::MGroupedMasked, KernelType::Kernel1D2D,
expected_m, n, k, num_groups, major_a, major_b,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), false,
device_runtime->get_num_sms());
const auto desc = GemmDesc {
.gemm_type = GemmType::MGroupedMasked,
.kernel_type = KernelType::Kernel1D2D,
.m = m, .n = n, .k = k, .num_groups = num_groups,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = false,
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims,
.expected_m = expected_m, .expected_n = n, .expected_k = k, .expected_num_groups = num_groups
};
const auto config = get_best_config<SM90ArchSpec>(desc);
// Requires no TMA splits
DG_HOST_ASSERT(config.smem_config.swizzle_a_mode == config.block_k);
DG_HOST_ASSERT(config.smem_config.swizzle_b_mode == config.block_k);
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), num_groups,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(major_b, b, n, k,
SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_d = make_tma_cd_desc(d, m, n,
SM90ArchSpec::get_cd_store_block_m(config.block_m),
SM90ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), num_groups,
config.smem_config.swizzle_cd_mode);
const auto& tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.block_m, config.block_k, num_groups, 0);
DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k);
DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k);
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), num_groups,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.storage_config.swizzle_b_mode);
const auto tensor_map_d = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), num_groups,
config.storage_config.swizzle_cd_mode);
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.layout.block_m, config.layout.block_k, num_groups, 0);
// Launch
const SM90FP8Gemm1D2DRuntime::Args& args = {
.major_sfb = major_sfb,
.m = m, .n = n, .k = k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
.epilogue_type = std::nullopt,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.epilogue_type = std::nullopt,
.major_sfb = major_sfb,
.sfb = sfb.data_ptr(),
.grouped_layout = masked_m.data_ptr(),
.tensor_map_a = tensor_map_a,
@@ -256,8 +282,8 @@ static void sm90_m_grouped_fp8_gemm_masked_1d2d(const torch::Tensor& a, const to
.tensor_map_d = tensor_map_d,
.tensor_map_sfa = tensor_map_sfa,
};
const auto& code = SM90FP8Gemm1D2DRuntime::generate(args);
const auto& runtime = compiler->build("sm90_fp8_m_grouped_gemm_masked_1d2d", code);
const auto code = SM90FP8Gemm1D2DRuntime::generate(args);
const auto runtime = compiler->build("sm90_fp8_m_grouped_gemm_masked_1d2d", code);
SM90FP8Gemm1D2DRuntime::launch(runtime, args);
}
@@ -271,51 +297,55 @@ static void sm90_fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa,
DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16);
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K);
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::Batched, KernelType::Kernel1D2D,
m, n, k, batch_size, major_a, major_b,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), c.has_value(),
device_runtime->get_num_sms());
const auto desc = GemmDesc {
.gemm_type = GemmType::Batched,
.kernel_type = KernelType::Kernel1D2D,
.m = m, .n = n, .k = k, .num_groups = batch_size,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = c.has_value(),
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims
};
const auto config = get_best_config<SM90ArchSpec>(desc);
// Requires no TMA splits
DG_HOST_ASSERT(config.smem_config.swizzle_a_mode == config.block_k);
DG_HOST_ASSERT(config.smem_config.swizzle_b_mode == config.block_k);
const int& load_block_m = SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m);
const auto& tensor_map_a = make_tma_3d_desc(a, k, m, batch_size,
config.block_k, load_block_m, 1,
a.stride(1),
a.stride(0),
config.smem_config.swizzle_a_mode);
DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k);
DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k);
const int load_block_m = config.storage_config.load_block_m;
const auto tensor_map_a = make_tma_3d_desc(a, k, m, batch_size,
config.layout.block_k, load_block_m, 1,
a.stride(1),
a.stride(0),
config.storage_config.swizzle_a_mode);
const int& load_block_n = SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n);
const auto& tensor_map_b = make_tma_3d_desc(b, k, n, batch_size,
config.block_k, load_block_n, 1,
b.stride(1),
b.stride(0),
config.smem_config.swizzle_b_mode);
const int load_block_n = config.storage_config.load_block_n;
const auto tensor_map_b = make_tma_3d_desc(b, k, n, batch_size,
config.layout.block_k, load_block_n, 1,
b.stride(1),
b.stride(0),
config.storage_config.swizzle_b_mode);
const int& store_block_m = SM90ArchSpec::get_cd_store_block_m(config.block_m);
const int& store_block_n = SM90ArchSpec::get_cd_store_block_n(config.block_n);
const auto& tensor_map_d = make_tma_3d_desc(d, n, m, batch_size,
store_block_n, store_block_m, 1,
d.stride(1), d.stride(0),
config.smem_config.swizzle_cd_mode);
const int store_block_m = config.storage_config.store_block_m;
const int store_block_n = config.storage_config.store_block_n;
const auto tensor_map_d = make_tma_3d_desc(d, n, m, batch_size,
store_block_n, store_block_m, 1,
d.stride(1), d.stride(0),
config.storage_config.swizzle_cd_mode);
const auto& tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.block_m, config.block_k, batch_size, 0);
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.layout.block_m, config.layout.block_k, batch_size, 0);
// Launch
const SM90FP8Gemm1D2DRuntime::Args& args = {
.major_sfb = major_sfb,
.m = m, .n = n, .k = k,
.num_groups = batch_size,
.compiled_dims = compiled_dims,
.epilogue_type = std::nullopt,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.epilogue_type = std::nullopt,
.major_sfb = major_sfb,
.sfb = sfb.data_ptr(),
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
@@ -323,8 +353,8 @@ static void sm90_fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa,
.tensor_map_d = tensor_map_d,
.tensor_map_sfa = tensor_map_sfa,
};
const auto& code = SM90FP8Gemm1D2DRuntime::generate(args);
const auto& runtime = compiler->build("sm90_fp8_gemm_1d2d", code);
const auto code = SM90FP8Gemm1D2DRuntime::generate(args);
const auto runtime = compiler->build("sm90_fp8_gemm_1d2d", code);
SM90FP8Gemm1D2DRuntime::launch(runtime, args);
}

36
csrc/jit_kernels/impls/sm90_tf32_hc_prenorm_gemm.hpp
View File

@@ -81,21 +81,21 @@ static void sm90_tf32_hc_prenorm_gemm(const torch::Tensor& a,
DG_HOST_ASSERT(n <= 32 and n % 8 == 0);
DG_HOST_ASSERT(k % block_k == 0);
const auto& swizzle_cd_mode = get_swizzle_mode(block_n, sizeof(float));
const auto& tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a, m, k,
block_m, block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(cute::UMMA::Major::K))), 1,
get_swizzle_mode(block_k, a.element_size()), 0,
true);
const auto& tensor_map_b = make_tma_b_desc(cute::UMMA::Major::K, b, n, k,
block_n, block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(cute::UMMA::Major::K))), 1,
get_swizzle_mode(block_k, b.element_size()), 0,
true);
const auto& tensor_map_d = num_splits == 1 ? make_tma_cd_desc(d, m, n,
block_m, block_n,
static_cast<int>(d.stride(-2)), 1,
swizzle_cd_mode)
const auto swizzle_cd_mode = get_swizzle_mode(block_n, sizeof(float));
const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a, m, k,
block_m, block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(cute::UMMA::Major::K))), 1,
get_swizzle_mode(block_k, a.element_size()), 0,
true);
const auto tensor_map_b = make_tma_b_desc(cute::UMMA::Major::K, b, n, k,
block_n, block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(cute::UMMA::Major::K))), 1,
get_swizzle_mode(block_k, b.element_size()), 0,
true);
const auto tensor_map_d = num_splits == 1 ? make_tma_cd_desc(d, m, n,
block_m, block_n,
static_cast<int>(d.stride(-2)), 1,
swizzle_cd_mode)
: make_tma_3d_desc(d, n, m, num_splits,
block_n, block_m, 1,
static_cast<int>(d.stride(-2)),
@@ -138,14 +138,14 @@ static void sm90_tf32_hc_prenorm_gemm(const torch::Tensor& a,
.num_stages = num_stages,
.num_math_threads = num_math_threads,
.num_tma_threads = num_tma_threads,
.launch_args = LaunchArgs(num_splits * ceil_div(m, block_m), num_threads, smem_size, 1),
.launch_args = LaunchArgs(num_splits * ceil_div(m, block_m), num_threads, smem_size),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_d = tensor_map_d,
.sqr_sum = sqr_sum.data_ptr<float>()
};
const auto& code = SM90BF16HCPrenormGemmRuntime::generate(args);
const auto& runtime = compiler->build("sm90_tf32_hc_prenorm_gemm", code);
const auto code = SM90BF16HCPrenormGemmRuntime::generate(args);
const auto runtime = compiler->build("sm90_tf32_hc_prenorm_gemm", code);
SM90BF16HCPrenormGemmRuntime::launch(runtime, args);
}

14
csrc/jit_kernels/impls/smxx_clean_logits.hpp
View File

@@ -17,7 +17,8 @@ public:
int* cu_seq_len_k_start;
int* cu_seq_len_k_end;
float* logits;
void* logits;
at::ScalarType logits_dtype;
int block_kv;
int num_warps;
@@ -33,10 +34,10 @@ using namespace deep_gemm;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&smxx_clean_logits<
{}, {}, {}
{}, {}, {}, {}
>);
}};
)", args.next_n, args.block_kv, args.num_warps);
)", args.next_n, args.block_kv, args.num_warps, to_string(args.logits_dtype));
}
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
@@ -65,14 +66,15 @@ static void smxx_clean_logits(const torch::Tensor& logits,
.stride_logits = stride_logits,
.cu_seq_len_k_start = cu_seq_len_k_start.has_value() ? cu_seq_len_k_start.value().data_ptr<int>() : nullptr,
.cu_seq_len_k_end = cu_seq_len_k_end.data_ptr<int>(),
.logits = logits.data_ptr<float>(),
.logits = logits.data_ptr(),
.logits_dtype = logits.scalar_type(),
.block_kv = block_kv,
.num_warps = num_warps,
.launch_args = LaunchArgs(device_runtime->get_num_sms(),
num_warps * 32, smem_size)
};
const auto& code = SMXXCleanLogitsRuntime::generate(args);
const auto& runtime = compiler->build("smxx_clean_logits", code);
const auto code = SMXXCleanLogitsRuntime::generate(args);
const auto runtime = compiler->build("smxx_clean_logits", code);
SMXXCleanLogitsRuntime::launch(runtime, args);
}

73
csrc/jit_kernels/impls/smxx_cublaslt.hpp
View File

@@ -46,7 +46,7 @@ static void call_cublaslt_api(const cublasOperation_t& trans_a,
DG_CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_SCALE_TYPE, &scale_type, sizeof(scale_type)));
#if DG_CUBLASLT_ADVANCED_FEATURES_COMPATIBLE
const int& math_sms = device_runtime->get_num_sms();
const int math_sms = device_runtime->get_num_sms();
DG_CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_SM_COUNT_TARGET, &math_sms, sizeof(math_sms)));
#endif
@@ -57,10 +57,10 @@ static void call_cublaslt_api(const cublasOperation_t& trans_a,
#endif
// Get cuBLASLt handle, workspace, and stream
const auto& handle = device_runtime->get_cublaslt_handle();
const auto& workspace = device_runtime->get_cublaslt_workspace();
const auto& workspace_bytes = workspace.nbytes();
const auto& stream = at::cuda::getCurrentCUDAStream();
const auto handle = device_runtime->get_cublaslt_handle();
const auto workspace = device_runtime->get_cublaslt_workspace();
const auto workspace_bytes = workspace.nbytes();
const auto stream = at::cuda::getCurrentCUDAStream();
// Algorithm selection
cublasLtMatmulPreference_t pref;
@@ -77,7 +77,7 @@ static void call_cublaslt_api(const cublasOperation_t& trans_a,
DG_HOST_ASSERT(num_heuristic_results == 1 and "Unable to find any algorithm for the GEMM");
// Call: D = alpha * (A @ B) + beta * C
const float& alpha = 1.0, beta = accumulate ? 1.0 : 0.0;
const float alpha = 1.0, beta = accumulate ? 1.0 : 0.0;
DG_CUBLASLT_CHECK(cublasLtMatmul(handle, // Light handle
desc, // Operation description
&alpha, // Alpha
@@ -99,47 +99,36 @@ static void call_cublaslt_api(const cublasOperation_t& trans_a,
}
static void cublaslt_gemm(const torch::Tensor& lhs, const torch::Tensor& rhs,
const std::optional<torch::Tensor>& acc,
const torch::Tensor& out,
const int& m, const int& n, const int& k,
const cute::UMMA::Major& a_major, const cute::UMMA::Major& b_major) {
const auto& trans_a = b_major == cute::UMMA::Major::K ? CUBLAS_OP_T : CUBLAS_OP_N;
const auto& trans_b = a_major == cute::UMMA::Major::K ? CUBLAS_OP_N : CUBLAS_OP_T;
// Duplicate the accumulator if necessary
// TODO: remove this
if (acc.has_value()) {
if (acc->data_ptr() == out.data_ptr()) {
DG_HOST_ASSERT(acc->sizes() == out.sizes() and acc->strides() == out.strides());
} else {
out.copy_(acc.value());
}
}
const cute::UMMA::Major& a_major, const cute::UMMA::Major& b_major,
const bool& accumulate) {
const auto trans_a = b_major == cute::UMMA::Major::K ? CUBLAS_OP_T : CUBLAS_OP_N;
const auto trans_b = a_major == cute::UMMA::Major::K ? CUBLAS_OP_N : CUBLAS_OP_T;
// Matrix layouts
const auto& cuda_type_a = at::cuda::ScalarTypeToCudaDataType(rhs.scalar_type());
const auto& cuda_type_b = at::cuda::ScalarTypeToCudaDataType(lhs.scalar_type());
const auto& cuda_type_d = at::cuda::ScalarTypeToCudaDataType(out.scalar_type());
const auto& layout_a = b_major == cute::UMMA::Major::K ? get_cublaslt_layout(cuda_type_a, k, n, rhs.stride(0))
: get_cublaslt_layout(cuda_type_a, n, k, rhs.stride(1));
const auto& layout_b = a_major == cute::UMMA::Major::K ? get_cublaslt_layout(cuda_type_b, k, m, lhs.stride(0))
: get_cublaslt_layout(cuda_type_b, m, k, lhs.stride(1));
const auto& layout_d = get_cublaslt_layout(cuda_type_d, n, m, out.stride(0));
const auto cuda_type_a = at::cuda::ScalarTypeToCudaDataType(rhs.scalar_type());
const auto cuda_type_b = at::cuda::ScalarTypeToCudaDataType(lhs.scalar_type());
const auto cuda_type_d = at::cuda::ScalarTypeToCudaDataType(out.scalar_type());
const auto layout_a = b_major == cute::UMMA::Major::K ? get_cublaslt_layout(cuda_type_a, k, n, rhs.stride(0))
: get_cublaslt_layout(cuda_type_a, n, k, rhs.stride(1));
const auto layout_b = a_major == cute::UMMA::Major::K ? get_cublaslt_layout(cuda_type_b, k, m, lhs.stride(0))
: get_cublaslt_layout(cuda_type_b, m, k, lhs.stride(1));
const auto layout_d = get_cublaslt_layout(cuda_type_d, n, m, out.stride(0));
call_cublaslt_api(trans_a, trans_b, layout_a, layout_b, layout_d, lhs, rhs, out, acc.has_value());
call_cublaslt_api(trans_a, trans_b, layout_a, layout_b, layout_d, lhs, rhs, out, accumulate);
}
static void cublaslt_bhr_hdr_bhd(const torch::Tensor& lhs, const torch::Tensor& rhs, const torch::Tensor& out,
const int& b, const int& h, const int& r, const int& d) {
const auto& m = d, n = b, k = r;
const auto& trans_a = CUBLAS_OP_T;
const auto& trans_b = CUBLAS_OP_N;
const auto m = d, n = b, k = r;
const auto trans_a = CUBLAS_OP_T;
const auto trans_b = CUBLAS_OP_N;
// Matrix layouts
const auto& layout_a = get_cublaslt_layout(CUDA_R_16BF, k, m, rhs.stride(1), h, rhs.stride(0));
const auto& layout_b = get_cublaslt_layout(CUDA_R_16BF, k, n, lhs.stride(0), h, lhs.stride(1));
const auto& layout_d = get_cublaslt_layout(CUDA_R_16BF, m, n, out.stride(0), h, out.stride(1));
const auto layout_a = get_cublaslt_layout(CUDA_R_16BF, k, m, rhs.stride(1), h, rhs.stride(0));
const auto layout_b = get_cublaslt_layout(CUDA_R_16BF, k, n, lhs.stride(0), h, lhs.stride(1));
const auto layout_d = get_cublaslt_layout(CUDA_R_16BF, m, n, out.stride(0), h, out.stride(1));
call_cublaslt_api(trans_a, trans_b, layout_a, layout_b, layout_d, lhs, rhs, out, false);
}
@@ -147,14 +136,14 @@ static void cublaslt_bhr_hdr_bhd(const torch::Tensor& lhs, const torch::Tensor&
static void cublaslt_bhd_hdr_bhr(const torch::Tensor& lhs, const torch::Tensor& rhs, const torch::Tensor& out,
const int& b, const int& h, const int& r, const int& d) {
const auto& m = r, n = b, k = d;
const auto& trans_a = CUBLAS_OP_N;
const auto& trans_b = CUBLAS_OP_N;
const auto m = r, n = b, k = d;
const auto trans_a = CUBLAS_OP_N;
const auto trans_b = CUBLAS_OP_N;
// Matrix layouts
const auto& layout_a = get_cublaslt_layout(CUDA_R_16BF, m, k, rhs.stride(1), h, rhs.stride(0));
const auto& layout_b = get_cublaslt_layout(CUDA_R_16BF, k, n, lhs.stride(0), h, lhs.stride(1));
const auto& layout_d = get_cublaslt_layout(CUDA_R_16BF, m, n, out.stride(0), h, out.stride(1));
const auto layout_a = get_cublaslt_layout(CUDA_R_16BF, m, k, rhs.stride(1), h, rhs.stride(0));
const auto layout_b = get_cublaslt_layout(CUDA_R_16BF, k, n, lhs.stride(0), h, lhs.stride(1));
const auto layout_d = get_cublaslt_layout(CUDA_R_16BF, m, n, out.stride(0), h, out.stride(1));
call_cublaslt_api(trans_a, trans_b, layout_a, layout_b, layout_d, lhs, rhs, out, false);
}

328
csrc/jit_kernels/impls/smxx_fp8_fp4_mqa_logits.hpp Normal file
View File

@@ -0,0 +1,328 @@
#pragma once
#include "../../jit/compiler.hpp"
#include "../../jit/device_runtime.hpp"
#include "../../jit/kernel_runtime.hpp"
#include "../heuristics/sm90.hpp"
#include "../heuristics/sm100.hpp"
#include "runtime_utils.hpp"
namespace deep_gemm {
class SMXXFP8MQALogitsRuntime final: public LaunchRuntime<SMXXFP8MQALogitsRuntime> {
public:
struct Args {
int seq_len;
int seq_len_kv;
int max_seqlen_k;
int stride_logits;
int num_heads, head_dim;
bool is_compressed_logits;
int num_q_stages;
int num_kv_stages;
int block_q;
int block_kv;
int* cu_seq_len_k_start;
int* cu_seq_len_k_end;
void* logits;
CUtensorMap tensor_map_q;
CUtensorMap tensor_map_kv;
CUtensorMap tensor_map_kv_scales;
CUtensorMap tensor_map_weights;
at::ScalarType logits_dtype;
int num_specialized_threads;
int num_math_threads;
LaunchArgs launch_args;
};
static std::string generate_impl(const Args& args) {
// TODO: optimize performance by tuning args
// Block sizes are fixed in this kernel
DG_HOST_ASSERT(128 % args.num_heads == 0);
const auto arch = device_runtime->get_arch(true);
return fmt::format(R"(
#include <deep_gemm/impls/sm{}_fp8_mqa_logits.cuh>
using namespace deep_gemm;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&sm{}_fp8_mqa_logits<
{}, {},
{},
{}, {},
{}, {},
{},
{}, {},
{}
>);
}};
)", arch, arch,
args.num_heads, args.head_dim,
args.is_compressed_logits,
args.block_q, args.block_kv,
args.num_q_stages, args.num_kv_stages,
args.launch_args.grid_dim.first,
args.num_specialized_threads, args.num_math_threads,
to_string(args.logits_dtype));
}
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.seq_len, args.seq_len_kv,
args.max_seqlen_k, args.stride_logits,
args.cu_seq_len_k_start, args.cu_seq_len_k_end,
args.logits,
args.tensor_map_q, args.tensor_map_kv,
args.tensor_map_kv_scales, args.tensor_map_weights
));
}
};
static void smxx_fp8_mqa_logits(const torch::Tensor& q,
const torch::Tensor& kv, const torch::Tensor& kv_scales,
const torch::Tensor& weights,
const torch::Tensor& cu_seq_len_k_start,
const torch::Tensor& cu_seq_len_k_end,
const torch::Tensor& logits,
const at::ScalarType& logits_dtype,
const int& seq_len, const int& seq_len_kv,
const int& max_seqlen_k, const int& stride_logits,
const int& num_heads, const int& head_dim,
const int& block_q, const int& block_kv) {
constexpr int num_specialized_threads = 128;
constexpr int num_q_stages = 3, num_kv_stages = 3;
const int num_math_threads = (device_runtime->get_arch_major() == 10 ? 256 : 512);
// Use compressed logits format when max_seqlen_k is specified
const bool is_compressed_logits = (max_seqlen_k > 0);
// Construct TMAs
DG_HOST_ASSERT(head_dim == 32 or head_dim == 64 or head_dim == 128);
const auto tensor_map_q = make_tma_2d_desc(q, head_dim, seq_len * num_heads,
head_dim, block_q * num_heads, head_dim, head_dim);
const auto tensor_map_kv = make_tma_2d_desc(kv, head_dim, seq_len_kv,
head_dim, block_kv, head_dim, head_dim);
// According to the driver API, the minimal alignment is 256 bytes
// So it is safe for us to do a 16-byte OOB
const auto tensor_map_kv_scales = make_tma_2d_desc(kv_scales,
get_tma_aligned_size(seq_len_kv, static_cast<int>(kv_scales.element_size())),
1, block_kv, 1, 0, 0);
const auto tensor_map_weights = make_tma_2d_desc(weights, num_heads, seq_len,
num_heads, block_q, num_heads, 0);
// Calculate shared memory size
int smem_size = 0;
const int smem_q_size_per_stage = block_q * num_heads * head_dim * static_cast<int>(q.element_size());
const int smem_weight_size_per_stage = block_q * num_heads * static_cast<int>(weights.element_size());
const int smem_kv_size_per_stage = block_kv * head_dim * static_cast<int>(kv.element_size());
const int kv_scale_size_per_stage = block_kv * static_cast<int>(kv_scales.element_size());
smem_size += num_q_stages * smem_q_size_per_stage;
smem_size += num_kv_stages * smem_kv_size_per_stage;
smem_size += num_q_stages * smem_weight_size_per_stage;
smem_size += num_kv_stages * kv_scale_size_per_stage;
smem_size += (num_q_stages * 2 + num_kv_stages * 2 + (num_math_threads / 128) * 2) * 8;
smem_size += 4;
DG_HOST_ASSERT(smem_size <= SM90ArchSpec::smem_capacity);
DG_HOST_ASSERT(smem_size <= SM100ArchSpec::smem_capacity);
// Launch
const SMXXFP8MQALogitsRuntime::Args args = {
.seq_len = seq_len,
.seq_len_kv = seq_len_kv,
.max_seqlen_k = max_seqlen_k,
.stride_logits = stride_logits,
.num_heads = num_heads, .head_dim = head_dim,
.is_compressed_logits = is_compressed_logits,
.num_q_stages = num_q_stages,
.num_kv_stages = num_kv_stages,
.block_q = block_q,
.block_kv = block_kv,
.cu_seq_len_k_start = cu_seq_len_k_start.data_ptr<int>(),
.cu_seq_len_k_end = cu_seq_len_k_end.data_ptr<int>(),
.logits = logits.data_ptr(),
.tensor_map_q = tensor_map_q,
.tensor_map_kv = tensor_map_kv,
.tensor_map_kv_scales = tensor_map_kv_scales,
.tensor_map_weights = tensor_map_weights,
.logits_dtype = logits_dtype,
.num_specialized_threads = num_specialized_threads,
.num_math_threads = num_math_threads,
.launch_args = LaunchArgs(device_runtime->get_num_sms(),
num_specialized_threads + num_math_threads,
smem_size)
};
const auto code = SMXXFP8MQALogitsRuntime::generate(args);
const auto runtime = compiler->build("smxx_fp8_mqa_logits", code);
SMXXFP8MQALogitsRuntime::launch(runtime, args);
}
class SM100FP4MQALogitsRuntime final: public LaunchRuntime<SM100FP4MQALogitsRuntime> {
public:
struct Args {
int seq_len;
int seq_len_kv;
int max_seqlen_k;
int stride_logits;
int num_heads, head_dim;
bool is_compressed_logits;
int num_q_stages;
int num_kv_stages;
int block_q;
int block_kv;
int* cu_seq_len_k_start;
int* cu_seq_len_k_end;
void* logits;
CUtensorMap tensor_map_q;
CUtensorMap tensor_map_sf_q;
CUtensorMap tensor_map_kv;
CUtensorMap tensor_map_sf_kv;
CUtensorMap tensor_map_weights;
at::ScalarType logits_dtype;
int num_specialized_threads;
int num_math_threads;
LaunchArgs launch_args;
};
static std::string generate_impl(const Args& args) {
// TODO: optimize performance by tuning args
// Block sizes are fixed in this kernel
DG_HOST_ASSERT(128 % args.num_heads == 0);
const auto arch = device_runtime->get_arch(true);
return fmt::format(R"(
#include <deep_gemm/impls/sm100_fp4_mqa_logits.cuh>
using namespace deep_gemm;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&sm100_fp4_mqa_logits<
{}, {},
{},
{}, {},
{}, {},
{},
{}, {},
{}
>);
}};
)", args.num_heads, args.head_dim,
args.is_compressed_logits,
args.block_q, args.block_kv,
args.num_q_stages, args.num_kv_stages,
args.launch_args.grid_dim.first,
args.num_specialized_threads, args.num_math_threads,
to_string(args.logits_dtype));
}
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.seq_len, args.seq_len_kv,
args.max_seqlen_k, args.stride_logits,
args.cu_seq_len_k_start, args.cu_seq_len_k_end,
args.logits,
args.tensor_map_q, args.tensor_map_sf_q,
args.tensor_map_kv, args.tensor_map_sf_kv,
args.tensor_map_weights
));
}
};
static void sm100_fp4_mqa_logits(const torch::Tensor& q, const torch::Tensor& sf_q,
const torch::Tensor& kv, const torch::Tensor& sf_kv,
const torch::Tensor& weights,
const torch::Tensor& cu_seq_len_k_start,
const torch::Tensor& cu_seq_len_k_end,
const torch::Tensor& logits,
const at::ScalarType& logits_dtype,
const int& seq_len, const int& seq_len_kv,
const int& max_seqlen_k, const int& stride_logits,
const int& num_heads, const int& head_dim,
const int& block_q, const int& block_kv) {
constexpr int num_specialized_threads = 128;
const int num_math_threads = 2 * 128;
constexpr int num_q_stages = 3, num_kv_stages = 6, num_tmem_stages = 3;
// Use compressed logits format when max_seqlen_k is specified
const bool is_compressed_logits = (max_seqlen_k > 0);
// Construct TMAs
// `head_dim` must be 128 for 64B swizzling
DG_HOST_ASSERT(head_dim == 128);
const auto tensor_map_q = make_tma_2d_desc(q, head_dim, seq_len * num_heads,
head_dim, block_q * num_heads,
static_cast<int>(q.stride(1)),
head_dim / 2, 0, false, false);
const auto tensor_map_sf_q = make_tma_2d_desc(sf_q, num_heads, seq_len,
num_heads, block_q,
static_cast<int>(sf_q.stride(0)), 0);
const auto tensor_map_weights = make_tma_2d_desc(weights, num_heads, seq_len,
num_heads, block_q,
static_cast<int>(weights.stride(0)), 0);
const auto tensor_map_kv = make_tma_2d_desc(kv, head_dim, seq_len_kv,
head_dim, block_kv,
static_cast<int>(kv.stride(0)),
head_dim / 2, 0, false, false);
// According to the driver API, the minimal alignment is 256 bytes
// So it is safe for us to do a 16-byte OOB
const auto tensor_map_sf_kv = make_tma_2d_desc(sf_kv,
get_tma_aligned_size(seq_len_kv, static_cast<int>(sf_kv.element_size())), 1,
block_kv, 1, 0, 0);
// Calculate shared memory size
const int smem_q_size_per_stage = block_q * num_heads * head_dim / 2;
const int smem_sf_q_size_per_stage = align(block_q * num_heads, 128) * sizeof(int);
const int smem_kv_size_per_stage = block_kv * head_dim / 2;
const int smem_sf_kv_size_per_stage = align(block_kv, 128) * sizeof(int);
const int smem_weight_size_per_stage = block_q * num_heads * sizeof(float);
const int smem_barriers = (num_q_stages + num_kv_stages + num_tmem_stages) * 2 * 8;
const int smem_tmem_ptr = 4;
const int smem_size = num_q_stages * (smem_q_size_per_stage + smem_sf_q_size_per_stage + smem_weight_size_per_stage) +
num_kv_stages * (smem_kv_size_per_stage + smem_sf_kv_size_per_stage) +
smem_barriers + smem_tmem_ptr;
DG_HOST_ASSERT(smem_size <= SM100ArchSpec::smem_capacity);
// Launch
const SM100FP4MQALogitsRuntime::Args args = {
.seq_len = seq_len,
.seq_len_kv = seq_len_kv,
.max_seqlen_k = max_seqlen_k,
.stride_logits = stride_logits,
.num_heads = num_heads, .head_dim = head_dim,
.is_compressed_logits = is_compressed_logits,
.num_q_stages = num_q_stages,
.num_kv_stages = num_kv_stages,
.block_q = block_q,
.block_kv = block_kv,
.cu_seq_len_k_start = cu_seq_len_k_start.data_ptr<int>(),
.cu_seq_len_k_end = cu_seq_len_k_end.data_ptr<int>(),
.logits = logits.data_ptr(),
.tensor_map_q = tensor_map_q,
.tensor_map_sf_q = tensor_map_sf_q,
.tensor_map_kv = tensor_map_kv,
.tensor_map_sf_kv = tensor_map_sf_kv,
.tensor_map_weights = tensor_map_weights,
.logits_dtype = logits_dtype,
.num_specialized_threads = num_specialized_threads,
.num_math_threads = num_math_threads,
.launch_args = LaunchArgs(device_runtime->get_num_sms(),
num_specialized_threads + num_math_threads,
smem_size)
};
const auto code = SM100FP4MQALogitsRuntime::generate(args);
const auto runtime = compiler->build("sm100_fp4_mqa_logits", code);
SM100FP4MQALogitsRuntime::launch(runtime, args);
}
} // namespace deep_gemm

445
csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp Normal file
View File

@@ -0,0 +1,445 @@
#pragma once
#include "../../jit/compiler.hpp"
#include "../../jit/device_runtime.hpp"
#include "../../jit/kernel_runtime.hpp"
#include "../heuristics/sm90.hpp"
#include "runtime_utils.hpp"
namespace deep_gemm {
class SMXXPagedMQALogitsMetadataRuntime final: public LaunchRuntime<SMXXPagedMQALogitsMetadataRuntime> {
public:
struct Args {
int aligned_batch_size;
int split_kv;
int num_sms;
int batch_size;
int next_n;
bool is_context_lens_2d;
int* context_lens;
int* schedule_metadata;
LaunchArgs launch_args;
};
static std::string generate_impl(const Args& args) {
return fmt::format(R"(
#include <deep_gemm/scheduler/paged_mqa_logits.cuh>
using namespace deep_gemm;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&sched::smxx_paged_mqa_logits_metadata<
{}, {}, {}
>);
}};
)", args.aligned_batch_size, args.split_kv, args.num_sms);
}
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.batch_size,
args.next_n,
args.is_context_lens_2d,
args.context_lens,
args.schedule_metadata
));
}
};
static void smxx_paged_mqa_logits_metadata(const torch::Tensor& context_lens,
const torch::Tensor& schedule_metadata,
const int& batch_size, const int& next_n,
const int& block_kv, const int& num_sms,
const bool& is_context_lens_2d) {
constexpr int split_kv = 256;
constexpr int num_threads = 32;
const int aligned_batch_size = align(batch_size, 32);
DG_HOST_ASSERT(split_kv % block_kv == 0);
// Calculate shared memory size
const int smem_size = aligned_batch_size * static_cast<int>(sizeof(int));
DG_HOST_ASSERT(smem_size <= SM90ArchSpec::smem_capacity);
DG_HOST_ASSERT(smem_size <= SM100ArchSpec::smem_capacity);
// Launch
const SMXXPagedMQALogitsMetadataRuntime::Args& args = {
.aligned_batch_size = aligned_batch_size,
.split_kv = split_kv,
.num_sms = num_sms,
.batch_size = batch_size,
.next_n = next_n,
.is_context_lens_2d = is_context_lens_2d,
.context_lens = context_lens.data_ptr<int>(),
.schedule_metadata = schedule_metadata.data_ptr<int>(),
.launch_args = LaunchArgs(1, num_threads, smem_size)
};
const auto code = SMXXPagedMQALogitsMetadataRuntime::generate(args);
const auto runtime = compiler->build("smxx_paged_mqa_logits_metadata", code);
SMXXPagedMQALogitsMetadataRuntime::launch(runtime, args);
}
class SMXXFP8PagedMQALogitsRuntime final: public LaunchRuntime<SMXXFP8PagedMQALogitsRuntime> {
public:
struct Args {
int batch_size;
int next_n;
int num_heads;
int head_dim;
int block_kv;
bool is_context_lens_2d;
int block_table_stride;
int logits_stride;
int num_q_stages;
int num_kv_stages;
int split_kv;
int* context_lens;
void* logits;
int* block_table;
int* schedule_meta;
CUtensorMap tensor_map_q;
CUtensorMap tensor_map_kv;
CUtensorMap tensor_map_kv_scales;
CUtensorMap tensor_map_weights;
at::ScalarType logits_dtype;
int num_specialized_threads;
int num_math_threads;
LaunchArgs launch_args;
};
static std::string generate_impl(const Args& args) {
// TODO: optimize performance by tuning args
// Block sizes are fixed in this kernel
DG_HOST_ASSERT(128 % args.num_heads == 0);
const auto arch = device_runtime->get_arch(true);
return fmt::format(R"(
#include <deep_gemm/impls/sm{}_fp8_paged_mqa_logits.cuh>
using namespace deep_gemm;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&sm{}_fp8_paged_mqa_logits<
{}, {},
{}, {},
{},
{}, {},
{},
{}, {},
{}
>);
}};
)", arch, arch,
args.next_n, args.num_heads,
args.head_dim, args.block_kv,
args.is_context_lens_2d,
args.num_q_stages, args.num_kv_stages,
args.split_kv,
args.num_specialized_threads, args.num_math_threads,
to_string(args.logits_dtype));
}
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.batch_size,
args.logits_stride, args.block_table_stride,
args.context_lens, args.logits,
args.block_table, args.schedule_meta,
args.tensor_map_q, args.tensor_map_kv,
args.tensor_map_kv_scales, args.tensor_map_weights
));
}
};
static void smxx_fp8_paged_mqa_logits(const torch::Tensor& q,
const torch::Tensor& kv_cache,
const torch::Tensor& kv_cache_scales,
const torch::Tensor& weights,
const torch::Tensor& context_lens,
const torch::Tensor& logits,
const torch::Tensor& block_table,
const torch::Tensor& schedule_meta,
const at::ScalarType& logits_dtype,
const int& batch_size, const int& next_n,
const int& num_heads, const int& head_dim,
const int& num_kv_blocks, const int& block_kv,
const bool& is_context_lens_2d,
const int& logits_stride,
const int& block_table_stride,
const int& num_sms,
const int& split_kv) {
const int num_specialized_threads = 128;
const int mma_m = (device_runtime->get_arch_major() == 10 ? 128 : 64);
const int num_math_warp_groups = split_kv / mma_m;
const int num_math_threads = num_math_warp_groups * 128;
const int num_q_stages = 3, num_kv_stages = (device_runtime->get_arch_major() == 10 ? 4 : 3);
DG_HOST_ASSERT(split_kv % mma_m == 0 and logits_stride % split_kv == 0);
// Construct TMAs
const int next_n_atom = (next_n % 2 == 0) ? 2 : 1;
const auto tensor_map_q = make_tma_2d_desc(q, head_dim, batch_size * next_n * num_heads,
head_dim, next_n_atom * num_heads,
static_cast<int>(q.stride(2)),
head_dim);
const auto tensor_map_kv = make_tma_3d_desc(kv_cache, head_dim, block_kv, num_kv_blocks,
head_dim, block_kv, 1,
static_cast<int>(kv_cache.stride(1)),
static_cast<int>(kv_cache.stride(0)),
head_dim);
const auto tensor_map_kv_scales = make_tma_2d_desc(kv_cache_scales, block_kv, num_kv_blocks,
block_kv, 1,
static_cast<int>(kv_cache_scales.stride(0)), 0);
const auto tensor_map_weights = make_tma_2d_desc(weights, num_heads, batch_size * next_n,
num_heads, next_n_atom,
static_cast<int>(weights.stride(0)), 0);
// Calculate shared memory size
int smem_size = 0;
if (device_runtime->get_arch_major() == 9) {
const int swizzle_alignment = head_dim * 8;
const int smem_q_size_per_stage = next_n * num_heads * head_dim * static_cast<int>(q.element_size());
const int aligned_smem_weight_size_per_stage = align(next_n * num_heads * static_cast<int>(weights.element_size()), swizzle_alignment);
const int smem_q_pipe_size = num_q_stages * (smem_q_size_per_stage + aligned_smem_weight_size_per_stage) + align(num_q_stages * 8 * 2, swizzle_alignment);
const int smem_kv_size_per_stage = block_kv * head_dim * static_cast<int>(kv_cache.element_size());
const int aligned_smem_kv_scale_size_per_stage = align(block_kv * static_cast<int>(kv_cache_scales.element_size()), swizzle_alignment);
const int smem_kv_pipe_size = num_kv_stages * (smem_kv_size_per_stage + aligned_smem_kv_scale_size_per_stage) + align(num_kv_stages * 8 * 2, swizzle_alignment);
// Allocate some shared memory for UMMA barriers and tensor memory pointer, although it is not used in SM90
const int smem_umma_barriers = num_math_warp_groups * 2 * 8;
const int smem_tmem_ptr = 4;
smem_size = smem_q_pipe_size + num_math_warp_groups * smem_kv_pipe_size + smem_umma_barriers + smem_tmem_ptr;
DG_HOST_ASSERT(smem_size <= SM90ArchSpec::smem_capacity);
DG_HOST_ASSERT(next_n == 1 or next_n == 2);
} else {
const int smem_q_size_per_stage = next_n_atom * num_heads * head_dim * static_cast<int>(q.element_size());
const int smem_kv_size_per_stage = split_kv * head_dim * static_cast<int>(kv_cache.element_size());
const int smem_kv_scale_size_per_stage = split_kv * static_cast<int>(kv_cache_scales.element_size());
const int smem_weight_size_per_stage = next_n_atom * num_heads * static_cast<int>(weights.element_size());
const int smem_barriers = (num_q_stages + num_kv_stages) * 2 * 8;
const int smem_umma_barriers = num_math_warp_groups * 2 * 8;
const int smem_tmem_ptr = 4;
smem_size = num_q_stages * (smem_q_size_per_stage + smem_weight_size_per_stage) +
num_kv_stages * (smem_kv_size_per_stage + smem_kv_scale_size_per_stage) +
smem_barriers + smem_umma_barriers + smem_tmem_ptr;
DG_HOST_ASSERT(smem_size <= SM100ArchSpec::smem_capacity);
}
// Launch
const SMXXFP8PagedMQALogitsRuntime::Args args = {
.batch_size = batch_size,
.next_n = next_n,
.num_heads = num_heads,
.head_dim = head_dim,
.block_kv = block_kv,
.is_context_lens_2d = is_context_lens_2d,
.block_table_stride = block_table_stride,
.logits_stride = logits_stride,
.num_q_stages = num_q_stages,
.num_kv_stages = num_kv_stages,
.split_kv = split_kv,
.context_lens = context_lens.data_ptr<int>(),
.logits = logits.data_ptr(),
.block_table = block_table.data_ptr<int>(),
.schedule_meta = schedule_meta.data_ptr<int>(),
.tensor_map_q = tensor_map_q,
.tensor_map_kv = tensor_map_kv,
.tensor_map_kv_scales = tensor_map_kv_scales,
.tensor_map_weights = tensor_map_weights,
.logits_dtype = logits_dtype,
.num_specialized_threads = num_specialized_threads,
.num_math_threads = num_math_threads,
.launch_args = LaunchArgs(num_sms,
num_specialized_threads + num_math_threads,
smem_size)
};
const auto code = SMXXFP8PagedMQALogitsRuntime::generate(args);
const auto runtime = compiler->build("smxx_fp8_paged_mqa_logits", code);
SMXXFP8PagedMQALogitsRuntime::launch(runtime, args);
}
class SM100FP4PagedMQALogitsRuntime final: public LaunchRuntime<SM100FP4PagedMQALogitsRuntime> {
public:
struct Args {
int batch_size;
int next_n;
int num_heads;
int head_dim;
int block_kv;
bool is_context_lens_2d;
int block_table_stride;
int logits_stride;
int num_q_stages;
int num_kv_stages;
int split_kv;
int* context_lens;
void* logits;
int* block_table;
int* schedule_meta;
CUtensorMap tensor_map_q;
CUtensorMap tensor_map_sf_q;
CUtensorMap tensor_map_kv;
CUtensorMap tensor_map_sf_kv;
CUtensorMap tensor_map_weights;
at::ScalarType logits_dtype;
int num_specialized_threads;
int num_math_threads;
LaunchArgs launch_args;
};
static std::string generate_impl(const Args& args) {
return fmt::format(R"(
#include <deep_gemm/impls/sm100_fp4_paged_mqa_logits.cuh>
using namespace deep_gemm;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&sm100_fp4_paged_mqa_logits<
{}, {},
{}, {},
{},
{}, {},
{},
{}, {},
{}
>);
}};
)", args.next_n, args.num_heads,
args.head_dim, args.block_kv,
args.is_context_lens_2d,
args.num_q_stages, args.num_kv_stages,
args.split_kv,
args.num_specialized_threads, args.num_math_threads,
to_string(args.logits_dtype));
}
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.batch_size,
args.logits_stride, args.block_table_stride,
args.context_lens, args.logits,
args.block_table, args.schedule_meta,
args.tensor_map_q, args.tensor_map_sf_q,
args.tensor_map_kv, args.tensor_map_sf_kv,
args.tensor_map_weights
));
}
};
static void sm100_fp4_paged_mqa_logits(const torch::Tensor& q,
const torch::Tensor& sf_q,
const torch::Tensor& kv_cache,
const torch::Tensor& kv_cache_sf,
const torch::Tensor& weights,
const torch::Tensor& context_lens,
const torch::Tensor& logits,
const torch::Tensor& block_table,
const torch::Tensor& schedule_meta,
const at::ScalarType& logits_dtype,
const int& batch_size, const int& next_n,
const int& num_heads, const int& head_dim,
const int& num_kv_blocks, const int& block_kv,
const bool& is_context_lens_2d,
const int& logits_stride,
const int& block_table_stride,
const int& num_sms,
const int& split_kv) {
const int num_specialized_threads = 128;
const int num_math_threads = 2 * 128;
DG_HOST_ASSERT(split_kv == 256 and logits_stride % split_kv == 0);
// TODO: tuning num_stages
const int num_q_stages = 3, num_kv_stages = 6, num_tmem_stages = 3;
const int next_n_atom = (next_n % 2 == 0) ? 2 : 1;
// `head_dim` must be 128 for 64B swizzling
DG_HOST_ASSERT(head_dim == 128);
// Using 2D TMA as tensor q is asserted contiguous
const auto tensor_map_q = make_tma_2d_desc(q, head_dim, batch_size * next_n * num_heads,
head_dim, next_n_atom * num_heads,
static_cast<int>(q.stride(2)),
head_dim / 2, 0, false, false);
// NOTES: `sf_q` is a 3D tensor, while `weights` is a 2D tensor
const auto tensor_map_sf_q = make_tma_2d_desc(sf_q, num_heads, batch_size * next_n,
num_heads, next_n_atom,
static_cast<int>(sf_q.stride(1)), 0);
const auto tensor_map_weights = make_tma_2d_desc(weights, num_heads, batch_size * next_n,
num_heads, next_n_atom,
static_cast<int>(weights.stride(0)), 0);
const auto tensor_map_kv = make_tma_3d_desc(kv_cache, head_dim, block_kv, num_kv_blocks,
head_dim, block_kv, 1,
static_cast<int>(kv_cache.stride(1)),
static_cast<int>(kv_cache.stride(0)),
head_dim / 2, 0, false, false);
const auto tensor_map_sf_kv = make_tma_2d_desc(kv_cache_sf, block_kv, num_kv_blocks,
block_kv, 1,
static_cast<int>(kv_cache_sf.stride(0)), 0);
// Calculate shared memory size
const int smem_q_size_per_stage = next_n_atom * num_heads * head_dim / 2;
const int smem_sf_q_size_per_stage = align(next_n_atom * num_heads, 128) * sizeof(int);
const int smem_kv_size_per_stage = split_kv * head_dim / 2;
const int smem_sf_kv_size_per_stage = align(split_kv, 128) * sizeof(int);
const int smem_weight_size_per_stage = next_n_atom * num_heads * sizeof(float);
const int smem_barriers = (num_q_stages + num_kv_stages + num_tmem_stages) * 2 * 8;
const int smem_tmem_ptr = 4;
const int smem_size = num_q_stages * (smem_q_size_per_stage + smem_sf_q_size_per_stage + smem_weight_size_per_stage) +
num_kv_stages * (smem_kv_size_per_stage + smem_sf_kv_size_per_stage) +
smem_barriers + smem_tmem_ptr;
DG_HOST_ASSERT(smem_size <= SM100ArchSpec::smem_capacity);
// Launch
const SM100FP4PagedMQALogitsRuntime::Args args = {
.batch_size = batch_size,
.next_n = next_n,
.num_heads = num_heads,
.head_dim = head_dim,
.block_kv = block_kv,
.is_context_lens_2d = is_context_lens_2d,
.block_table_stride = block_table_stride,
.logits_stride = logits_stride,
.num_q_stages = num_q_stages,
.num_kv_stages = num_kv_stages,
.split_kv = split_kv,
.context_lens = context_lens.data_ptr<int>(),
.logits = logits.data_ptr(),
.block_table = block_table.data_ptr<int>(),
.schedule_meta = schedule_meta.data_ptr<int>(),
.tensor_map_q = tensor_map_q,
.tensor_map_sf_q = tensor_map_sf_q,
.tensor_map_kv = tensor_map_kv,
.tensor_map_sf_kv = tensor_map_sf_kv,
.tensor_map_weights = tensor_map_weights,
.logits_dtype = logits_dtype,
.num_specialized_threads = num_specialized_threads,
.num_math_threads = num_math_threads,
.launch_args = LaunchArgs(num_sms,
num_specialized_threads + num_math_threads,
smem_size)
};
const auto code = SM100FP4PagedMQALogitsRuntime::generate(args);
const auto runtime = compiler->build("sm100_fp4_paged_mqa_logits", code);
SM100FP4PagedMQALogitsRuntime::launch(runtime, args);
}
} // namespace deep_gemm

75
csrc/jit_kernels/impls/smxx_layout.hpp
View File

@@ -72,7 +72,7 @@ static void __instantiate_kernel() {{
class PackFP32IntoUE8M0Runtime final: public LaunchRuntime<PackFP32IntoUE8M0Runtime> {
public:
struct Args {
int num_groups, mn, sf_k, packed_sf_k;
int num_groups, mn, sf_k, packed_sf_k, gran_k;
int block_mn, block_packed_sf_k;
void *sf, *out, *ks;
@@ -95,32 +95,32 @@ static void __instantiate_kernel() {{
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.sf, args.out, args.ks, args.mn, args.sf_k, args.packed_sf_k));
args.sf, args.out, args.ks, args.mn, args.sf_k, args.packed_sf_k, args.gran_k));
}
};
static std::tuple<int, int, int, int, int, torch::Tensor> preprocess_sf(const torch::Tensor& sf) {
// NOTES: for the extreme performance, you may rewrite/fuse this function in CUDA
const auto& dim = sf.dim();
const auto dim = sf.dim();
DG_HOST_ASSERT(dim == 2 or dim == 3);
DG_HOST_ASSERT(sf.scalar_type() == torch::kFloat);
const auto& batched_sf = dim == 2 ? sf.unsqueeze(0) : sf;
const auto batched_sf = dim == 2 ? sf.unsqueeze(0) : sf;
const auto& [num_groups, mn, sf_k] = get_shape<3>(batched_sf);
const auto& tma_aligned_mn = get_tma_aligned_size(mn, static_cast<int>(sf.element_size()));
const auto [num_groups, mn, sf_k] = get_shape<3>(batched_sf);
const auto tma_aligned_mn = get_tma_aligned_size(mn, static_cast<int>(sf.element_size()));
return {dim, num_groups, mn, sf_k, tma_aligned_mn, batched_sf};
}
static torch::Tensor get_mn_major_tma_aligned_tensor(const torch::Tensor& sf) {
const auto& [dim, num_groups, mn, sf_k, tma_aligned_mn, batched_sf] = preprocess_sf(sf);
const auto [dim, num_groups, mn, sf_k, tma_aligned_mn, batched_sf] = preprocess_sf(sf);
// The last kernel already gives a column-major TMA aligned layout
if ((batched_sf.stride(0) == tma_aligned_mn * sf_k or dim == 2) and batched_sf.stride(1) == 1 and batched_sf.stride(2) == tma_aligned_mn)
return (dim == 2) ? batched_sf.squeeze(0) : batched_sf;
const auto& out = torch::empty_strided({num_groups, mn, sf_k},
{tma_aligned_mn * sf_k, 1, tma_aligned_mn},
batched_sf.options());
const auto out = torch::empty_strided({num_groups, mn, sf_k},
{tma_aligned_mn * sf_k, 1, tma_aligned_mn},
batched_sf.options());
if (not batched_sf.is_contiguous()) {
// Fallback to PyTorch's slow copy if not contiguous
@@ -129,7 +129,7 @@ static torch::Tensor get_mn_major_tma_aligned_tensor(const torch::Tensor& sf) {
} else {
constexpr int block_mn = 64;
constexpr int num_threads = 512;
const auto& smem_size = block_mn * (sf_k + (1 - (sf_k % 2))) * static_cast<int>(sizeof(float));
const auto smem_size = block_mn * (sf_k + (1 - (sf_k % 2))) * static_cast<int>(sizeof(float));
const TransposeFP32Runtime::Args& args = {
.mn = mn,
.sf_k = sf_k,
@@ -139,25 +139,25 @@ static torch::Tensor get_mn_major_tma_aligned_tensor(const torch::Tensor& sf) {
.launch_args = LaunchArgs({ceil_div(mn, block_mn), num_groups}, num_threads, smem_size)
};
const auto& code = TransposeFP32Runtime::generate(args);
const auto& runtime = compiler->build("transpose_fp32", code);
const auto code = TransposeFP32Runtime::generate(args);
const auto runtime = compiler->build("transpose_fp32", code);
TransposeFP32Runtime::launch(runtime, args);
}
return (dim == 2) ? out.squeeze(0) : out;
}
static torch::Tensor get_mn_major_tma_aligned_packed_ue8m0_tensor_torch(const torch::Tensor& sf) {
const auto& sf_reshaped = (sf.dim() == 2) ? sf.unsqueeze(0) : sf;
const auto sf_reshaped = (sf.dim() == 2) ? sf.unsqueeze(0) : sf;
// First, convert into UE8M0 `uint8_t`
const auto& ue8m0_tensor = sf_reshaped.view(torch::kInt32).bitwise_right_shift(23).to(torch::kUInt8);
const auto ue8m0_tensor = sf_reshaped.view(torch::kInt32).bitwise_right_shift(23).to(torch::kUInt8);
// Second, make padded packed tensors
const auto& [num_groups, mn, k] = get_shape<3>(sf_reshaped);
const auto& aligned_mn = get_tma_aligned_size(mn, 4);
const auto& aligned_k = align(k, 4);
const auto [num_groups, mn, k] = get_shape<3>(sf_reshaped);
const auto aligned_mn = get_tma_aligned_size(mn, 4);
const auto aligned_k = align(k, 4);
const auto& options = torch::TensorOptions().device(sf.device()).dtype(torch::kUInt8);
const auto options = torch::TensorOptions().device(sf.device()).dtype(torch::kUInt8);
auto padded = torch::zeros({num_groups, aligned_mn, aligned_k}, options);
// ReSharper disable once CppExpressionWithoutSideEffects
padded.slice(1, 0, mn).slice(2, 0, k).copy_(ue8m0_tensor);
@@ -172,11 +172,11 @@ static torch::Tensor get_mn_major_tma_aligned_packed_ue8m0_tensor_torch(const to
}
static torch::Tensor get_mn_major_tma_aligned_packed_ue8m0_tensor(const torch::Tensor& sf) {
const auto& [dim, num_groups, mn, sf_k, tma_aligned_mn, batched_sf] = preprocess_sf(sf);
const auto& packed_sf_k = ceil_div(sf_k, 4);
const auto& out = torch::empty_strided({num_groups, mn, packed_sf_k},
{packed_sf_k * tma_aligned_mn, 1, tma_aligned_mn},
at::TensorOptions().device(batched_sf.device()).dtype(torch::kInt));
const auto [dim, num_groups, mn, sf_k, tma_aligned_mn, batched_sf] = preprocess_sf(sf);
const auto packed_sf_k = ceil_div(sf_k, 4);
const auto out = torch::empty_strided({num_groups, mn, packed_sf_k},
{packed_sf_k * tma_aligned_mn, 1, tma_aligned_mn},
at::TensorOptions().device(batched_sf.device()).dtype(torch::kInt));
// Launch the kernel
if (batched_sf.is_contiguous()) {
if ((mn * sf_k) % 4 != 0 and num_groups > 1)
@@ -193,8 +193,8 @@ static torch::Tensor get_mn_major_tma_aligned_packed_ue8m0_tensor(const torch::T
.launch_args = LaunchArgs({ceil_div(mn, block_mn), num_groups}, num_threads, block_mn * sf_k * 4)
};
const auto& code = TransposeAndPackFP32IntoUE8M0Runtime::generate(args);
const auto& runtime = compiler->build("transpose_and_pack_fp32_into_ue8m0", code);
const auto code = TransposeAndPackFP32IntoUE8M0Runtime::generate(args);
const auto runtime = compiler->build("transpose_and_pack_fp32_into_ue8m0", code);
TransposeAndPackFP32IntoUE8M0Runtime::launch(runtime, args);
} else {
if (mn % 4 != 0 or num_groups > 1)
@@ -217,8 +217,8 @@ static torch::Tensor get_mn_major_tma_aligned_packed_ue8m0_tensor(const torch::T
.launch_args = LaunchArgs({ceil_div(mn, block_mn), ceil_div(packed_sf_k, block_packed_sf_k)}, num_threads)
};
const auto& code = PackFP32IntoUE8M0Runtime::generate(args);
const auto& runtime = compiler->build("pack_fp32_into_ue8m0", code);
const auto code = PackFP32IntoUE8M0Runtime::generate(args);
const auto runtime = compiler->build("pack_fp32_into_ue8m0", code);
PackFP32IntoUE8M0Runtime::launch(runtime, args);
}
return (dim == 2) ? out.squeeze(0) : out;
@@ -226,18 +226,20 @@ static torch::Tensor get_mn_major_tma_aligned_packed_ue8m0_tensor(const torch::T
static torch::Tensor get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(const torch::Tensor& sf,
const torch::Tensor& ks_tensor,
const std::vector<int>& ks) {
const auto& [sf_k, mn] = get_shape<2>(sf);
const auto& num_groups = static_cast<int>(ks.size());
const std::vector<int>& ks,
const int gran_k) {
DG_HOST_ASSERT(gran_k == 32 or gran_k == 128);
const auto [sf_k, mn] = get_shape<2>(sf);
const auto num_groups = static_cast<int>(ks.size());
int ref_sf_k = 0, packed_sf_k = 0;
for (const auto& k: ks)
ref_sf_k += ceil_div(k, 128), packed_sf_k += ceil_div(k, 512);
for (const auto k: ks)
ref_sf_k += ceil_div(k, gran_k), packed_sf_k += ceil_div(k, gran_k * 4);
DG_HOST_ASSERT(sf.is_contiguous());
DG_HOST_ASSERT(ref_sf_k == sf_k);
DG_HOST_ASSERT(num_groups <= 128 and mn % 4 == 0);
const auto& out = torch::empty({packed_sf_k, mn}, at::TensorOptions().device(sf.device()).dtype(torch::kInt));
const auto out = torch::empty({packed_sf_k, mn}, at::TensorOptions().device(sf.device()).dtype(torch::kInt));
constexpr int block_mn = 128;
constexpr int block_packed_sf_k = 16;
@@ -247,6 +249,7 @@ static torch::Tensor get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(cons
.mn = mn,
.sf_k = sf_k,
.packed_sf_k = packed_sf_k,
.gran_k = gran_k,
.block_mn = block_mn,
.block_packed_sf_k = block_packed_sf_k,
.sf = sf.data_ptr(),
@@ -255,8 +258,8 @@ static torch::Tensor get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(cons
.launch_args = LaunchArgs({ceil_div(mn, block_mn), ceil_div(packed_sf_k, block_packed_sf_k)}, num_threads)
};
const auto& code = PackFP32IntoUE8M0Runtime::generate(args);
const auto& runtime = compiler->build("pack_fp32_into_ue8m0", code);
const auto code = PackFP32IntoUE8M0Runtime::generate(args);
const auto runtime = compiler->build("pack_fp32_into_ue8m0", code);
PackFP32IntoUE8M0Runtime::launch(runtime, args);
return out;
}