leavelet/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

2
csrc/jit/cache.hpp
View File

@@ -9,7 +9,7 @@
namespace deep_gemm {
class KernelRuntimeCache {
std::unordered_map<std::filesystem::path, std::shared_ptr<KernelRuntime>> cache;
std::unordered_map<std::string, std::shared_ptr<KernelRuntime>> cache;
public:
// TODO: consider cache capacity

149
csrc/jit/compiler.hpp
View File

@@ -1,14 +1,17 @@
#pragma once
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <filesystem>
#include <fstream>
#include <nvrtc.h>
#include <regex>
#include <string>
#include "../utils/exception.hpp"
#include "../utils/format.hpp"
#include "../utils/hash.hpp"
#include "../utils/lazy_init.hpp"
#include "../utils/system.hpp"
#include "cache.hpp"
#include "device_runtime.hpp"
@@ -16,10 +19,13 @@
namespace deep_gemm {
class Compiler {
std::string library_version;
std::filesystem::path library_root_path;
public:
static std::filesystem::path library_root_path;
static std::filesystem::path library_include_path;
static std::filesystem::path cuda_home;
static std::string library_version;
std::string get_library_version() const {
static std::string get_library_version() {
std::stringstream ss;
for (const auto& f: collect_files(library_include_path / "deep_gemm")) {
std::ifstream in(f, std::ios::binary);
@@ -28,16 +34,23 @@ class Compiler {
return get_hex_digest(ss.str());
}
public:
static void prepare_init(const std::string& library_root_path,
const std::string& cuda_home_path_by_torch) {
Compiler::library_root_path = library_root_path;
Compiler::library_include_path = Compiler::library_root_path / "include";
Compiler::cuda_home = cuda_home_path_by_torch;
Compiler::library_version = get_library_version();
}
std::string signature, flags;
std::filesystem::path library_include_path;
std::filesystem::path cache_dir_path;
explicit Compiler(const std::filesystem::path& library_root_path) {
// Static library paths
this->library_root_path = library_root_path;
this->library_include_path = library_root_path / "include";
this->library_version = get_library_version();
Compiler() {
// Check `prepare_init`
DG_HOST_ASSERT(not library_root_path.empty());
DG_HOST_ASSERT(not library_include_path.empty());
DG_HOST_ASSERT(not cuda_home.empty());
DG_HOST_ASSERT(not library_version.empty());
// Cache settings
cache_dir_path = std::filesystem::path(get_env<std::string>("HOME")) / ".deep_gemm";
@@ -46,10 +59,11 @@ public:
// The compiler flags applied to all derived compilers
signature = "unknown-compiler";
std::string ptxas_flags = "--ptxas-options=--register-usage-level=10";
if (get_env<int>("DG_JIT_PTXAS_VERBOSE", 0))
ptxas_flags += ",--verbose";
flags = fmt::format("-std=c++20 --diag-suppress=39,161,174,177,186,940 {}", ptxas_flags);
flags = fmt::format("-std=c++{} --diag-suppress=39,161,174,177,186,940 "
"--ptxas-options=--register-usage-level=10",
get_env<int>("DG_JIT_CPP_STANDARD", 20));
if (get_env("DG_JIT_DEBUG", 0) or get_env("DG_JIT_PTXAS_VERBOSE", 0))
flags += " --ptxas-options=--verbose";
}
virtual ~Compiler() = default;
@@ -102,6 +116,11 @@ public:
virtual void compile(const std::string &code, const std::filesystem::path& dir_path, const std::filesystem::path &cubin_path) const = 0;
};
DG_DECLARE_STATIC_VAR_IN_CLASS(Compiler, library_root_path);
DG_DECLARE_STATIC_VAR_IN_CLASS(Compiler, library_include_path);
DG_DECLARE_STATIC_VAR_IN_CLASS(Compiler, cuda_home);
DG_DECLARE_STATIC_VAR_IN_CLASS(Compiler, library_version);
class NVCCCompiler final: public Compiler {
std::filesystem::path nvcc_path;
@@ -125,11 +144,9 @@ class NVCCCompiler final: public Compiler {
}
public:
NVCCCompiler(const std::filesystem::path& library_root_path,
const std::filesystem::path& cuda_home_path_by_torch):
Compiler(library_root_path) {
NVCCCompiler() {
// Override the compiler signature
nvcc_path = cuda_home_path_by_torch / "bin" / "nvcc";
nvcc_path = cuda_home / "bin" / "nvcc";
if (const auto& env_nvcc_path = get_env<std::string>("DG_JIT_NVCC_COMPILER"); not env_nvcc_path.empty())
nvcc_path = env_nvcc_path;
const auto& [nvcc_major, nvcc_minor] = get_nvcc_version();
@@ -150,10 +167,10 @@ public:
// Compile
const auto& command = fmt::format("{} {} -o {} {}", nvcc_path.c_str(), code_path.c_str(), cubin_path.c_str(), flags);
if (get_env("DG_JIT_DEBUG", 0) or get_env("DG_JIT_PRINT_COMPILER_COMMAND", 0))
printf("Running NVCC command: %s", command.c_str());
printf("Running NVCC command: %s\n", command.c_str());
const auto& [return_code, output] = call_external_command(command);
if (return_code != 0) {
printf("NVCC compilation failed: %s", output.c_str());
printf("NVCC compilation failed: %s\n", output.c_str());
DG_HOST_ASSERT(false and "NVCC compilation failed");
}
@@ -163,6 +180,96 @@ public:
}
};
static std::shared_ptr<Compiler> compiler = nullptr;
class NVRTCCompiler final: public Compiler {
public:
NVRTCCompiler() {
// Override the compiler signature
int major, minor;
DG_NVRTC_CHECK(nvrtcVersion(&major, &minor));
signature = fmt::format("NVRTC{}.{}", major, minor);
// Build include directories list
std::string include_dirs;
include_dirs += fmt::format("-I{} ", library_include_path.string());
include_dirs += fmt::format("-I{} ", (cuda_home / "include").string());
// Add PCH support for version 12.8 and above
// NOTES: PCH is vital for compilation speed
std::string pch_flags;
if (major > 12 or (major == 12 and minor >= 8)) {
pch_flags = "--pch ";
if (get_env<int>("DG_JIT_DEBUG", 0))
pch_flags += "--pch-verbose=true ";
}
// Override the compiler flags
flags = fmt::format("{} {}--gpu-architecture=sm_{}a -default-device {}",
flags, include_dirs, device_runtime->get_arch(), pch_flags);
}
void compile(const std::string &code, const std::filesystem::path& dir_path, const std::filesystem::path &cubin_path) const override {
// Write the code into the cache directory
const auto& code_path = dir_path / "kernel.cu";
put(code_path, code);
// Parse compilation options
std::istringstream iss(flags);
std::vector<std::string> options;
std::string option;
while (iss >> option)
options.push_back(option);
// Convert to C-style string array for NVRTC
std::vector<const char*> option_cstrs;
for (const auto& opt: options)
option_cstrs.push_back(opt.c_str());
// Print compiler command if requested
if (get_env<int>("DG_JIT_DEBUG", 0) or get_env<int>("DG_JIT_PRINT_COMPILER_COMMAND", 0)) {
printf("Compiling JIT runtime with NVRTC options: ");
for (const auto& opt: options)
printf("%s ", opt.c_str());
printf("\n");
}
// Create NVRTC program and compile
nvrtcProgram program;
DG_NVRTC_CHECK(nvrtcCreateProgram(&program, code.c_str(), "kernel.cu", 0, nullptr, nullptr));
const auto& compile_result = nvrtcCompileProgram(program, static_cast<int>(option_cstrs.size()), option_cstrs.data());
// Get and print compiler log
size_t log_size;
DG_NVRTC_CHECK(nvrtcGetProgramLogSize(program, &log_size));
if (get_env<int>("DG_JIT_DEBUG", 0) or compile_result != NVRTC_SUCCESS) {
if (compile_result != NVRTC_SUCCESS)
DG_HOST_ASSERT(log_size > 1);
if (log_size > 1) {
std::string compilation_log(log_size, '\0');
DG_NVRTC_CHECK(nvrtcGetProgramLog(program, compilation_log.data()));
printf("NVRTC log: %s\n", compilation_log.c_str());
}
}
// Get CUBIN size and data
size_t cubin_size;
DG_NVRTC_CHECK(nvrtcGetCUBINSize(program, &cubin_size));
std::string cubin_data(cubin_size, '\0');
DG_NVRTC_CHECK(nvrtcGetCUBIN(program, cubin_data.data()));
// Write into the file system
put(cubin_path, cubin_data);
// Cleanup
DG_NVRTC_CHECK(nvrtcDestroyProgram(&program));
}
};
static auto compiler = LazyInit<Compiler>([]() -> std::shared_ptr<Compiler> {
if (get_env<int>("DG_JIT_USE_NVRTC", 0)) {
return std::make_shared<NVRTCCompiler>();
} else {
return std::make_shared<NVCCCompiler>();
}
});
} // namespace deep_gemm

14
csrc/jit/device_runtime.hpp
View File

@@ -3,11 +3,12 @@
#include <ATen/cuda/CUDAContext.h>
#include "../utils/exception.hpp"
#include "../utils/lazy_init.hpp"
namespace deep_gemm {
class DeviceRuntime {
int num_sms = 0;
int num_sms = 0, tc_util = 0;
std::shared_ptr<cudaDeviceProp> cached_prop;
public:
@@ -43,8 +44,17 @@ public:
num_sms = get_prop()->multiProcessorCount;
return num_sms;
}
void set_tc_util(const int& new_tc_util) {
DG_HOST_ASSERT(0 <= new_tc_util and new_tc_util <= 100);
tc_util = new_tc_util;
}
int get_tc_util() const {
return tc_util == 0 ? 100 : tc_util;
}
};
static auto device_runtime = std::make_shared<DeviceRuntime>();
static auto device_runtime = LazyInit<DeviceRuntime>([](){ return std::make_shared<DeviceRuntime>(); });
} // namespace deep_gemm

135
csrc/jit/handle.hpp Normal file
View File

@@ -0,0 +1,135 @@
#pragma once
#include <cuda.h>
#include <cuda_runtime.h>
#include <filesystem>
#include "../utils/exception.hpp"
namespace deep_gemm {
#if CUDART_VERSION >= 12080 or defined(DG_JIT_USE_DRIVER_API)
// Use CUDA runtime API
using LibraryHandle = cudaLibrary_t;
using KernelHandle = cudaKernel_t;
using LaunchConfigHandle = cudaLaunchConfig_t;
using LaunchAttrHandle = cudaLaunchAttribute;
#define DG_CUDA_UNIFIED_CHECK DG_CUDA_RUNTIME_CHECK
static KernelHandle load_kernel(const std::filesystem::path& cubin_path, const std::string& func_name,
LibraryHandle *library_opt = nullptr) {
LibraryHandle library;
KernelHandle kernel{};
DG_CUDA_RUNTIME_CHECK(cudaLibraryLoadFromFile(&library, cubin_path.c_str(), nullptr, nullptr, 0, nullptr, nullptr, 0));
DG_CUDA_RUNTIME_CHECK(cudaLibraryGetKernel(&kernel, library, func_name.c_str()));
if (library_opt != nullptr)
*library_opt = library;
return kernel;
}
static void unload_library(const LibraryHandle& library) {
const auto& error = cudaLibraryUnload(library);
DG_HOST_ASSERT(error == cudaSuccess or error == cudaErrorCudartUnloading);
}
static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel,
const cudaStream_t& stream, const int& smem_size,
const dim3& grid_dim, const dim3& block_dim, const int& cluster_dim) {
if (smem_size > 0)
DG_CUDA_RUNTIME_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
LaunchConfigHandle config;
config.gridDim = grid_dim;
config.blockDim = block_dim;
config.dynamicSmemBytes = smem_size;
config.stream = stream;
config.numAttrs = 0;
config.attrs = nullptr;
// NOTES: must use `static` or the `attr` will be deconstructed
static LaunchAttrHandle attr;
if (cluster_dim > 1) {
attr.id = cudaLaunchAttributeClusterDimension;
attr.val.clusterDim = {static_cast<unsigned>(cluster_dim), 1, 1};
config.attrs = &attr;
config.numAttrs = 1;
}
return config;
}
template<typename... ActTypes>
static auto launch_kernel(const KernelHandle& kernel, const LaunchConfigHandle& config, ActTypes&&... args) {
void *ptr_args[] = { &args... };
return cudaLaunchKernelExC(&config, kernel, ptr_args);
}
#else
// Use CUDA driver API
using LibraryHandle = CUmodule;
using KernelHandle = CUfunction;
using LaunchConfigHandle = CUlaunchConfig;
using LaunchAttrHandle = CUlaunchAttribute;
#define DG_CUDA_UNIFIED_CHECK DG_CUDA_DRIVER_CHECK
static KernelHandle load_kernel(const std::filesystem::path& cubin_path, const std::string& func_name,
LibraryHandle *library_opt = nullptr) {
LibraryHandle library;
KernelHandle kernel;
DG_CUDA_DRIVER_CHECK(cuModuleLoad(&library, cubin_path.c_str()));
DG_CUDA_DRIVER_CHECK(cuModuleGetFunction(&kernel, library, func_name.c_str()));
if (library_opt != nullptr)
*library_opt = library;
return kernel;
}
static void unload_library(const LibraryHandle& library) {
const auto& error = cuModuleUnload(library);
DG_HOST_ASSERT(error == CUDA_SUCCESS or error == CUDA_ERROR_DEINITIALIZED);
}
static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel,
const cudaStream_t& stream, const int& smem_size,
const dim3& grid_dim, const dim3& block_dim, const int& cluster_dim) {
if (smem_size > 0)
DG_CUDA_DRIVER_CHECK(cuFuncSetAttribute(kernel, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem_size));
LaunchConfigHandle config;
config.gridDimX = grid_dim.x;
config.gridDimY = grid_dim.y;
config.gridDimZ = grid_dim.z;
config.blockDimX = block_dim.x;
config.blockDimY = block_dim.y;
config.blockDimZ = block_dim.z;
config.sharedMemBytes = smem_size;
config.hStream = stream;
config.numAttrs = 0;
config.attrs = nullptr;
// NOTES: must use `static` or the `attr` will be deconstructed
static LaunchAttrHandle attr;
if (cluster_dim > 1) {
attr.id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION;
attr.value.clusterDim.x = cluster_dim;
attr.value.clusterDim.y = 1;
attr.value.clusterDim.z = 1;
config.attrs = &attr;
config.numAttrs = 1;
}
return config;
}
template<typename... ActTypes>
static auto launch_kernel(const KernelHandle& kernel, const LaunchConfigHandle& config, ActTypes&&... args) {
void *ptr_args[] = { &args... };
return cuLaunchKernelEx(&config, kernel, ptr_args, nullptr);
}
#endif
} // namespace deep_gemm

62
csrc/jit/kernel_runtime.hpp
View File

@@ -1,12 +1,10 @@
#pragma once
#include <cuda_runtime.h>
#include <filesystem>
#include "../utils/exception.hpp"
#include "../utils/format.hpp"
#include "../utils/system.hpp"
#include "device_runtime.hpp"
#include "handle.hpp"
namespace deep_gemm {
@@ -23,19 +21,17 @@ struct LaunchArgs {
grid_dim(grid_dim), num_threads(num_threads), smem_size(smem_size), cluster_dim(cluster_dim) {}
};
template <typename T>
concept HasLaunchArgs = requires (const T& t) {
{ t.launch_args } -> std::convertible_to<decltype(t.launch_args)>;
};
class KernelRuntime final {
public:
static std::filesystem::path cuda_home;
cudaLibrary_t library;
cudaKernel_t kernel;
LibraryHandle library;
KernelHandle kernel;
explicit KernelRuntime(const std::filesystem::path& dir_path) {
// Check `prepare_init`
DG_HOST_ASSERT(not cuda_home.empty());
// NOLINT(*-pro-type-member-init)
const auto& cuobjdump_path = cuda_home / "bin" / "cuobjdump";
const auto& cubin_path = dir_path / "kernel.cubin";
@@ -50,7 +46,8 @@ public:
std::istringstream iss(symbols);
std::vector<std::string> symbol_names;
for (std::string line; std::getline(iss, line); ) {
if (line.find("STT_FUNC") == 0 and std::ranges::none_of(illegal_names, [&](const auto& name) { return line.find(name) != std::string::npos; })) {
if (line.find("STT_FUNC") == 0 and std::none_of(illegal_names.begin(), illegal_names.end(),
[&](const auto& name) { return line.find(name) != std::string::npos; })) {
const auto& last_space = line.rfind(' ');
symbol_names.push_back(line.substr(last_space + 1));
}
@@ -64,11 +61,10 @@ public:
// Load from the library
DG_HOST_ASSERT(symbol_names.size() == 1);
DG_CUDA_RUNTIME_CHECK(cudaLibraryLoadFromFile(&library, cubin_path.c_str(), nullptr, nullptr, 0, nullptr, nullptr, 0));
DG_CUDA_RUNTIME_CHECK(cudaLibraryGetKernel(&kernel, library, symbol_names[0].c_str()));
kernel = load_kernel(cubin_path, symbol_names[0], &library);
}
static void set_cuda_home(const std::string& cuda_home_path_by_torch) {
static void prepare_init(const std::string& cuda_home_path_by_torch) {
cuda_home = cuda_home_path_by_torch;
}
@@ -78,18 +74,16 @@ public:
}
~KernelRuntime() noexcept(false) {
const auto& error = cudaLibraryUnload(library);
DG_HOST_ASSERT(error == cudaSuccess or error == cudaErrorCudartUnloading);
unload_library(library);
}
};
// Declare after defining
decltype(KernelRuntime::cuda_home) KernelRuntime::cuda_home;
DG_DECLARE_STATIC_VAR_IN_CLASS(KernelRuntime, cuda_home);
template <typename Derived>
class LaunchRuntime {
public:
template <typename Args> requires HasLaunchArgs<Args>
template <typename Args>
static std::string generate(const Args& args) {
const auto& code = Derived::generate_impl(args);
if (get_env<int>("DG_JIT_DEBUG", 0))
@@ -97,34 +91,18 @@ public:
return code;
}
template <typename Args> requires HasLaunchArgs<Args>
template <typename Args>
static void launch(const std::shared_ptr<KernelRuntime>& kernel_runtime, const Args& args) {
const auto& kernel = kernel_runtime->kernel;
const auto& stream = at::cuda::getCurrentCUDAStream();
const LaunchArgs& launch_args = args.launch_args;
// Set dynamic shared memory size
if (launch_args.smem_size > 0)
DG_CUDA_RUNTIME_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, launch_args.smem_size));
// Launch config
cudaLaunchConfig_t config;
config.gridDim = {static_cast<unsigned>(launch_args.grid_dim.first),
static_cast<unsigned>(launch_args.grid_dim.second),
1};
config.blockDim = {static_cast<unsigned>(launch_args.num_threads), 1, 1};
config.dynamicSmemBytes = launch_args.smem_size;
config.stream = stream;
config.numAttrs = 0;
// Clusters
cudaLaunchAttribute attr;
if (launch_args.cluster_dim > 1) {
attr.id = cudaLaunchAttributeClusterDimension;
attr.val.clusterDim = {static_cast<unsigned>(launch_args.cluster_dim), 1, 1};
config.attrs = &attr;
config.numAttrs = 1;
}
const dim3& grid_dim = {static_cast<unsigned>(launch_args.grid_dim.first),
static_cast<unsigned>(launch_args.grid_dim.second),
1};
const dim3& block_dim = {static_cast<unsigned>(launch_args.num_threads), 1, 1};
auto config = construct_launch_config(kernel, stream, launch_args.smem_size,
grid_dim, block_dim, launch_args.cluster_dim);
// Launch in the derived class
if (get_env<int>("DG_JIT_DEBUG")) {

16
csrc/jit_kernels/heuristics/common.hpp
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);
}
}

9
csrc/jit_kernels/heuristics/sm100.hpp
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,

24
csrc/jit_kernels/impls/sm100_fp8_gemm_1d1d.hpp
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()) {

16
csrc/jit_kernels/impls/sm100_fp8_gemm_1d2d.hpp
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,

15
csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp
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,

57
csrc/jit_kernels/impls/smxx_layout.hpp
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);

44
csrc/python_api.cpp
View File

@@ -5,10 +5,10 @@
#include "jit/device_runtime.hpp"
#include "utils/layout.hpp"
#include "jit_kernels/impls/smxx_layout.hpp"
#include "jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp"
#include "jit_kernels/impls/sm100_fp8_gemm_1d1d.hpp"
#include "jit_kernels/impls/sm100_fp8_gemm_1d2d.hpp"
#include "jit_kernels/impls/smxx_layout.hpp"
#ifndef TORCH_EXTENSION_NAME
#define TORCH_EXTENSION_NAME deep_gemm_cpp
@@ -17,8 +17,8 @@
namespace deep_gemm {
torch::Tensor transform_sf_into_required_layout(const torch::Tensor& sf,
const int& mn, const int& k,
const std::optional<int>& num_groups,
const std::tuple<int, int, int>& recipe,
const std::optional<int>& num_groups,
const bool& is_sfa,
const bool& disable_ue8m0_cast) {
const auto& gran_mn = is_sfa ? std::get<0>(recipe) : std::get<1>(recipe);
@@ -121,8 +121,8 @@ void fp8_gemm_nt(const std::pair<torch::Tensor, torch::Tensor>& a,
// Transform SFA and SFB into compute-required layout
if (not recipe.has_value())
recipe = get_default_recipe(a.second.scalar_type(), b.second.scalar_type());
const auto& sfa = transform_sf_into_required_layout(a.second, m, k, std::nullopt, recipe.value(), true, disable_ue8m0_cast);
const auto& sfb = transform_sf_into_required_layout(b.second, n, k, std::nullopt, recipe.value(), false, disable_ue8m0_cast);
const auto& sfa = transform_sf_into_required_layout(a.second, m, k, recipe.value(), std::nullopt, true, disable_ue8m0_cast);
const auto& sfb = transform_sf_into_required_layout(b.second, n, k, recipe.value(), std::nullopt, false, disable_ue8m0_cast);
// Dispatch into different implements
const auto& arch_major = device_runtime->get_arch_major();
@@ -133,7 +133,7 @@ void fp8_gemm_nt(const std::pair<torch::Tensor, torch::Tensor>& a,
} else if (arch_major == 10 and sfa.scalar_type() == torch::kFloat) {
sm100_fp8_gemm_1d2d(a.first, sfa, b.first, sfb, c, d, m, n, k, major_a, major_b, compiled_dims);
} else {
DG_HOST_UNREACHABLE("Unknown kernel or scaling factor types");
DG_HOST_UNREACHABLE("Unsupported architecture or scaling factor types");
}
}
@@ -208,8 +208,8 @@ void m_grouped_fp8_gemm_nt_contiguous(const std::pair<torch::Tensor, torch::Tens
// Transform SFA and SFB into compute-required layout
if (not recipe.has_value())
recipe = get_default_recipe(a.second.scalar_type(), b.second.scalar_type());
const auto& sfa = transform_sf_into_required_layout(a.second, m, k, std::nullopt, recipe.value(), true, disable_ue8m0_cast);
const auto& sfb = transform_sf_into_required_layout(b.second, n, k, num_groups, recipe.value(), false, disable_ue8m0_cast);
const auto& sfa = transform_sf_into_required_layout(a.second, m, k, recipe.value(), std::nullopt, true, disable_ue8m0_cast);
const auto& sfb = transform_sf_into_required_layout(b.second, n, k, recipe.value(), num_groups, false, disable_ue8m0_cast);
// Dispatch implementation
const auto& arch_major = device_runtime->get_arch_major();
@@ -223,7 +223,7 @@ void m_grouped_fp8_gemm_nt_contiguous(const std::pair<torch::Tensor, torch::Tens
sm100_m_grouped_fp8_gemm_contiguous_1d2d(a.first, sfa, b.first, sfb, d, m_indices,
num_groups, m, n, k, major_a, major_b, compiled_dims);
} else {
DG_HOST_UNREACHABLE("Unknown kernel or scaling factor types");
DG_HOST_UNREACHABLE("Unsupported architecture or scaling factor types");
}
}
@@ -271,8 +271,8 @@ void fp8_m_grouped_gemm_nt_masked(const std::pair<torch::Tensor, torch::Tensor>&
// Transform scaling factors
if (not recipe.has_value())
recipe = get_default_recipe(a.second.scalar_type(), b.second.scalar_type());
const auto& sfa = transform_sf_into_required_layout(a.second, m, k, num_groups, recipe.value(), true, disable_ue8m0_cast);
const auto& sfb = transform_sf_into_required_layout(b.second, n, k, num_groups, recipe.value(), false, disable_ue8m0_cast);
const auto& sfa = transform_sf_into_required_layout(a.second, m, k, recipe.value(), num_groups, true, disable_ue8m0_cast);
const auto& sfb = transform_sf_into_required_layout(b.second, n, k, recipe.value(), num_groups, false, disable_ue8m0_cast);
// Dispatch implementation
const auto& arch_major = device_runtime->get_arch_major();
@@ -286,7 +286,7 @@ void fp8_m_grouped_gemm_nt_masked(const std::pair<torch::Tensor, torch::Tensor>&
sm100_fp8_m_grouped_gemm_masked_1d2d(a.first, sfa, b.first, sfb, d, masked_m,
num_groups, m, n, k, expected_m, major_a, major_b, compiled_dims);
} else {
DG_HOST_UNREACHABLE("Unsupported kernel or scaling factor types");
DG_HOST_UNREACHABLE("Unsupported architecture or scaling factor types");
}
}
@@ -339,18 +339,23 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.doc() = "DeepGEMM C++ library";
// Runtime
m.def("set_num_sms", [&](const int& new_num_sms) {
device_runtime->set_num_sms(new_num_sms);
});
m.def("get_num_sms", [&]() {
return device_runtime->get_num_sms();
});
m.def("set_num_sms", [&](const int& new_num_sms) {
device_runtime->set_num_sms(new_num_sms);
m.def("set_tc_util", [&](const int& new_tc_util) {
device_runtime->set_tc_util(new_tc_util);
});
m.def("get_tc_util", [&]() {
return device_runtime->get_tc_util();
});
// JIT
m.def("init", [&](const std::string& library_root_path, const std::string& cuda_home_path_by_torch) {
DG_HOST_ASSERT(get_env("DG_JIT_USE_NVRTC", 0) == 0 and "Currently only support NVCC");
compiler = std::make_shared<NVCCCompiler>(library_root_path, cuda_home_path_by_torch);
KernelRuntime::set_cuda_home(cuda_home_path_by_torch);
Compiler::prepare_init(library_root_path, cuda_home_path_by_torch);
KernelRuntime::prepare_init(cuda_home_path_by_torch);
});
// Stable kernel APIs with automatic arch/layout dispatch
@@ -391,7 +396,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
py::arg("ks_tensor"), py::arg("c") = std::nullopt,
py::arg("recipe") = std::make_tuple(1, 1, 128),
py::arg("compiled_dims") = "mn");
m.def("transform_sf_into_required_layout", &transform_sf_into_required_layout);
// Layout kernels
m.def("transform_sf_into_required_layout", &transform_sf_into_required_layout,
py::arg("sf"), py::arg("mn"), py::arg("k"), py::arg("recipe"),
py::arg("num_groups") = std::nullopt, py::arg("is_sfa") = false,
py::arg("disable_ue8m0_cast") = false);
// Raw kernels or functions
m.def("get_tma_aligned_size", &get_tma_aligned_size);

10
csrc/utils/exception.hpp
View File

@@ -35,6 +35,16 @@ do { \
#define DG_HOST_UNREACHABLE(reason) (throw DGException("Assertion", __FILE__, __LINE__, reason))
#endif
#ifndef DG_NVRTC_CHECK
#define DG_NVRTC_CHECK(cmd) \
do { \
const auto& e = (cmd); \
if (e != NVRTC_SUCCESS) { \
throw DGException("NVRTC", __FILE__, __LINE__, nvrtcGetErrorString(e)); \
} \
} while (0)
#endif
#ifndef DG_CUDA_DRIVER_CHECK
#define DG_CUDA_DRIVER_CHECK(cmd) \
do { \

27
csrc/utils/lazy_init.hpp Normal file
View File

@@ -0,0 +1,27 @@
#pragma once
#include <functional>
#include <memory>
#define DG_DECLARE_STATIC_VAR_IN_CLASS(cls, name) decltype(cls::name) cls::name
namespace deep_gemm {
template <typename T>
class LazyInit {
public:
explicit LazyInit(std::function<std::shared_ptr<T>()> factory)
: factory(std::move(factory)) {}
T* operator -> () {
if (ptr == nullptr)
ptr = factory();
return ptr.get();
}
private:
std::shared_ptr<T> ptr;
std::function<std::shared_ptr<T>()> factory;
};
} // namespace deep_gemm

4
csrc/utils/system.hpp
View File

@@ -38,8 +38,8 @@ static std::tuple<int, std::string> call_external_command(std::string command) {
std::string output;
while (fgets(buffer.data(), buffer.size(), pipe.get()))
output += buffer.data();
const auto exit_code = pclose(pipe.release());
return {WEXITSTATUS(exit_code), output};
const auto& exit_code = WEXITSTATUS(pclose(pipe.release()));
return {exit_code, output};
}
static std::vector<std::filesystem::path> collect_files(const std::filesystem::path& root) {