biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Bugfix][CPU] Fix thread num for shared memory communication (#33317)

Browse Source 
Signed-off-by: jiang1.li <jiang1.li@intel.com>
Signed-off-by: Li, Jiang <bigpyj64@gmail.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Li, Jiang
2026-01-29 19:26:58 +08:00
committed by GitHub
parent 40c35038d2
commit 8311f083bd
3 changed files with 25 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
csrc/cpu/shm.cpp
View File

@@ -237,10 +237,10 @@ struct ThreadSHMContext {
class SHMManager {
public:
explicit SHMManager(const std::string& name, const int rank,
const int group_size)
const int group_size, const int thread_num)
: _rank(rank),
_group_size(group_size),
_thread_num(omp_get_max_threads()),
_thread_num(thread_num),
_shm_names({""}),
_shared_mem_ptrs({nullptr}),
_shm_ctx(nullptr) {
@@ -282,11 +282,11 @@ class SHMManager {
}
static int64_t create_singleton_instance(const std::string& name,
const int group_size,
const int rank) {
const int group_size, const int rank,
const int thread_num) {
std::lock_guard<std::mutex> guard(SingletonInstancesLock);
SingletonInstances.emplace_back(
std::make_unique<SHMManager>(name, rank, group_size));
std::make_unique<SHMManager>(name, rank, group_size, thread_num));
return static_cast<int64_t>(SingletonInstances.size() - 1);
}
@@ -854,8 +854,9 @@ std::vector<torch::Tensor> shm_recv_tensor_list(int64_t handle, int64_t src) {
}
int64_t init_shm_manager(const std::string& name, const int64_t group_size,
const int64_t rank) {
return SHMManager::create_singleton_instance(name, group_size, rank);
const int64_t rank, const int64_t thread_num) {
return SHMManager::create_singleton_instance(name, group_size, rank,
thread_num);
}
std::string join_shm_manager(int64_t handle, const std::string& name) {