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

Make various updates and fixes:

Browse Source 
- Add support for legacy CUDA versions; now compatible with CUDA 12.3 and newer
- Add support for NVRTC compilation
- Other fixes and code refactoring
This commit is contained in:
Ray Wang
2025-08-02 19:52:22 -07:00
parent aff9da0aba
commit d9c363f86f
36 changed files with 592 additions and 362 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/jit_kernels/heuristics/common.hpp
+11 -5
View File
@@ -62,8 +62,9 @@ struct GemmConfig {
int block_m, block_n, block_k;
int num_stages, num_last_stages;
// Runtime configs
// Templated device configs
int num_sms;
int tc_util;
// Structured configs
MulticastConfig multicast_config;
@@ -265,30 +266,35 @@ static GemmConfig get_best_config(const GemmType& gemm_type, const KernelType& 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)
};
// Only SM100 BF16 kernels support tensor core control
if (config.tc_util < 100)
DG_HOST_ASSERT(device_runtime->get_arch_major() == 10 and ab_dtype == torch::kBFloat16);
// 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,
ab_dtype, cd_dtype, with_accumulation, num_sms);
static std::set<decltype(key)> printed;
if (not printed.contains(key)) {
printf("Gemm type: %d, kernel type: %d, M: %d, N: %d, K: %d, groups: %d, "
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, AB 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, threads: %d\n",
"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), c10::toString(ab_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.thread_config.num_threads);
best_smem_config.swizzle_cd_mode, config.num_sms, config.thread_config.num_threads, config.tc_util);
printed.insert(key);
}
}
csrc/jit_kernels/heuristics/sm100.hpp
+7 -2
View File
@@ -44,6 +44,11 @@ struct SM100ArchSpec {
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const at::ScalarType& ab_dtype, const at::ScalarType& cd_dtype,
const int& block_m, const int& block_n) {
// TODO: consider more carefully for BF16 GEMMs
// 2SM BF16 UMMA does not support `N % 32 != 0`
if (ab_dtype == torch::kBFloat16 and block_n % 32 != 0)
return false;
// Layout A/D does not support `block_m == 64` and `block_n % 16 != 0`
if (block_m == 64 or block_n % 16 != 0)
return false;
@@ -68,7 +73,7 @@ struct SM100ArchSpec {
// 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 % 64 == 0;
return major_b == cute::UMMA::Major::K or (block_n * c10::elementSize(ab_dtype)) % 64 == 0;
}
static bool is_num_stages_legal(const at::ScalarType& ab_dtype, const at::ScalarType& cd_dtype,
@@ -93,7 +98,7 @@ struct SM100ArchSpec {
static ThreadConfig get_thread_config(const KernelType& kernel_type,
const int& block_m, const int& block_n) {
return ThreadConfig::sm100(128, kernel_type == KernelType::Kernel1D1D ? 128 : block_m);
return ThreadConfig::sm100(128, kernel_type == KernelType::Kernel1D2D ? block_m : 128);
}
static int get_smem_cd_size(const KernelType& kernel_type,
csrc/jit_kernels/impls/sm100_fp8_gemm_1d1d.hpp
+8 -16
View File
@@ -32,13 +32,6 @@ public:
static std::string generate_impl(const Args& args) {
return fmt::format(R"(
#ifdef __CUDACC_RTC__
#include <deep_gemm/nvrtc_std.cuh>
#else
#include <cuda.h>
#include <string>
#endif
#include <deep_gemm/impls/sm100_fp8_gemm_1d1d.cuh>
using namespace deep_gemm;
@@ -54,7 +47,7 @@ static void __instantiate_kernel() {{
{}, {},
{}, {},
{},
{}, {}
{}, {}, {}
>);
}};
)",
@@ -66,14 +59,13 @@ static void __instantiate_kernel() {{
args.gemm_config.num_stages, args.gemm_config.num_last_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,
to_string(args.gemm_config.gemm_type),
args.gemm_config.with_accumulation,
to_string(args.gemm_config.cd_dtype));
args.gemm_config.num_sms,
to_string(args.gemm_config.gemm_type), args.gemm_config.with_accumulation, to_string(args.gemm_config.cd_dtype));
}
static void launch_impl(const cudaKernel_t& kernel, const cudaLaunchConfig_t& config, Args args) {
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
// TODO: optimize `args` copy
DG_CUDA_RUNTIME_CHECK(cudaLaunchKernelEx(&config, kernel,
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.grouped_layout, args.m, args.n, args.k,
args.tensor_map_a, args.tensor_map_b,
args.tensor_map_sfa, args.tensor_map_sfb,
@@ -286,7 +278,7 @@ static void fp8_k_grouped_gemm_1d1d(const torch::Tensor& a, const torch::Tensor&
const auto& num_groups = static_cast<int>(ks.size());
// Get config using max K for better performance
const auto& max_k = *std::ranges::max_element(ks);
const auto& max_k = *std::max_element(ks.begin(), ks.end());
const auto& config = get_best_config<SM100ArchSpec>(
GemmType::KGroupedContiguous, KernelType::Kernel1D1D,
m, n, max_k, num_groups, cute::UMMA::Major::MN, cute::UMMA::Major::MN,
@@ -316,9 +308,9 @@ static void fp8_k_grouped_gemm_1d1d(const torch::Tensor& a, const torch::Tensor&
static_cast<int>(cd.stride(1)), num_groups,
config.smem_config.swizzle_cd_mode);
const auto& tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, sum_sf_k * 512,
config.block_m, config.block_k, num_groups, 0);
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 * 512,
config.block_n, config.block_k, num_groups, 0);
config.block_n, config.block_k, 1, 0);
// Duplicate the accumulator if necessary
if (c.has_value()) {
csrc/jit_kernels/impls/sm100_fp8_gemm_1d2d.hpp
+5 -11
View File
@@ -30,13 +30,6 @@ public:
static std::string generate_impl(const Args& args) {
return fmt::format(R"(
#ifdef __CUDACC_RTC__
#include <deep_gemm/nvrtc_std.cuh>
#else
#include <cuda.h>
#include <string>
#endif
#include <deep_gemm/impls/sm100_fp8_gemm_1d2d.cuh>
using namespace deep_gemm;
@@ -51,6 +44,7 @@ static void __instantiate_kernel() {{
{}, {},
{}, {},
{}, {},
{},
{}, {}
>);
}};
@@ -63,13 +57,13 @@ static void __instantiate_kernel() {{
args.gemm_config.num_stages, args.gemm_config.num_last_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,
to_string(args.gemm_config.gemm_type),
to_string(args.gemm_config.cd_dtype));
args.gemm_config.num_sms,
to_string(args.gemm_config.gemm_type), to_string(args.gemm_config.cd_dtype));
}
static void launch_impl(const cudaKernel_t& kernel, const cudaLaunchConfig_t& config, Args args) {
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
// TODO: optimize `args` copy
DG_CUDA_RUNTIME_CHECK(cudaLaunchKernelEx(&config, kernel,
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.sfb, args.grouped_layout,
args.m, args.n, args.k,
args.tensor_map_a, args.tensor_map_b,
csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp
+4 -11
View File
@@ -29,13 +29,6 @@ public:
static std::string generate_impl(const Args& args) {
return fmt::format(R"(
#ifdef __CUDACC_RTC__
#include <deep_gemm/nvrtc_std.cuh>
#else
#include <cuda.h>
#include <string>
#endif
#include <deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh>
using namespace deep_gemm;
@@ -49,7 +42,7 @@ static void __instantiate_kernel() {{
{}, {},
{}, {},
{}, {},
{}
{}, {}
>);
}};
)",
@@ -61,12 +54,12 @@ static void __instantiate_kernel() {{
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,
to_string(args.gemm_config.gemm_type));
args.gemm_config.num_sms, to_string(args.gemm_config.gemm_type));
}
static void launch_impl(const cudaKernel_t& kernel, const cudaLaunchConfig_t& config, Args args) {
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
// TODO: optimize `args` copy
DG_CUDA_RUNTIME_CHECK(cudaLaunchKernelEx(&config, kernel,
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.sfb, args.grouped_layout,
args.m, args.n, args.k,
args.tensor_map_a, args.tensor_map_b,
csrc/jit_kernels/impls/smxx_layout.hpp
+6 -51
View File
@@ -22,13 +22,6 @@ public:
static std::string generate_impl(const Args& args) {
return fmt::format(R"(
#ifdef __CUDACC_RTC__
#include <deep_gemm/nvrtc_std.cuh>
#else
#include <cuda.h>
#include <string>
#endif
#include <deep_gemm/impls/smxx_layout.cuh>
using namespace deep_gemm;
@@ -41,8 +34,8 @@ static void __instantiate_kernel() {{
)", args.launch_args.num_threads, args.block_mn, args.sf_k);
}
static void launch_impl(const cudaKernel_t& kernel, const cudaLaunchConfig_t& config, Args args) {
DG_CUDA_RUNTIME_CHECK(cudaLaunchKernelEx(&config, kernel, args.sf, args.out, static_cast<uint32_t>(args.mn)));
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, args.sf, args.out, static_cast<uint32_t>(args.mn)));
}
};
@@ -58,13 +51,6 @@ public:
static std::string generate_impl(const Args& args) {
return fmt::format(R"(
#ifdef __CUDACC_RTC__
#include <deep_gemm/nvrtc_std.cuh>
#else
#include <cuda.h>
#include <string>
#endif
#include <deep_gemm/impls/smxx_layout.cuh>
using namespace deep_gemm;
@@ -77,8 +63,8 @@ static void __instantiate_kernel() {{
)", args.num_groups, args.launch_args.num_threads, args.block_mn, args.block_packed_sf_k);
}
static void launch_impl(const cudaKernel_t& kernel, const cudaLaunchConfig_t& config, Args args) {
DG_CUDA_RUNTIME_CHECK(cudaLaunchKernelEx(&config, 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));
}
};
@@ -108,43 +94,16 @@ static torch::Tensor get_mn_major_tma_aligned_tensor(const torch::Tensor& sf) {
return (dim == 2) ? aligned_sf.squeeze(0) : aligned_sf;
}
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;
// First, convert into UE8M0 `uint8_t`
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& 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);
padded = padded.view(-1).view(torch::kInt32).view({num_groups, aligned_mn, aligned_k / 4});
// Finally, transpose
auto out = torch::empty_strided({num_groups, aligned_mn, aligned_k / 4},
{aligned_mn * (aligned_k / 4), 1, aligned_mn},
at::TensorOptions().device(sf.device()).dtype(torch::kInt32));
out = out.copy_(padded).slice(1, 0, mn);
return (sf.dim() == 2) ? out.squeeze(0) : out;
}
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));
DG_HOST_ASSERT(num_groups == 1 or (mn * sf_k) % 4 == 0);
// Launch the kernel
if (batched_sf.is_contiguous()) {
// Fallback to slow PyTorch impl for non-supported cases
if ((mn * sf_k) % 4 != 0 and num_groups > 1)
return get_mn_major_tma_aligned_packed_ue8m0_tensor_torch(sf);
constexpr int block_mn = 48;
constexpr int num_threads = 512;
const TransposeAndPackFP32IntoUE8M0Runtime::Args& args = {
@@ -160,10 +119,6 @@ static torch::Tensor get_mn_major_tma_aligned_packed_ue8m0_tensor(const torch::T
const auto& runtime = compiler->build("transpose_and_pack_fp32_into_ue8m0", code);
TransposeAndPackFP32IntoUE8M0Runtime::launch(runtime, args);
} else {
// Fallback to slow PyTorch impl for non-supported cases
if (mn % 4 != 0 or num_groups > 1)
return get_mn_major_tma_aligned_packed_ue8m0_tensor_torch(sf);
DG_HOST_ASSERT(mn % 4 == 0 and num_groups == 1);
DG_HOST_ASSERT(batched_sf.stride(1) == 1 and batched_sf.stride(2) == mn);