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

fix: try add wrapper for preload

Browse Source
This commit is contained in:
孙远航
2025-11-19 14:50:02 +00:00
parent d4ffb6a9bf
commit 3e765cd3aa
1 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
csrc/python_api.cpp
View File

@@ -94,6 +94,11 @@ void init_wrapper(const std::string& library_root_path, const std::string& cuda_
deep_gemm::KernelRuntime::prepare_init(cuda_home_path_by_python); deep_gemm::KernelRuntime::prepare_init(cuda_home_path_by_python);
} }
void preload_kernels_wrapper() {
const auto& cache_dir = deep_gemm::Compiler::get_cache_dir();
deep_gemm::kernel_runtime_cache->preload_all(cache_dir);
}
// Scalar layout utility wrappers (int64_t signatures for PyTorch registration) // Scalar layout utility wrappers (int64_t signatures for PyTorch registration)
int64_t get_tma_aligned_size_wrapper(int64_t x, int64_t element_size); int64_t get_tma_aligned_size_wrapper(int64_t x, int64_t element_size);
int64_t get_mk_alignment_for_contiguous_layout_wrapper(); int64_t get_mk_alignment_for_contiguous_layout_wrapper();
@@ -249,8 +254,7 @@ TORCH_LIBRARY(deep_gemm, m) {
m.def("preload_kernels() -> ()"); m.def("preload_kernels() -> ()");
m.impl("preload_kernels", []() { m.impl("preload_kernels", []() {
const auto& cache_dir = deep_gemm::Compiler::get_cache_dir(); deep_gemm_wrappers::preload_kernels_wrapper();
deep_gemm::kernel_runtime_cache->preload_all(cache_dir);
}); });
// layout APIs // layout APIs