From 3e765cd3aa064a35f368615b564e6efc86894339 Mon Sep 17 00:00:00 2001
From: leavelet <leavelet@163.com>
Date: Wed, 19 Nov 2025 14:50:02 +0000
Subject: [PATCH] fix: try add wrapper for preload

---
 csrc/python_api.cpp | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/csrc/python_api.cpp b/csrc/python_api.cpp
index d0848fd..4ed41d3 100644
--- a/csrc/python_api.cpp
+++ b/csrc/python_api.cpp
@@ -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);
 }
 
+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)
 int64_t get_tma_aligned_size_wrapper(int64_t x, int64_t element_size);
 int64_t get_mk_alignment_for_contiguous_layout_wrapper();
@@ -249,8 +254,7 @@ TORCH_LIBRARY(deep_gemm, m) {
 
     m.def("preload_kernels() -> ()");
     m.impl("preload_kernels", []() {
-        const auto& cache_dir = deep_gemm::Compiler::get_cache_dir();
-        deep_gemm::kernel_runtime_cache->preload_all(cache_dir);
+        deep_gemm_wrappers::preload_kernels_wrapper();
     });
 
     // layout APIs