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

10
csrc/python_api.cpp
View File

@@ -214,6 +214,16 @@ TORCH_LIBRARY(deep_gemm, m) {
return static_cast<int64_t>(deep_gemm::device_runtime->get_num_sms());
});
m.def("set_compile_mode(int new_compile_mode) -> ()");
m.impl("set_compile_mode", [](int64_t new_compile_mode) {
deep_gemm::device_runtime->set_compile_mode(static_cast<int>(new_compile_mode));
});
m.def("get_compile_mode() -> int");
m.impl("get_compile_mode", []() -> int64_t {
return static_cast<int64_t>(deep_gemm::device_runtime->get_compile_mode());
});
m.def("set_tc_util(int new_tc_util) -> ()");
m.impl("set_tc_util", [](int64_t new_tc_util) {
deep_gemm::device_runtime->set_tc_util(static_cast<int>(new_tc_util));