mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-26 16:35:05 +00:00
18 lines
626 B
C++
18 lines
626 B
C++
#include <torch/extension.h>
|
|
|
|
#include "cuda_utils.hpp"
|
|
#include "engine/pybind.hpp"
|
|
#include "graph_timer/graph_timer.hpp"
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
|
// Cuda utils
|
|
pybind11::module cuda_utils = m.def_submodule("cuda_utils", "DeepInfer cuda utils");
|
|
deepinfer::cuda::utils::pymodule_init(cuda_utils);
|
|
// Engine
|
|
pybind11::module engine = m.def_submodule("engine");
|
|
deepinfer::engine::pymodule_init(engine);
|
|
// CUDA graph 内探针计时
|
|
pybind11::module graph_timer = m.def_submodule("graph_timer", "CUDA graph timer");
|
|
deepinfer::graph_timer::pymodule_init(graph_timer);
|
|
}
|