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

20
csrc/apis/runtime.hpp
View File

@@ -4,6 +4,7 @@
#include "../jit/compiler.hpp"
#endif
#include "../jit/device_runtime.hpp"
#include "../jit_kernels/heuristics/runtime.hpp"
namespace deep_gemm::runtime {
@@ -20,10 +21,29 @@ static void register_apis(pybind11::module_& m) {
m.def("get_tc_util", [&]() {
return device_runtime->get_tc_util();
});
m.def("set_pdl", [&](const bool& new_enable_pdl) {
device_runtime->set_pdl(new_enable_pdl);
});
m.def("get_pdl", [&]() {
return device_runtime->get_pdl();
});
m.def("set_ignore_compile_dims", [&](const bool& new_value) {
heuristics_runtime->set_ignore_compile_dims(new_value);
});
m.def("set_block_size_multiple_of", [&](const std::variant<int, std::tuple<int, int>>& new_value) {
if (std::holds_alternative<int>(new_value)) {
auto x = std::get<int>(new_value);
heuristics_runtime->set_block_size_multiple_of(x, x);
} else {
auto [x, y] = std::get<std::tuple<int, int>>(new_value);
heuristics_runtime->set_block_size_multiple_of(x, y);
}
});
m.def("init", [&](const std::string& library_root_path, const std::string& cuda_home_path_by_python) {
#if DG_TENSORMAP_COMPATIBLE
Compiler::prepare_init(library_root_path, cuda_home_path_by_python);
KernelRuntime::prepare_init(cuda_home_path_by_python);
IncludeParser::prepare_init(library_root_path);
#endif
});
}