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

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

Browse Source 
* Merge with private repo

* Update README

* Update README

* Update README

* Add PyTorch requirements

* Fix sync scopes for MQA logits (#256)

* Update README
This commit is contained in:
Chenggang Zhao
2026-04-17 09:45:14 +08:00
committed by GitHub
parent d30fc36c8f
commit 7f2a703ed5
109 changed files with 12101 additions and 3219 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/jit/device_runtime.hpp
+46 -7
View File
@@ -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>(); });