From d4ffb6a9bfa9864e0a48f45a5c0c4a642704e13e Mon Sep 17 00:00:00 2001 From: leavelet Date: Wed, 19 Nov 2025 14:37:49 +0000 Subject: [PATCH] feat: add support for deepgemm sglang --- csrc/python_api.cpp | 8 +++++++- deep_gemm/__init__.py | 7 ++++--- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/csrc/python_api.cpp b/csrc/python_api.cpp index 14f3b15..d0848fd 100644 --- a/csrc/python_api.cpp +++ b/csrc/python_api.cpp @@ -247,6 +247,12 @@ TORCH_LIBRARY(deep_gemm, m) { deep_gemm_wrappers::init_wrapper(library_root_path, cuda_home_path_by_torch); }); + 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); + }); + // layout APIs m.def("transform_sf_into_required_layout(Tensor sf, int mn, int k, int[] recipe, int? num_groups=None, bool is_sfa=False, bool disable_ue8m0_cast=False) -> Tensor", deep_gemm_wrappers::transform_sf_into_required_layout_wrapper); @@ -540,4 +546,4 @@ int64_t deep_gemm_wrappers::get_mk_alignment_for_contiguous_layout_wrapper() { return static_cast(deep_gemm::get_mk_alignment_for_contiguous_layout()); } -REGISTER_EXTENSION(deep_gemm_cpp) \ No newline at end of file +REGISTER_EXTENSION(deep_gemm_cpp) diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py index f088511..b1bab8a 100644 --- a/deep_gemm/__init__.py +++ b/deep_gemm/__init__.py @@ -56,6 +56,7 @@ set_compile_mode = _wrap_op('set_compile_mode') get_compile_mode = _wrap_op('get_compile_mode') set_tc_util = _wrap_op('set_tc_util') get_tc_util = _wrap_op('get_tc_util') +preload_kernels = _wrap_op('preload_kernels') fp8_gemm_nt = _wrap_op('fp8_gemm_nt') fp8_gemm_nn = _wrap_op('fp8_gemm_nn') @@ -109,7 +110,7 @@ from .utils import * def _verify_ops_loaded(): expected_ops = [ - 'init', 'set_num_sms', 'get_num_sms', 'set_tc_util', 'get_tc_util', + 'init', 'preload_kernels', 'set_num_sms', 'get_num_sms', 'set_tc_util', 'get_tc_util', 'fp8_gemm_nt', 'fp8_gemm_nn', 'fp8_gemm_tn', 'fp8_gemm_tt', 'm_grouped_fp8_gemm_nt_contiguous', 'm_grouped_fp8_gemm_nn_contiguous', 'm_grouped_fp8_gemm_nt_masked', 'k_grouped_fp8_gemm_nt_contiguous', @@ -141,10 +142,10 @@ if __debug__: # Preload cached kernels if enabled via environment variable if os.environ.get('DG_PRELOAD_KERNELS', '0') == '1': try: - deep_gemm_cpp.preload_kernels() + preload_kernels() except Exception as e: # Preloading failure should not block initialization import warnings warnings.warn(f"Failed to preload DeepGEMM kernels: {e}") -__version__ = '2.1.2' \ No newline at end of file +__version__ = '2.2.1' \ No newline at end of file