[Public release 26/04] Introducing Mega MoE, FP4 Indexer and other features/fixes (#304)
* Merge with private repo * Update README * Update README * Update README * Add PyTorch requirements * Fix sync scopes for MQA logits (#256) * Update README
This commit is contained in:
Chenggang Zhao
+1
-1
@@ -17,7 +17,7 @@ public:
|
||||
|
||||
std::shared_ptr<KernelRuntime> get(const std::filesystem::path& dir_path) {
|
||||
// Hit the runtime cache
|
||||
if (const auto& iterator = cache.find(dir_path); iterator != cache.end())
|
||||
if (const auto iterator = cache.find(dir_path); iterator != cache.end())
|
||||
return iterator->second;
|
||||
|
||||
if (KernelRuntime::check_validity(dir_path))
|
||||
|
||||
+70
-61
@@ -2,6 +2,7 @@
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <fcntl.h>
|
||||
#include <filesystem>
|
||||
#include <fstream>
|
||||
#include <nvrtc.h>
|
||||
@@ -15,6 +16,7 @@
|
||||
#include "../utils/system.hpp"
|
||||
#include "cache.hpp"
|
||||
#include "device_runtime.hpp"
|
||||
#include "include_parser.hpp"
|
||||
|
||||
namespace deep_gemm {
|
||||
|
||||
@@ -23,29 +25,13 @@ 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;
|
||||
static std::filesystem::path cuobjdump_path;
|
||||
|
||||
static std::string get_library_version() {
|
||||
std::vector<char> buffer;
|
||||
for (const auto& f: collect_files(library_include_path / "deep_gemm")) {
|
||||
std::ifstream in(f, std::ios::binary);
|
||||
DG_HOST_ASSERT(in.is_open());
|
||||
|
||||
// Append into the buffer
|
||||
buffer.insert(buffer.end(),
|
||||
std::istreambuf_iterator<char>(in),
|
||||
std::istreambuf_iterator<char>());
|
||||
}
|
||||
return get_hex_digest(buffer);
|
||||
}
|
||||
|
||||
static void prepare_init(const std::string& library_root_path,
|
||||
const std::string& cuda_home_path_by_python) {
|
||||
Compiler::library_root_path = library_root_path;
|
||||
Compiler::library_include_path = Compiler::library_root_path / "include";
|
||||
Compiler::cuda_home = cuda_home_path_by_python;
|
||||
Compiler::library_version = get_library_version();
|
||||
Compiler::cuobjdump_path = Compiler::cuda_home / "bin" / "cuobjdump";
|
||||
}
|
||||
|
||||
@@ -57,12 +43,11 @@ public:
|
||||
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());
|
||||
DG_HOST_ASSERT(not cuobjdump_path.empty());
|
||||
|
||||
// Cache settings
|
||||
cache_dir_path = std::filesystem::path(get_env<std::string>("HOME")) / ".deep_gemm";
|
||||
if (const auto& env_cache_dir_path = get_env<std::string>("DG_JIT_CACHE_DIR"); not env_cache_dir_path.empty())
|
||||
if (const auto env_cache_dir_path = get_env<std::string>("DG_JIT_CACHE_DIR"); not env_cache_dir_path.empty())
|
||||
cache_dir_path = env_cache_dir_path;
|
||||
|
||||
// The compiler flags applied to all derived compilers
|
||||
@@ -82,58 +67,79 @@ public:
|
||||
return make_dirs(cache_dir_path / "tmp");
|
||||
}
|
||||
|
||||
std::filesystem::path get_tmp_file_path() const {
|
||||
return make_tmp_dir() / get_uuid();
|
||||
static void fsync_path(const std::filesystem::path& path) {
|
||||
const auto fd = ::open(path.c_str(), O_RDONLY);
|
||||
if (fd >= 0) {
|
||||
::fsync(fd);
|
||||
::close(fd);
|
||||
}
|
||||
}
|
||||
|
||||
void put(const std::filesystem::path& path, const std::string& data) const {
|
||||
const auto tmp_file_path = get_tmp_file_path();
|
||||
// Recursively fsync a directory: files and subdirectories first (bottom-up), then the directory itself
|
||||
// NOTES: ensures data and directory entries are visible on other nodes in distributed filesystems
|
||||
static void fsync_dir(const std::filesystem::path& dir_path) { // NOLINT(*-no-recursion)
|
||||
for (const auto& entry: std::filesystem::directory_iterator(dir_path)) {
|
||||
if (entry.is_directory())
|
||||
fsync_dir(entry.path());
|
||||
else if (entry.is_regular_file())
|
||||
fsync_path(entry.path());
|
||||
}
|
||||
fsync_path(dir_path);
|
||||
}
|
||||
|
||||
// Write into the temporary file
|
||||
std::ofstream out(tmp_file_path, std::ios::binary);
|
||||
static void put(const std::filesystem::path& path, const std::string& data) {
|
||||
std::ofstream out(path, std::ios::binary);
|
||||
DG_HOST_ASSERT(out.write(data.data(), data.size()));
|
||||
out.close();
|
||||
|
||||
// Atomically replace
|
||||
std::filesystem::rename(tmp_file_path, path);
|
||||
// NOTES: fsync to ensure the data is visible to other processes (e.g., NVCC)
|
||||
// on distributed filesystems, where `close()` alone does not guarantee persistence
|
||||
fsync_path(path);
|
||||
}
|
||||
|
||||
std::shared_ptr<KernelRuntime> build(const std::string& name, const std::string& code) const {
|
||||
const auto kernel_signature = fmt::format("{}$${}$${}$${}$${}", name, library_version, signature, flags, code);
|
||||
const auto kernel_signature = fmt::format("{}$${}$${}$${}", name, signature, flags, code);
|
||||
const auto dir_path = cache_dir_path / "cache" / fmt::format("kernel.{}.{}", name, get_hex_digest(kernel_signature));
|
||||
|
||||
// Hit the runtime cache
|
||||
if (const auto& runtime = kernel_runtime_cache->get(dir_path); runtime != nullptr)
|
||||
if (const auto runtime = kernel_runtime_cache->get(dir_path); runtime != nullptr)
|
||||
return runtime;
|
||||
|
||||
// Create the kernel directory
|
||||
make_dirs(dir_path);
|
||||
// Compile into a temporary directory, then atomically rename the whole directory
|
||||
// NOTES: renaming a directory is atomic on both local and distributed filesystems,
|
||||
// avoiding the stale inode issue that occurs when renaming individual files
|
||||
const auto tmp_dir_path = make_tmp_dir() / get_uuid();
|
||||
make_dirs(tmp_dir_path);
|
||||
|
||||
// Compile into a temporary CUBIN
|
||||
const auto tmp_cubin_path = get_tmp_file_path();
|
||||
// Compile into the temporary directory
|
||||
const auto tmp_cubin_path = tmp_dir_path / "kernel.cubin";
|
||||
if (get_env<int>("DG_JIT_DUMP_ASM") or get_env<int>("DG_JIT_DUMP_PTX")) {
|
||||
// Dump PTX if needed
|
||||
const auto tmp_ptx_path = get_tmp_file_path();
|
||||
compile(code, dir_path, tmp_cubin_path, tmp_ptx_path);
|
||||
|
||||
// Replace into the cache directory
|
||||
std::filesystem::rename(tmp_ptx_path, dir_path / "kernel.ptx");
|
||||
const auto tmp_ptx_path = tmp_dir_path / "kernel.ptx";
|
||||
compile(code, tmp_dir_path, tmp_cubin_path, tmp_ptx_path);
|
||||
} else {
|
||||
compile(code, dir_path, tmp_cubin_path);
|
||||
compile(code, tmp_dir_path, tmp_cubin_path);
|
||||
}
|
||||
|
||||
// Replace into the cache directory
|
||||
const auto cubin_path = dir_path / "kernel.cubin";
|
||||
std::filesystem::rename(tmp_cubin_path, cubin_path);
|
||||
|
||||
// Disassemble if needed
|
||||
if (get_env<int>("DG_JIT_DUMP_ASM") or get_env<int>("DG_JIT_DUMP_SASS")) {
|
||||
// Dump into a temporary SASS
|
||||
const auto tmp_sass_path = get_tmp_file_path();
|
||||
disassemble(cubin_path, tmp_sass_path);
|
||||
const auto tmp_sass_path = tmp_dir_path / "kernel.sass";
|
||||
disassemble(tmp_cubin_path, tmp_sass_path);
|
||||
}
|
||||
|
||||
// Replace into the current directory
|
||||
std::filesystem::rename(tmp_sass_path, dir_path / "kernel.sass");
|
||||
// Fsync before rename to ensure visibility on distributed filesystems
|
||||
fsync_dir(tmp_dir_path);
|
||||
|
||||
// Atomically rename the temporary directory to the final cache path
|
||||
// NOTES: if another rank already created dir_path, rename will fail — that's fine
|
||||
make_dirs(dir_path.parent_path());
|
||||
std::error_code error_code;
|
||||
std::filesystem::rename(tmp_dir_path, dir_path, error_code);
|
||||
if (error_code) {
|
||||
// Another rank beat us, then clean up our dir and use the existing one
|
||||
// NOTES: avoid `std::filesystem::remove_all` here — it can segfault on
|
||||
// distributed filesystems, when concurrent processes operate
|
||||
// on the same parent directory, causing stale directory entries
|
||||
safe_remove_all(tmp_dir_path);
|
||||
}
|
||||
|
||||
// Put into the runtime cache
|
||||
@@ -160,7 +166,6 @@ public:
|
||||
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);
|
||||
DG_DECLARE_STATIC_VAR_IN_CLASS(Compiler, cuobjdump_path);
|
||||
|
||||
class NVCCCompiler final: public Compiler {
|
||||
@@ -170,8 +175,8 @@ class NVCCCompiler final: public Compiler {
|
||||
DG_HOST_ASSERT(std::filesystem::exists(nvcc_path));
|
||||
|
||||
// Call the version command
|
||||
const auto& command = std::string(nvcc_path) + " --version";
|
||||
const auto& [return_code, output] = call_external_command(command);
|
||||
const auto command = std::string(nvcc_path) + " --version";
|
||||
const auto [return_code, output] = call_external_command(command);
|
||||
DG_HOST_ASSERT(return_code == 0);
|
||||
|
||||
// The version should be at least 12.3, for the best performance with 12.9
|
||||
@@ -189,14 +194,14 @@ public:
|
||||
NVCCCompiler() {
|
||||
// Override the compiler signature
|
||||
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())
|
||||
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();
|
||||
const auto [nvcc_major, nvcc_minor] = get_nvcc_version();
|
||||
signature = fmt::format("NVCC{}.{}", nvcc_major, nvcc_minor);
|
||||
|
||||
// The override the compiler flags
|
||||
// Only NVCC >= 12.9 supports arch-specific family suffix
|
||||
const auto& arch = device_runtime->get_arch(false, nvcc_major > 12 or nvcc_minor >= 9);
|
||||
const auto arch = device_runtime->get_arch(false, nvcc_major > 12 or nvcc_minor >= 9);
|
||||
flags = fmt::format("{} -I{} --gpu-architecture=sm_{} "
|
||||
"--compiler-options=-fPIC,-O3,-fconcepts,-Wno-deprecated-declarations,-Wno-abi "
|
||||
"-O3 --expt-relaxed-constexpr --expt-extended-lambda",
|
||||
@@ -207,14 +212,17 @@ public:
|
||||
const std::filesystem::path &cubin_path,
|
||||
const std::optional<std::filesystem::path> &ptx_path) const override {
|
||||
// Write the code into the cache directory
|
||||
const auto& code_path = dir_path / "kernel.cu";
|
||||
const auto code_path = dir_path / "kernel.cu";
|
||||
put(code_path, code);
|
||||
|
||||
// Compile
|
||||
const auto& command = fmt::format("{} {} -cubin -o {} {}", nvcc_path.c_str(), code_path.c_str(), cubin_path.c_str(), flags);
|
||||
// Avoid cwd files shadowing C++ standard library headers
|
||||
const auto compile_dir = make_tmp_dir();
|
||||
const auto command = fmt::format("cd {} && {} {} -cubin -o {} {}",
|
||||
compile_dir.c_str(), 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\n", command.c_str());
|
||||
const auto& [return_code, output] = call_external_command(command);
|
||||
const auto [return_code, output] = call_external_command(command);
|
||||
if (return_code != 0) {
|
||||
printf("NVCC compilation failed: %s\n", output.c_str());
|
||||
DG_HOST_ASSERT(false and "NVCC compilation failed");
|
||||
@@ -222,7 +230,8 @@ public:
|
||||
|
||||
// Compile to PTX if needed
|
||||
if (ptx_path.has_value()) {
|
||||
const auto ptx_command = fmt::format("{} {} -ptx -o {} {}", nvcc_path.c_str(), code_path.c_str(), ptx_path->c_str(), flags);
|
||||
const auto ptx_command = fmt::format("cd {} && {} {} -ptx -o {} {}",
|
||||
compile_dir.c_str(), nvcc_path.c_str(), code_path.c_str(), ptx_path->c_str(), flags);
|
||||
if (get_env("DG_JIT_DEBUG", 0) or get_env("DG_JIT_PRINT_COMPILER_COMMAND", 0))
|
||||
printf("Running NVCC PTX command: %s\n", ptx_command.c_str());
|
||||
const auto [ptx_return_code, ptx_output] = call_external_command(ptx_command);
|
||||
@@ -267,7 +276,7 @@ public:
|
||||
|
||||
// Override the compiler flags
|
||||
// Only NVRTC >= 12.9 supports arch-specific family suffix
|
||||
const auto& arch = device_runtime->get_arch(false, major > 12 or minor >= 9);
|
||||
const auto arch = device_runtime->get_arch(false, major > 12 or minor >= 9);
|
||||
flags = fmt::format("{} {}--gpu-architecture=sm_{} -default-device {} --device-int128",
|
||||
flags, include_dirs, arch, pch_flags);
|
||||
}
|
||||
@@ -276,7 +285,7 @@ public:
|
||||
const std::filesystem::path &cubin_path,
|
||||
const std::optional<std::filesystem::path> &ptx_path) const override {
|
||||
// Write the code into the cache directory
|
||||
const auto& code_path = dir_path / "kernel.cu";
|
||||
const auto code_path = dir_path / "kernel.cu";
|
||||
put(code_path, code);
|
||||
|
||||
// Parse compilation options
|
||||
@@ -302,7 +311,7 @@ public:
|
||||
// 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());
|
||||
const auto compile_result = nvrtcCompileProgram(program, static_cast<int>(option_cstrs.size()), option_cstrs.data());
|
||||
|
||||
// Get and print compiler log
|
||||
size_t log_size;
|
||||
|
||||
@@ -7,10 +7,13 @@
|
||||
#include "../utils/exception.hpp"
|
||||
#include "../utils/lazy_init.hpp"
|
||||
|
||||
#define PYTORCH_SUPPORTS_GET_CUBLASLT_HANDLE (TORCH_VERSION_MAJOR > 2 or (TORCH_VERSION_MAJOR == 2 and TORCH_VERSION_MINOR >= 3))
|
||||
|
||||
namespace deep_gemm {
|
||||
|
||||
class DeviceRuntime {
|
||||
int num_sms = 0, tc_util = 0;
|
||||
bool enable_pdl = false;
|
||||
std::shared_ptr<cudaDeviceProp> cached_prop;
|
||||
|
||||
// cuBLASLt utils
|
||||
@@ -18,24 +21,52 @@ class DeviceRuntime {
|
||||
|
||||
public:
|
||||
// Create the cuBLASLt handle ourselves
|
||||
cublasLtHandle_t cublaslt_handle{};
|
||||
std::shared_ptr<torch::Tensor> cublaslt_workspace;
|
||||
cublasLtHandle_t cublaslt_handle;
|
||||
torch::Tensor cublaslt_workspace;
|
||||
bool use_pytorch_managed_cublaslt_handle;
|
||||
bool use_temp_cublaslt_workspace;
|
||||
|
||||
explicit DeviceRuntime() {
|
||||
cublaslt_workspace = std::make_shared<torch::Tensor>(torch::empty({kCublasLtWorkspaceSize}, dtype(torch::kByte).device(at::kCUDA)));
|
||||
DG_CUBLASLT_CHECK(cublasLtCreate(&cublaslt_handle));
|
||||
|
||||
// Whether to use PyTorch cuBLASLt
|
||||
// By default, we don't use it,
|
||||
// as `at::cuda::getCurrentCUDABlasLtHandle` has large CPU overhead with some PyTorch versions
|
||||
use_pytorch_managed_cublaslt_handle = get_env<int>("DG_USE_PYTORCH_CUBLASLT_HANDLE", 0) > 0;
|
||||
#if not PYTORCH_SUPPORTS_GET_CUBLASLT_HANDLE
|
||||
DG_HOST_ASSERT(not use_pytorch_managed_cublaslt_handle and "PyTorch does not support to get cuBLASLt handle");
|
||||
#endif
|
||||
|
||||
// Whether to create workspace tensor on each call instead of holding one.
|
||||
// Enabled by compute-sanitizer tests, which trigger `cudaErrorCudartUnloading`
|
||||
// when the workspace tensor is destructed after CUDA driver shutdown.
|
||||
use_temp_cublaslt_workspace = get_env<int>("DG_USE_TEMP_CUBLASLT_WORKSPACE", 0) > 0;
|
||||
|
||||
if (not use_pytorch_managed_cublaslt_handle)
|
||||
DG_CUBLASLT_CHECK(cublasLtCreate(&cublaslt_handle));
|
||||
|
||||
if (not use_temp_cublaslt_workspace)
|
||||
cublaslt_workspace = torch::empty({kCublasLtWorkspaceSize}, dtype(torch::kByte).device(at::kCUDA));
|
||||
}
|
||||
|
||||
~DeviceRuntime() noexcept(false) {
|
||||
DG_CUBLASLT_CHECK(cublasLtDestroy(cublaslt_handle));
|
||||
if (not use_pytorch_managed_cublaslt_handle)
|
||||
DG_CUBLASLT_CHECK(cublasLtDestroy(cublaslt_handle));
|
||||
}
|
||||
|
||||
cublasLtHandle_t get_cublaslt_handle() const {
|
||||
#if PYTORCH_SUPPORTS_GET_CUBLASLT_HANDLE
|
||||
if (use_pytorch_managed_cublaslt_handle)
|
||||
return at::cuda::getCurrentCUDABlasLtHandle();
|
||||
#endif
|
||||
|
||||
// Self-managed handle
|
||||
return cublaslt_handle;
|
||||
}
|
||||
|
||||
torch::Tensor get_cublaslt_workspace() const {
|
||||
return *cublaslt_workspace;
|
||||
if (use_temp_cublaslt_workspace)
|
||||
return torch::empty({kCublasLtWorkspaceSize}, dtype(torch::kByte).device(at::kCUDA));
|
||||
return cublaslt_workspace;
|
||||
}
|
||||
|
||||
std::shared_ptr<cudaDeviceProp> get_prop() {
|
||||
@@ -56,7 +87,7 @@ public:
|
||||
|
||||
std::string get_arch(const bool& number_only = false,
|
||||
const bool& support_arch_family = false) {
|
||||
const auto& [major, minor] = get_arch_pair();
|
||||
const auto [major, minor] = get_arch_pair();
|
||||
if (major == 10 and minor != 1) {
|
||||
if (number_only)
|
||||
return "100";
|
||||
@@ -92,6 +123,14 @@ public:
|
||||
int get_tc_util() const {
|
||||
return tc_util == 0 ? 100 : tc_util;
|
||||
}
|
||||
|
||||
void set_pdl(const bool& new_enable_pdl) {
|
||||
enable_pdl = new_enable_pdl;
|
||||
}
|
||||
|
||||
bool get_pdl() const {
|
||||
return enable_pdl;
|
||||
}
|
||||
};
|
||||
|
||||
static auto device_runtime = LazyInit<DeviceRuntime>([](){ return std::make_shared<DeviceRuntime>(); });
|
||||
|
||||
+73
-19
@@ -24,7 +24,7 @@ static void* get_driver_handle() {
|
||||
#define DECL_LAZY_CUDA_DRIVER_FUNCTION(name) \
|
||||
template <typename... Args> \
|
||||
static auto lazy_##name(Args&&... args) -> decltype(name(args...)) { \
|
||||
using FuncType = decltype(&name); \
|
||||
using FuncType = decltype(&(name)); \
|
||||
static FuncType func = nullptr; \
|
||||
if (func == nullptr) { \
|
||||
func = reinterpret_cast<FuncType>(dlsym(get_driver_handle(), #name)); \
|
||||
@@ -39,6 +39,9 @@ DECL_LAZY_CUDA_DRIVER_FUNCTION(cuFuncSetAttribute);
|
||||
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuModuleLoad);
|
||||
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuModuleUnload);
|
||||
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuModuleGetFunction);
|
||||
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuLibraryLoadFromFile);
|
||||
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuLibraryUnload);
|
||||
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuKernelGetFunction);
|
||||
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuLaunchKernelEx);
|
||||
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuTensorMapEncodeTiled);
|
||||
|
||||
@@ -65,13 +68,13 @@ static KernelHandle load_kernel(const std::filesystem::path& cubin_path, const s
|
||||
}
|
||||
|
||||
static void unload_library(const LibraryHandle& library) {
|
||||
const auto& error = cudaLibraryUnload(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) {
|
||||
const dim3& grid_dim, const dim3& block_dim, const int& cluster_dim, const bool& enable_pdl) {
|
||||
if (smem_size > 0)
|
||||
DG_CUDA_RUNTIME_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
|
||||
|
||||
@@ -80,17 +83,27 @@ static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel,
|
||||
config.blockDim = block_dim;
|
||||
config.dynamicSmemBytes = smem_size;
|
||||
config.stream = stream;
|
||||
config.numAttrs = 0;
|
||||
config.attrs = nullptr;
|
||||
|
||||
// Create attributes
|
||||
// NOTES: must use `static` or the `attr` will be deconstructed
|
||||
static LaunchAttrHandle attr;
|
||||
static LaunchAttrHandle attrs[2];
|
||||
config.numAttrs = 0;
|
||||
config.attrs = attrs;
|
||||
|
||||
// Cluster size
|
||||
if (cluster_dim > 1) {
|
||||
auto& attr = attrs[config.numAttrs ++];
|
||||
attr.id = cudaLaunchAttributeClusterDimension;
|
||||
attr.val.clusterDim = {static_cast<unsigned>(cluster_dim), 1, 1};
|
||||
config.attrs = &attr;
|
||||
config.numAttrs = 1;
|
||||
}
|
||||
|
||||
// Dependent kernel launch
|
||||
if (enable_pdl) {
|
||||
auto& attr = attrs[config.numAttrs ++];
|
||||
attr.id = cudaLaunchAttributeProgrammaticStreamSerialization;
|
||||
attr.val.programmaticStreamSerializationAllowed = 1;
|
||||
}
|
||||
|
||||
return config;
|
||||
}
|
||||
|
||||
@@ -103,19 +116,46 @@ static auto launch_kernel(const KernelHandle& kernel, const LaunchConfigHandle&
|
||||
#else
|
||||
|
||||
// Use CUDA driver API
|
||||
using LibraryHandle = CUmodule;
|
||||
using KernelHandle = CUfunction;
|
||||
using LaunchConfigHandle = CUlaunchConfig;
|
||||
using LaunchAttrHandle = CUlaunchAttribute;
|
||||
|
||||
// `cuLibraryEnumerateKernels` is supported since CUDA Driver API 12.4
|
||||
#if CUDA_VERSION >= 12040
|
||||
#define DG_JIT_USE_LIBRARY_ENUM_KERNELS
|
||||
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuLibraryGetKernelCount);
|
||||
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuLibraryEnumerateKernels);
|
||||
using LibraryHandle = CUlibrary;
|
||||
#else
|
||||
using LibraryHandle = CUmodule;
|
||||
#endif
|
||||
|
||||
#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_opt = nullptr) {
|
||||
LibraryHandle library;
|
||||
KernelHandle kernel;
|
||||
|
||||
#ifdef DG_JIT_USE_LIBRARY_ENUM_KERNELS
|
||||
DG_CUDA_DRIVER_CHECK(lazy_cuLibraryLoadFromFile(&library, cubin_path.c_str(), nullptr, nullptr, 0, nullptr, nullptr, 0));
|
||||
unsigned int num_kernels;
|
||||
DG_CUDA_DRIVER_CHECK(lazy_cuLibraryGetKernelCount(&num_kernels, library));
|
||||
if (num_kernels != 1) {
|
||||
const auto dir_path = cubin_path.parent_path();
|
||||
printf("Corrupted JIT cache directory (expected 1 kernel, found %u): %s, "
|
||||
"please run `rm -rf %s` and restart your task.\n",
|
||||
num_kernels, dir_path.c_str(), dir_path.c_str());
|
||||
DG_HOST_ASSERT(false and "Corrupted JIT cache directory");
|
||||
}
|
||||
|
||||
CUkernel cu_kernel;
|
||||
DG_CUDA_DRIVER_CHECK(lazy_cuLibraryEnumerateKernels(&cu_kernel, 1, library));
|
||||
DG_CUDA_DRIVER_CHECK(lazy_cuKernelGetFunction(&kernel, cu_kernel));
|
||||
#else
|
||||
DG_CUDA_DRIVER_CHECK(lazy_cuModuleLoad(&library, cubin_path.c_str()));
|
||||
DG_CUDA_DRIVER_CHECK(lazy_cuModuleGetFunction(&kernel, library, func_name.c_str()));
|
||||
#endif
|
||||
|
||||
if (library_opt != nullptr)
|
||||
*library_opt = library;
|
||||
@@ -123,13 +163,17 @@ static KernelHandle load_kernel(const std::filesystem::path& cubin_path, const s
|
||||
}
|
||||
|
||||
static void unload_library(const LibraryHandle& library) {
|
||||
const auto& error = lazy_cuModuleUnload(library);
|
||||
#ifdef DG_JIT_USE_LIBRARY_ENUM_KERNELS
|
||||
const auto error = lazy_cuLibraryUnload(library);
|
||||
#else
|
||||
const auto error = lazy_cuModuleUnload(library);
|
||||
#endif
|
||||
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) {
|
||||
const dim3& grid_dim, const dim3& block_dim, const int& cluster_dim, const bool& enable_pdl) {
|
||||
if (smem_size > 0)
|
||||
DG_CUDA_DRIVER_CHECK(lazy_cuFuncSetAttribute(kernel, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem_size));
|
||||
|
||||
@@ -142,19 +186,29 @@ static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel,
|
||||
config.blockDimZ = block_dim.z;
|
||||
config.sharedMemBytes = smem_size;
|
||||
config.hStream = stream;
|
||||
config.numAttrs = 0;
|
||||
config.attrs = nullptr;
|
||||
|
||||
|
||||
// Create attributes
|
||||
// NOTES: must use `static` or the `attr` will be deconstructed
|
||||
static LaunchAttrHandle attr;
|
||||
static LaunchAttrHandle attrs[2];
|
||||
config.numAttrs = 0;
|
||||
config.attrs = attrs;
|
||||
|
||||
// Cluster size
|
||||
if (cluster_dim > 1) {
|
||||
auto& attr = attrs[config.numAttrs ++];
|
||||
attr.id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION;
|
||||
attr.value.clusterDim.x = cluster_dim;
|
||||
attr.value.clusterDim.x = static_cast<unsigned>(cluster_dim);
|
||||
attr.value.clusterDim.y = 1;
|
||||
attr.value.clusterDim.z = 1;
|
||||
config.attrs = &attr;
|
||||
config.numAttrs = 1;
|
||||
}
|
||||
|
||||
// Dependent kernel launch
|
||||
if (enable_pdl) {
|
||||
auto& attr = attrs[config.numAttrs ++];
|
||||
attr.id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION;
|
||||
attr.value.programmaticStreamSerializationAllowed = 1;
|
||||
}
|
||||
|
||||
return config;
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
#pragma once
|
||||
|
||||
#include <filesystem>
|
||||
#include <regex>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "../utils/format.hpp"
|
||||
#include "../utils/system.hpp"
|
||||
|
||||
namespace deep_gemm {
|
||||
|
||||
class IncludeParser {
|
||||
std::unordered_map<std::string, std::optional<std::string>> cache;
|
||||
|
||||
static std::vector<std::string> get_includes(const std::string& code, const std::filesystem::path& file_path = "") {
|
||||
std::vector<std::string> includes;
|
||||
const std::regex pattern(R"(#\s*include\s*[<"][^>"]+[>"])");
|
||||
std::sregex_iterator iter(code.begin(), code.end(), pattern);
|
||||
const std::sregex_iterator end;
|
||||
|
||||
// TODO: parse relative paths as well
|
||||
for (; iter != end; ++ iter) {
|
||||
const auto include_str = iter->str();
|
||||
const int len = include_str.length();
|
||||
if (include_str.substr(0, 10) == "#include <" and include_str[len - 1] == '>' and include_str[10] != ' ' and include_str[len - 2] != ' ') {
|
||||
std::string filename = include_str.substr(10, len - 11);
|
||||
if (filename.substr(0, 9) == "deep_gemm") // We only parse `<deep_gemm/*>`
|
||||
includes.push_back(filename);
|
||||
} else {
|
||||
std::string error_info = fmt::format("Non-standard include: {}", include_str);
|
||||
if (file_path != "")
|
||||
error_info += fmt::format(" ({})", file_path.string());
|
||||
DG_HOST_UNREACHABLE(error_info);
|
||||
}
|
||||
}
|
||||
return includes;
|
||||
}
|
||||
|
||||
public:
|
||||
static std::filesystem::path library_include_path;
|
||||
|
||||
static void prepare_init(const std::string& library_root_path) {
|
||||
library_include_path = std::filesystem::path(library_root_path) / "include";
|
||||
}
|
||||
|
||||
std::string get_hash_value(const std::string& code, const bool& exclude_code = true) {
|
||||
std::stringstream ss;
|
||||
for (const auto& i: get_includes(code))
|
||||
ss << get_hash_value_by_path(library_include_path / i) << "$";
|
||||
if (not exclude_code)
|
||||
ss << "#" << get_hex_digest(code);
|
||||
return get_hex_digest(ss.str());
|
||||
}
|
||||
|
||||
std::string get_hash_value_by_path(const std::filesystem::path& path) {
|
||||
// Check whether hit in cache
|
||||
// ReSharper disable once CppUseAssociativeContains
|
||||
if (cache.count(path) > 0) {
|
||||
const auto opt = cache[path];
|
||||
if (not opt.has_value())
|
||||
DG_HOST_UNREACHABLE(fmt::format("Circular include may occur: {}", path.string()));
|
||||
return opt.value();
|
||||
}
|
||||
|
||||
// Read file and calculate hash recursively
|
||||
std::ifstream in(path);
|
||||
if (not in.is_open())
|
||||
DG_HOST_UNREACHABLE(fmt::format("Failed to open: {}", path.string()));
|
||||
std::string code((std::istreambuf_iterator<char>(in)), std::istreambuf_iterator<char>());
|
||||
cache[path] = std::nullopt;
|
||||
return (cache[path] = get_hash_value(code, false)).value();
|
||||
}
|
||||
};
|
||||
|
||||
DG_DECLARE_STATIC_VAR_IN_CLASS(IncludeParser, library_include_path);
|
||||
|
||||
static auto include_parser = std::make_shared<IncludeParser>();
|
||||
|
||||
} // namespace deep_gemm
|
||||
+76
-29
@@ -1,10 +1,13 @@
|
||||
#pragma once
|
||||
|
||||
#include <chrono>
|
||||
|
||||
#include "../utils/exception.hpp"
|
||||
#include "../utils/format.hpp"
|
||||
#include "../utils/system.hpp"
|
||||
#include "device_runtime.hpp"
|
||||
#include "handle.hpp"
|
||||
#include "include_parser.hpp"
|
||||
|
||||
namespace deep_gemm {
|
||||
|
||||
@@ -13,12 +16,13 @@ struct LaunchArgs {
|
||||
int num_threads;
|
||||
int smem_size;
|
||||
int cluster_dim;
|
||||
bool enable_pdl;
|
||||
|
||||
LaunchArgs(const int& grid_dim_x, const int& num_threads, const int& smem_size = 0, const int& cluster_dim = 1):
|
||||
grid_dim({grid_dim_x, 1}), num_threads(num_threads), smem_size(smem_size), cluster_dim(cluster_dim) {}
|
||||
LaunchArgs(const int& grid_dim_x, const int& num_threads, const int& smem_size = 0, const int& cluster_dim = 1, const bool& enable_pdl = true):
|
||||
grid_dim({grid_dim_x, 1}), num_threads(num_threads), smem_size(smem_size), cluster_dim(cluster_dim), enable_pdl(enable_pdl) {}
|
||||
|
||||
LaunchArgs(const std::pair<int, int>& grid_dim, const int& num_threads, const int& smem_size = 0, const int& cluster_dim = 1):
|
||||
grid_dim(grid_dim), num_threads(num_threads), smem_size(smem_size), cluster_dim(cluster_dim) {}
|
||||
LaunchArgs(const std::pair<int, int>& grid_dim, const int& num_threads, const int& smem_size = 0, const int& cluster_dim = 1, const bool& enable_pdl = true):
|
||||
grid_dim(grid_dim), num_threads(num_threads), smem_size(smem_size), cluster_dim(cluster_dim), enable_pdl(enable_pdl) {}
|
||||
};
|
||||
|
||||
class KernelRuntime final {
|
||||
@@ -33,36 +37,56 @@ public:
|
||||
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";
|
||||
const auto cuobjdump_path = cuda_home / "bin" / "cuobjdump";
|
||||
const auto cubin_path = dir_path / "kernel.cubin";
|
||||
if (get_env<int>("DG_JIT_DEBUG"))
|
||||
printf("Loading CUBIN: %s\n", cubin_path.c_str());
|
||||
|
||||
// Record start time
|
||||
std::chrono::high_resolution_clock::time_point start_time;
|
||||
if (get_env<int>("DG_JIT_DEBUG") or get_env<int>("DG_JIT_PRINT_LOAD_TIME"))
|
||||
start_time = std::chrono::high_resolution_clock::now();
|
||||
|
||||
#ifdef DG_JIT_USE_LIBRARY_ENUM_KERNELS
|
||||
// Load from the library
|
||||
kernel = load_kernel(cubin_path, {}, &library);
|
||||
#else
|
||||
// Find the only symbol
|
||||
// TODO: use kernel enumeration for newer drivers
|
||||
const std::vector<std::string> illegal_names = {"vprintf", "__instantiate_kernel", "__internal", "__assertfail"};
|
||||
const auto& [exit_code, symbols] = call_external_command(fmt::format("{} -symbols {}", cuobjdump_path.c_str(), cubin_path.c_str()));
|
||||
const auto [exit_code, symbols] = call_external_command(fmt::format("{} -symbols {}", cuobjdump_path.c_str(), cubin_path.c_str()));
|
||||
DG_HOST_ASSERT(exit_code == 0);
|
||||
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 line.find("STO_ENTRY") != std::string::npos 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(' ');
|
||||
[&](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));
|
||||
}
|
||||
}
|
||||
if (get_env<int>("DG_JIT_DEBUG")) {
|
||||
printf("Symbol names: ");
|
||||
|
||||
// Print symbols
|
||||
if (symbol_names.size() != 1 or get_env<int>("DG_JIT_DEBUG")) {
|
||||
printf("Symbols: ");
|
||||
printf(" > CUBIN: %s\n", cubin_path.c_str());
|
||||
printf(" > Raw symbols: %s\n", symbols.c_str());
|
||||
printf(" > Parsed symbols:\n");
|
||||
for (const auto& symbol: symbol_names)
|
||||
printf("%s, ", symbol.c_str());
|
||||
printf("\n");
|
||||
printf(" > %s, ", symbol.c_str());
|
||||
}
|
||||
DG_HOST_ASSERT(symbol_names.size() == 1);
|
||||
|
||||
// Load from the library
|
||||
DG_HOST_ASSERT(symbol_names.size() == 1);
|
||||
kernel = load_kernel(cubin_path, symbol_names[0], &library);
|
||||
#endif
|
||||
|
||||
// Print load time
|
||||
if (get_env<int>("DG_JIT_DEBUG") or get_env<int>("DG_JIT_PRINT_LOAD_TIME")) {
|
||||
std::chrono::duration<double, std::milli> load_time = std::chrono::high_resolution_clock::now() - start_time;
|
||||
printf("Load time (%s): %.2lf ms\n", dir_path.c_str(), load_time.count());
|
||||
}
|
||||
}
|
||||
|
||||
static void prepare_init(const std::string& cuda_home_path_by_python) {
|
||||
@@ -70,8 +94,19 @@ public:
|
||||
}
|
||||
|
||||
static bool check_validity(const std::filesystem::path& dir_path) {
|
||||
return std::filesystem::exists(dir_path / "kernel.cu") and
|
||||
std::filesystem::exists(dir_path / "kernel.cubin");
|
||||
if (not std::filesystem::exists(dir_path))
|
||||
return false;
|
||||
|
||||
// NOTES: if the directory exists, `kernel.cu` and `kernel.cubin` must both exist,
|
||||
// because the directory is created atomically via rename
|
||||
if (not std::filesystem::exists(dir_path / "kernel.cu") or
|
||||
not std::filesystem::exists(dir_path / "kernel.cubin")) {
|
||||
printf("Corrupted JIT cache directory (missing kernel.cu or kernel.cubin): %s, "
|
||||
"please run `rm -rf %s` and restart your task.\n",
|
||||
dir_path.c_str(), dir_path.c_str());
|
||||
DG_HOST_ASSERT(false and "Corrupted JIT cache directory");
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
~KernelRuntime() noexcept(false) {
|
||||
@@ -86,30 +121,42 @@ class LaunchRuntime {
|
||||
public:
|
||||
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))
|
||||
printf("Generated kernel code: %s\n", code.c_str());
|
||||
auto code = Derived::generate_impl(args);
|
||||
|
||||
// NOTES: we require that `generate_impl`'s includes never change
|
||||
static std::string include_hash;
|
||||
if (include_hash.empty())
|
||||
include_hash = include_parser->get_hash_value(code);
|
||||
|
||||
// TODO: optimize string concat performance
|
||||
code = fmt::format("// Includes' hash value: {}\n{}", include_hash, code);
|
||||
if (get_env<int>("DG_JIT_DEBUG"))
|
||||
printf("Generated kernel code:\n%s\n", code.c_str());
|
||||
return code;
|
||||
}
|
||||
|
||||
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;
|
||||
const auto kernel = kernel_runtime->kernel;
|
||||
const auto stream = at::cuda::getCurrentCUDAStream();
|
||||
LaunchArgs launch_args = args.launch_args;
|
||||
|
||||
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};
|
||||
// Allow runtime override from Python.
|
||||
// NOTES: the default is enabled.
|
||||
launch_args.enable_pdl = device_runtime->get_pdl();
|
||||
|
||||
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);
|
||||
grid_dim, block_dim, launch_args.cluster_dim, launch_args.enable_pdl);
|
||||
|
||||
// Launch in the derived class
|
||||
if (get_env<int>("DG_JIT_DEBUG")) {
|
||||
printf("Launch kernel with {%d, %d} x %d, shared memory: %d bytes, cluster: %d, stream: %ld\n",
|
||||
printf("Launch kernel with {%d, %d} x %d, shared memory: %d bytes, cluster: %d, pdl: %d, stream: %ld\n",
|
||||
launch_args.grid_dim.first, launch_args.grid_dim.second, launch_args.num_threads,
|
||||
launch_args.smem_size, launch_args.cluster_dim, stream.id());
|
||||
launch_args.smem_size, launch_args.cluster_dim, launch_args.enable_pdl, stream.id());
|
||||
}
|
||||
Derived::launch_impl(kernel, config, args);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user