leavelet/deepgemm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: support misc kernel launch

Browse Source
This commit is contained in:
Rain Jiang
2025-08-29 16:23:39 -07:00
committed by FlamingoPg
parent 301cbc1d75
commit f4adba8a66
9 changed files with 52 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
csrc/jit/device_runtime.hpp
View File

@@ -11,11 +11,7 @@ namespace deep_gemm {
class DeviceRuntime {
int num_sms = 0, tc_util = 0;
std::shared_ptr<cudaDeviceProp> cached_prop;
// cuBLASLt utils
static constexpr size_t kCublasLtWorkspaceSize = 32 * 1024 * 1024;
cublasLtHandle_t cublaslt_handle{};
std::shared_ptr<torch::Tensor> cublaslt_workspace;
int compile_mode = 0;
// cuBLASLt utils
static constexpr size_t kCublasLtWorkspaceSize = 32 * 1024 * 1024;
@@ -77,6 +73,15 @@ public:
return num_sms;
}
void set_compile_mode(const int& new_compile_mode) {
DG_HOST_ASSERT(0 <= new_compile_mode and new_compile_mode <= 1);
compile_mode = new_compile_mode;
}
int get_compile_mode() {
return compile_mode;
}
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;