A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 9, 10:00 UTC
[Bugfix][CPU] Fix thread num for shared memory communication (#33317)
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>
Li, Jiang committed
Jan 29, 2026 at 19:26 UTC
8311f083bd53ac761556f2fb8733e0cf22ef7024
3 files changed
+25
-10
csrc/cpu/shm.cpp
+8
-7
@@ -237,10 +237,10 @@ struct ThreadSHMContext {
237
class SHMManager {
238
public:
239
explicit SHMManager(const std::string& name, const int rank,
240
- const int group_size)
240
+ const int group_size, const int thread_num)
241
: _rank(rank),
242
_group_size(group_size),
243
- _thread_num(omp_get_max_threads()),
243
+ _thread_num(thread_num),
244
_shm_names({""}),
245
_shared_mem_ptrs({nullptr}),
246
_shm_ctx(nullptr) {
@@ -282,11 +282,11 @@ class SHMManager {
282
}
283
284
static int64_t create_singleton_instance(const std::string& name,
285
- const int group_size,
286
- const int rank) {
285
+ const int group_size, const int rank,
286
+ const int thread_num) {
287
std::lock_guard<std::mutex> guard(SingletonInstancesLock);
288
SingletonInstances.emplace_back(
289
- std::make_unique<SHMManager>(name, rank, group_size));
289
+ std::make_unique<SHMManager>(name, rank, group_size, thread_num));
290
return static_cast<int64_t>(SingletonInstances.size() - 1);
291
}
292
@@ -854,8 +854,9 @@ std::vector<torch::Tensor> shm_recv_tensor_list(int64_t handle, int64_t src) {
854
}
855
856
int64_t init_shm_manager(const std::string& name, const int64_t group_size,
857
- const int64_t rank) {
858
- return SHMManager::create_singleton_instance(name, group_size, rank);
857
+ const int64_t rank, const int64_t thread_num) {
858
+ return SHMManager::create_singleton_instance(name, group_size, rank,
859
+ thread_num);
860
}
861
862
std::string join_shm_manager(int64_t handle, const std::string& name) {
csrc/cpu/torch_bindings.cpp
+5
-3
@@ -35,7 +35,7 @@ void mla_decode_kvcache(torch::Tensor& out, torch::Tensor& query,
35
torch::Tensor& block_tables, torch::Tensor& seq_lens);
36
37
int64_t init_shm_manager(const std::string& name, const int64_t group_size,
38
- const int64_t rank);
38
+ const int64_t rank, const int64_t thread_num);
39
40
std::string join_shm_manager(int64_t handle, const std::string& name);
41
@@ -232,8 +232,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
232
233
// SHM CCL
234
#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__))
235
- ops.def("init_shm_manager(str name, int group_size, int rank) -> int",
236
- &init_shm_manager);
235
+ ops.def(
236
+ "init_shm_manager(str name, int group_size, int rank, int thread_num) -> "
237
+ "int",
238
+ &init_shm_manager);
239
ops.def("join_shm_manager(int handle, str name) -> str", &join_shm_manager);
240
ops.def("shm_allreduce(int handle, Tensor! data) -> ()");
241
ops.impl("shm_allreduce", torch::kCPU, &shm_allreduce);
vllm/distributed/device_communicators/cpu_communicator.py
+12
@@ -205,10 +205,22 @@ class _CPUSHMDistributed:
205
self.handle = self._init_cpu_shm()
206
207
def _init_cpu_shm(self) -> int:
208
+ thread_num_tensor = torch.tensor(
209
+ [torch.get_num_threads()],
210
+ dtype=torch.int64,
211
+ )
212
+ torch.distributed.all_reduce(
213
+ thread_num_tensor,
214
+ op=torch.distributed.ReduceOp.MIN,
215
+ group=self.communicator.device_group,
216
+ )
217
+ thread_num = thread_num_tensor.item()
218
+
219
handle = torch.ops._C.init_shm_manager(
220
self.group_name,
221
self.communicator.world_size,
222
self.communicator.rank,
223
+ thread_num,
224
)
225
torch.distributed.barrier(self.communicator.device_group)
226
torch.ops._C.join_shm_manager(