diff --git a/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache_infiniops.cc b/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache_infiniops.cc index 6fdac33e3..801358b95 100644 --- a/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache_infiniops.cc +++ b/csrc/infinicore/src/ops/mha_kvcache/mha_kvcache_infiniops.cc @@ -25,7 +25,7 @@ std::size_t implementation_index_for_device( return 0; } if (device_type == infini::ops::Device::Type::kMoore) { - return 8; + return 16; } return 16; } diff --git a/csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen_infiniops.cc b/csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen_infiniops.cc index 043ce070a..6c456423d 100644 --- a/csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen_infiniops.cc +++ b/csrc/infinicore/src/ops/multi_head_attention_varlen/mha_varlen_infiniops.cc @@ -23,7 +23,7 @@ std::size_t implementation_index_for_device( return 0; } if (device_type == infini::ops::Device::Type::kMoore) { - return 8; + return 16; } return 16; } diff --git a/csrc/pybind11/engine/engine.hpp b/csrc/pybind11/engine/engine.hpp index c5e85577c..24419cec6 100644 --- a/csrc/pybind11/engine/engine.hpp +++ b/csrc/pybind11/engine/engine.hpp @@ -121,19 +121,15 @@ inline void bind_infer_engine(py::module &m) { } return state_dict_tp_all; }) - .def("process_weights_after_loading", &InferEngine::process_weights_after_loading, "Process the weights after loading on all workers (e.g., for quantization)") - .def( - "forward", [](InferEngine &self, const InferEngine::Input &input) -> InferEngine::Output { + .def("process_weights_after_loading", &InferEngine::process_weights_after_loading, "Process the weights after loading on all workers (e.g., for quantization)", py::call_guard()) + .def("forward", [](InferEngine &self, const InferEngine::Input &input) -> InferEngine::Output { // IMPORTANT: Release the GIL before calling forward() to allow other Python threads // to run concurrently during inference (which may block for a long time). // Do NOT remove this — without it, the GIL is held throughout inference and will // deadlock or stall any other Python thread (e.g., request handling, scheduling). py::gil_scoped_release release; - return self.forward(input); - }, - "Run inference on all ranks with arbitrary arguments") - .def( - "reset_cache", [](InferEngine &self, std::shared_ptr cfg) { self.reset_cache(cfg ? cfg.get() : nullptr); }, py::arg("cache_config") = py::none()) + return self.forward(input); }, "Run inference on all ranks with arbitrary arguments") + .def("reset_cache", [](InferEngine &self, std::shared_ptr cfg) { self.reset_cache(cfg ? cfg.get() : nullptr); }, py::arg("cache_config") = py::none()) .def("get_kv_cache", &InferEngine::get_kv_cache, "Get per-rank kv cache list") .def("get_cache_config", [](const InferEngine &self) -> std::shared_ptr { auto cfg = self.get_cache_config();