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/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);