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

feat: add support for deepgemm sglang

Browse Source
This commit is contained in:
孙远航
2025-11-19 14:37:49 +00:00
parent d763aaf304
commit d4ffb6a9bf
2 changed files with 11 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
csrc/python_api.cpp
View File

@@ -247,6 +247,12 @@ TORCH_LIBRARY(deep_gemm, m) {
deep_gemm_wrappers::init_wrapper(library_root_path, cuda_home_path_by_torch); 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 // 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); 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<int64_t>(deep_gemm::get_mk_alignment_for_contiguous_layout()); return static_cast<int64_t>(deep_gemm::get_mk_alignment_for_contiguous_layout());
} }
REGISTER_EXTENSION(deep_gemm_cpp) REGISTER_EXTENSION(deep_gemm_cpp)

7
deep_gemm/__init__.py
View File

@@ -56,6 +56,7 @@ set_compile_mode = _wrap_op('set_compile_mode')
get_compile_mode = _wrap_op('get_compile_mode') get_compile_mode = _wrap_op('get_compile_mode')
set_tc_util = _wrap_op('set_tc_util') set_tc_util = _wrap_op('set_tc_util')
get_tc_util = _wrap_op('get_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_nt = _wrap_op('fp8_gemm_nt')
fp8_gemm_nn = _wrap_op('fp8_gemm_nn') fp8_gemm_nn = _wrap_op('fp8_gemm_nn')
@@ -109,7 +110,7 @@ from .utils import *
def _verify_ops_loaded(): def _verify_ops_loaded():
expected_ops = [ 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', '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_contiguous', 'm_grouped_fp8_gemm_nn_contiguous',
'm_grouped_fp8_gemm_nt_masked', 'k_grouped_fp8_gemm_nt_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 # Preload cached kernels if enabled via environment variable
if os.environ.get('DG_PRELOAD_KERNELS', '0') == '1': if os.environ.get('DG_PRELOAD_KERNELS', '0') == '1':
try: try:
deep_gemm_cpp.preload_kernels() preload_kernels()
except Exception as e: except Exception as e:
# Preloading failure should not block initialization # Preloading failure should not block initialization
import warnings import warnings
warnings.warn(f"Failed to preload DeepGEMM kernels: {e}") warnings.warn(f"Failed to preload DeepGEMM kernels: {e}")
__version__ = '2.1.2' __version__ = '2.2.1'