@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 550b2801ad
Fork 0

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 15, 16:00 UTC

[CPU][Bugfix] Using custom allreduce for CPU backend (#15934)

Signed-off-by: jiang1.li <jiang1.li@intel.com>

Li, Jiang committed Apr 2, 2025 at 22:46 UTC 550b2801ad57b8f1e1782037c40fd01b2632435e
8 files changed +1013 -16
cmake/cpu_extension.cmake
+1
@@ -197,6 +197,7 @@ set(VLLM_EXT_SRC
197 if (AVX512_FOUND AND NOT AVX512_DISABLED)
198 set(VLLM_EXT_SRC
199 "csrc/cpu/quant.cpp"
200 + "csrc/cpu/shm.cpp"
201 ${VLLM_EXT_SRC})
202 endif()
203
csrc/cpu/cpu_types_x86.hpp
+58
@@ -78,9 +78,14 @@ struct FP16Vec16 : public Vec<FP16Vec16> {
78
79 __m256i reg;
80
81 + // normal load
82 explicit FP16Vec16(const void* ptr)
83 : reg((__m256i)_mm256_loadu_si256((__m256i*)ptr)) {}
84
85 + // non-temproal load
86 + explicit FP16Vec16(bool, void* ptr)
87 + : reg(_mm256_stream_load_si256((__m256i*)ptr)) {}
88 +
89 explicit FP16Vec16(const FP32Vec16&);
90
91 void save(void* ptr) const { *reinterpret_cast<__m256i*>(ptr) = reg; }
@@ -110,9 +115,14 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
115
116 __m256i reg;
117
118 + // normal load
119 explicit BF16Vec16(const void* ptr)
120 : reg((__m256i)_mm256_loadu_si256((__m256i*)ptr)) {}
121
122 + // non-temproal load
123 + explicit BF16Vec16(bool, void* ptr)
124 + : reg(_mm256_stream_load_si256((__m256i*)ptr)) {}
125 +
126 explicit BF16Vec16(const FP32Vec16&);
127
128 void save(void* ptr) const { *reinterpret_cast<__m256i*>(ptr) = reg; }
@@ -313,8 +323,13 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
323
324 explicit FP32Vec16() : reg(_mm512_set1_ps(0.0)) {}
325
326 + // normal load
327 explicit FP32Vec16(const float* ptr) : reg(_mm512_loadu_ps(ptr)) {}
328
329 + // non-temproal load
330 + explicit FP32Vec16(bool, void* ptr)
331 + : reg((__m512)_mm512_stream_load_si512(ptr)) {}
332 +
333 explicit FP32Vec16(__m512 data) : reg(data) {}
334
335 explicit FP32Vec16(const FP32Vec4& data)
@@ -547,6 +562,33 @@ struct INT8Vec16 : public Vec<INT8Vec16> {
562 _mm_mask_storeu_epi8(ptr, mask, reg);
563 }
564 };
565 +
566 +struct INT8Vec64 : public Vec<INT8Vec64> {
567 + constexpr static int VEC_ELEM_NUM = 64;
568 + union AliasReg {
569 + __m512i reg;
570 + int8_t values[VEC_ELEM_NUM];
571 + };
572 +
573 + __m512i reg;
574 +
575 + // normal load
576 + explicit INT8Vec64(void* ptr) : reg(_mm512_loadu_epi8(ptr)) {}
577 +
578 + // non-temproal load
579 + explicit INT8Vec64(bool, void* ptr) : reg(_mm512_stream_load_si512(ptr)) {}
580 +
581 + void save(void* ptr) const { _mm512_storeu_epi8(ptr, reg); }
582 +
583 + void save(int8_t* ptr, const int elem_num) const {
584 + constexpr uint64_t M = 0xFFFFFFFFFFFFFFFF;
585 + __mmask64 mask = _cvtu64_mask64(M >> (64 - elem_num));
586 + _mm512_mask_storeu_epi8(ptr, mask, reg);
587 + }
588 +
589 + // non-temproal save
590 + void nt_save(int8_t* ptr) { _mm512_stream_si512((__m512i*)ptr, reg); }
591 +};
592 #endif
593
594 template <typename T>
@@ -657,6 +699,22 @@ inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
699
700 inline void prefetch(const void* addr) { _mm_prefetch(addr, _MM_HINT_T1); }
701
702 +#ifdef __AVX512F__
703 +inline void non_temporal_save(FP16Vec16& vec, void* ptr) {
704 + _mm256_stream_si256((__m256i*)ptr, vec.reg);
705 +}
706 +inline void non_temporal_save(BF16Vec32& vec, void* ptr) {
707 + _mm512_stream_si512((__m512i*)ptr, vec.reg);
708 +}
709 +inline void non_temporal_save(BF16Vec16& vec, void* ptr) {
710 + _mm256_stream_si256((__m256i*)ptr, vec.reg);
711 +}
712 +inline void non_temporal_save(FP32Vec16& vec, void* ptr) {
713 + _mm512_stream_ps((float*)ptr, vec.reg);
714 +}
715 +#endif
716 +
717 +inline void mem_barrier() { _mm_mfence(); }
718 }; // namespace vec_op
719
720 #endif
csrc/cpu/shm.cpp new
+781
@@ -0,0 +1,781 @@
1 +#include "cpu/cpu_types.hpp"
2 +
3 +#include <fcntl.h>
4 +#include <sys/mman.h>
5 +#include <sys/stat.h>
6 +#include <unistd.h>
7 +
8 +namespace {
9 +#define MAX_SHM_RANK_NUM 8
10 +#define MAX_THREAD_NUM 12
11 +#define PER_THREAD_SHM_BUFFER_BYTES (4 * 1024 * 1024)
12 +#define MIN_THREAD_PROCESS_SIZE (8 * 1024)
13 +#define MAX_P2P_SEND_TENSOR_NUM 8
14 +
15 +template <typename scalar_t>
16 +struct KernelVecType {
17 + using scalar_vec_t = void;
18 +};
19 +
20 +template <>
21 +struct KernelVecType<float> {
22 + using scalar_vec_t = vec_op::FP32Vec16;
23 +};
24 +
25 +template <>
26 +struct KernelVecType<c10::BFloat16> {
27 + using scalar_vec_t = vec_op::BF16Vec16;
28 +};
29 +
30 +template <>
31 +struct KernelVecType<c10::Half> {
32 + using scalar_vec_t = vec_op::FP16Vec16;
33 +};
34 +
35 +enum class ThreadSHMStat : char { THREAD_READY = 0, SHM_DATA_READY, DONE };
36 +
37 +struct ThreadSHMContext {
38 + volatile ThreadSHMStat thread_stats[MAX_SHM_RANK_NUM];
39 + int thread_id;
40 + int thread_num;
41 + int rank;
42 + int group_size;
43 + size_t _spinning_count;
44 + int swizzled_ranks[MAX_SHM_RANK_NUM];
45 + void* thread_shm_ptrs[MAX_SHM_RANK_NUM];
46 + ThreadSHMContext* shm_contexts[MAX_SHM_RANK_NUM];
47 +
48 + ThreadSHMContext(const int thread_id, const int thread_num, const int rank,
49 + const int group_size, void* thread_shm_ptr)
50 + : thread_id(thread_id),
51 + thread_num(thread_num),
52 + rank(rank),
53 + group_size(group_size),
54 + _spinning_count(0) {
55 + static_assert(sizeof(ThreadSHMContext) % 64 == 0);
56 + TORCH_CHECK(group_size <= MAX_SHM_RANK_NUM);
57 + TORCH_CHECK((size_t)this % 64 == 0);
58 + TORCH_CHECK((size_t)thread_shm_ptr % 64 == 0);
59 + for (int i = 0; i < MAX_SHM_RANK_NUM; ++i) {
60 + shm_contexts[i] = nullptr;
61 + thread_shm_ptrs[i] = nullptr;
62 + swizzled_ranks[i] = (i + rank) % group_size;
63 + thread_stats[i] = ThreadSHMStat::DONE;
64 + }
65 + set_context(rank, this, thread_shm_ptr);
66 + }
67 +
68 + void set_context(int rank, ThreadSHMContext* ptr, void* thread_shm_ptr) {
69 + TORCH_CHECK(rank < MAX_SHM_RANK_NUM);
70 + TORCH_CHECK(ptr);
71 + TORCH_CHECK(thread_shm_ptr);
72 + TORCH_CHECK_EQ(ptr->thread_num, thread_num);
73 + TORCH_CHECK_EQ(ptr->thread_id, thread_id);
74 + shm_contexts[rank] = ptr;
75 + thread_shm_ptrs[rank] = thread_shm_ptr;
76 + }
77 +
78 + template <typename T>
79 + T* get_thread_shm_ptr(int rank) {
80 + return reinterpret_cast<T*>(thread_shm_ptrs[rank]);
81 + }
82 +
83 + int get_swizzled_rank(int idx) { return swizzled_ranks[idx]; }
84 +
85 + void wait_for_all(ThreadSHMStat prev_stat) {
86 + for (int idx = 0; idx < group_size; ++idx) {
87 + int rank = get_swizzled_rank(idx);
88 + while (thread_stats[rank] == prev_stat) {
89 + ++_spinning_count;
90 + _mm_pause();
91 + }
92 + }
93 + vec_op::mem_barrier();
94 + }
95 +
96 + void wait_for_one(int rank, ThreadSHMStat prev_stat) {
97 + while (thread_stats[rank] == prev_stat) {
98 + ++_spinning_count;
99 + _mm_pause();
100 + }
101 + vec_op::mem_barrier();
102 + }
103 +
104 + void set_thread_stat(ThreadSHMStat stat) {
105 + for (int idx = 0; idx < group_size; ++idx) {
106 + int rank = get_swizzled_rank(idx);
107 + shm_contexts[rank]->thread_stats[this->rank] = stat;
108 + }
109 + }
110 +
111 + void set_thread_stat(int target_rank, ThreadSHMStat stat) {
112 + for (int idx = 0; idx < group_size; ++idx) {
113 + int rank = get_swizzled_rank(idx);
114 + shm_contexts[rank]->thread_stats[target_rank] = stat;
115 + }
116 + }
117 +
118 + // barrier for all ranks in the group, used for all2all ops
119 + // DONE -> THREAD_READY -> SHM_DATA_READY -> DONE -> ...
120 + void barrier(ThreadSHMStat next_stat) {
121 + if (next_stat == ThreadSHMStat::THREAD_READY) {
122 + set_thread_stat(ThreadSHMStat::THREAD_READY);
123 + wait_for_all(ThreadSHMStat::DONE);
124 + } else if (next_stat == ThreadSHMStat::SHM_DATA_READY) {
125 + set_thread_stat(ThreadSHMStat::SHM_DATA_READY);
126 + wait_for_all(ThreadSHMStat::THREAD_READY);
127 + } else if (next_stat == ThreadSHMStat::DONE) {
128 + set_thread_stat(ThreadSHMStat::DONE);
129 + wait_for_all(ThreadSHMStat::SHM_DATA_READY);
130 + } else {
131 + TORCH_CHECK(false, "Invalid next_stat to barrier.");
132 + }
133 + }
134 +
135 + std::string to_string() const {
136 + std::stringstream ss;
137 + ss << "SHMContext:";
138 + ss << "\nrank: " << rank;
139 + ss << "\ngroup_size: " << group_size;
140 + ss << "\nthread_num: " << thread_num;
141 + ss << "\nthread_id: " << thread_id;
142 +
143 + ss << "\nshm_ctx_stat_loop_seq: [";
144 + for (int i = 0; i < group_size; ++i) {
145 + ss << swizzled_ranks[i] << ", ";
146 + }
147 + ss << "]";
148 +
149 + ss << "\nshm_contexts: [";
150 + for (int i = 0; i < group_size; ++i) {
151 + if (shm_contexts[i]) {
152 + ss << shm_contexts[i]->rank << ", ";
153 + }
154 + }
155 + ss << "]";
156 +
157 + return ss.str();
158 + }
159 +};
160 +
161 +class SHMManager {
162 + public:
163 + explicit SHMManager(const std::string& name, const int rank,
164 + const int group_size)
165 + : _rank(rank),
166 + _group_size(group_size),
167 + _thread_num(std::min(torch::get_num_threads(), MAX_THREAD_NUM)),
168 + _shm_names({""}),
169 + _shared_mem_ptrs({nullptr}),
170 + _shm_ctx(nullptr) {
171 + _shm_names[rank] = get_shm_name(name, rank);
172 + _shared_mem_ptrs[rank] = init_shm(rank);
173 + _shm_ctx = reinterpret_cast<ThreadSHMContext*>(_shared_mem_ptrs[rank]);
174 +
175 + for (int i = 0; i < _thread_num; ++i) {
176 + ThreadSHMContext* ctx = new (_shm_ctx + i)
177 + ThreadSHMContext(i, _thread_num, _rank, _group_size,
178 + compute_thread_shm_ptr(_shm_ctx, i));
179 + }
180 + }
181 +
182 + void join(const std::string& name) {
183 + for (int rank_idx = 0; rank_idx < _group_size; ++rank_idx) {
184 + if (rank_idx != _rank) {
185 + TORCH_CHECK(_shm_names[rank_idx].empty());
186 + TORCH_CHECK(_shared_mem_ptrs[rank_idx] == nullptr);
187 + _shm_names[rank_idx] = get_shm_name(name, rank_idx);
188 + _shared_mem_ptrs[rank_idx] = init_shm(rank_idx);
189 + ThreadSHMContext* target_ctx =
190 + reinterpret_cast<ThreadSHMContext*>(_shared_mem_ptrs[rank_idx]);
191 + for (int thread_idx = 0; thread_idx < _thread_num; ++thread_idx) {
192 + _shm_ctx[thread_idx].set_context(
193 + rank_idx, target_ctx + thread_idx,
194 + compute_thread_shm_ptr(target_ctx, thread_idx));
195 + }
196 + }
197 + }
198 + }
199 +
200 + ~SHMManager() { destroy_shm(); }
201 +
202 + ThreadSHMContext* get_shm_ctx() const { return _shm_ctx; }
203 +
204 + static std::string get_shm_name(const std::string& name, int rank) {
205 + return name + "_" + std::to_string(rank);
206 + }
207 +
208 + static int64_t create_singleton_instance(const std::string& name,
209 + const int group_size,
210 + const int rank) {
211 + std::lock_guard<std::mutex> guard(SingletonInstancesLock);
212 + SingletonInstances.emplace_back(
213 + std::make_unique<SHMManager>(name, rank, group_size));
214 + return static_cast<int64_t>(SingletonInstances.size() - 1);
215 + }
216 +
217 + static SHMManager* get_singleton_instance(int64_t handle) {
218 + return SingletonInstances[handle].get();
219 + }
220 +
221 + protected:
222 + static std::vector<std::unique_ptr<SHMManager>> SingletonInstances;
223 + static std::mutex SingletonInstancesLock;
224 +
225 + private:
226 + static size_t round_to_alignment(size_t num) {
227 + return ((num + 63) / 64) * 64;
228 + }
229 +
230 + int8_t* compute_thread_shm_ptr(ThreadSHMContext* ctx, int thread_id) {
231 + int8_t* thread_shm_ptr =
232 + reinterpret_cast<int8_t*>(ctx) +
233 + round_to_alignment(_thread_num * sizeof(ThreadSHMContext));
234 + return thread_shm_ptr +
235 + thread_id * round_to_alignment(PER_THREAD_SHM_BUFFER_BYTES);
236 + }
237 +
238 + size_t compute_shm_size() {
239 + const size_t rounded_rank_buffer_size =
240 + round_to_alignment(PER_THREAD_SHM_BUFFER_BYTES) * _thread_num;
241 + const size_t rounded_thread_shm_ctx_size =
242 + round_to_alignment(_thread_num * sizeof(ThreadSHMContext));
243 + const size_t shm_size =
244 + rounded_thread_shm_ctx_size + rounded_rank_buffer_size;
245 + return shm_size;
246 + }
247 +
248 + void* init_shm(int target_rank) {
249 + const std::string& shm_name = _shm_names[target_rank];
250 + const int local_rank = _rank;
251 + const size_t shm_size = compute_shm_size();
252 +
253 + int fd = -1;
254 + if (local_rank == target_rank) {
255 + fd = shm_open(shm_name.c_str(), O_CREAT | O_EXCL | O_RDWR,
256 + S_IRUSR | S_IWUSR);
257 +
258 + if (fd == -1)
259 + TORCH_CHECK(false, "create shm in SHMManager failed. errno: " +
260 + std::to_string(errno));
261 +
262 + if (ftruncate(fd, shm_size) == -1)
263 + TORCH_CHECK(false, "ftruncate in SHMManager failed. errno: " +
264 + std::to_string(errno));
265 + } else {
266 + fd = shm_open(shm_name.c_str(), O_RDWR, S_IRUSR | S_IWUSR);
267 +
268 + if (fd == -1)
269 + TORCH_CHECK(false, "open shm in SHMManager failed. errno: " +
270 + std::to_string(errno));
271 + }
272 +
273 + void* shm_ptr = mmap(nullptr, shm_size, PROT_READ | PROT_WRITE,
274 + MAP_SHARED | MAP_POPULATE, fd, 0);
275 +
276 + if (shm_ptr == MAP_FAILED) {
277 + TORCH_CHECK(false,
278 + "mmap in SHMManager failed. errno: " + std::to_string(errno));
279 + }
280 +
281 + if (close(fd) != 0) {
282 + TORCH_CHECK(
283 + false, "close in SHMManager failed. errno: " + std::to_string(errno));
284 + }
285 +
286 + TORCH_CHECK((size_t)shm_ptr % 64 == 0);
287 +
288 + return shm_ptr;
289 + }
290 +
291 + void destroy_shm() {
292 + std::stringstream ss;
293 + ss << "local rank " << _rank << ": [";
294 + for (int thread_id = 0; thread_id < _thread_num; ++thread_id) {
295 + ss << _shm_ctx[thread_id]._spinning_count << ", ";
296 + }
297 + ss << "]\n";
298 +
299 + for (int i = 0; i < MAX_SHM_RANK_NUM; ++i) {
300 + if (_shared_mem_ptrs[i] != nullptr) {
301 + munmap(_shared_mem_ptrs[i], compute_shm_size());
302 + }
303 +
304 + if (!_shm_names[i].empty()) {
305 + shm_unlink(_shm_names[i].c_str());
306 + }
307 + }
308 + }
309 +
310 + int _rank;
311 + int _group_size;
312 + int _thread_num;
313 + std::array<std::string, MAX_SHM_RANK_NUM> _shm_names;
314 + std::array<void*, MAX_SHM_RANK_NUM> _shared_mem_ptrs;
315 + ThreadSHMContext* _shm_ctx;
316 +};
317 +
318 +namespace shm_cc_ops {
319 +template <typename scalar_t, typename F>
320 +void shm_cc_loop(ThreadSHMContext* ctx, int64_t elem_num, F&& inner_func) {
321 + int thread_num = ctx->thread_num;
322 + int64_t total_bytes = elem_num * sizeof(scalar_t);
323 + int64_t total_units_num =
324 + (total_bytes + MIN_THREAD_PROCESS_SIZE - 1) / MIN_THREAD_PROCESS_SIZE;
325 + int64_t per_thread_units_num =
326 + (total_units_num + thread_num - 1) / thread_num;
327 + int64_t per_unit_elem_num = MIN_THREAD_PROCESS_SIZE / sizeof(scalar_t);
328 + int64_t max_per_thread_iteration_elem_num =
329 + PER_THREAD_SHM_BUFFER_BYTES / sizeof(scalar_t);
330 + int64_t per_thread_elem_num = per_unit_elem_num * per_thread_units_num;
331 +
332 +#pragma omp parallel for schedule(static, 1)
333 + for (int i = 0; i < thread_num; ++i) {
334 + int64_t offset = i * per_thread_elem_num;
335 + int64_t end = std::min(elem_num, offset + per_thread_elem_num);
336 + int64_t curr_elem_num =
337 + std::min(max_per_thread_iteration_elem_num, end - offset);
338 + ThreadSHMContext* thread_ctx = ctx + i;
339 +
340 + while (curr_elem_num > 0) {
341 + inner_func(thread_ctx, offset, curr_elem_num);
342 +
343 + offset += max_per_thread_iteration_elem_num;
344 + curr_elem_num = std::min(max_per_thread_iteration_elem_num, end - offset);
345 + }
346 + }
347 +}
348 +}; // namespace shm_cc_ops
349 +
350 +namespace shm_cc_ops {
351 +
352 +void memcpy_from_shm(void* dst, void* src, const int64_t bytes) {
353 + const int64_t aligned_bytes = ((bytes >> 6) << 6); // 64 bytes aligned
354 + int64_t i = 0;
355 +#pragma GCC unroll 4
356 + for (; i < aligned_bytes; i += 64) {
357 + vec_op::INT8Vec64 data(
358 + true, (int8_t*)src + i); // stream loading shm to avoid caching
359 + data.save((int8_t*)dst + i);
360 + }
361 + if (aligned_bytes < bytes) {
362 + vec_op::INT8Vec64 data(true, (int8_t*)src + aligned_bytes);
363 + data.save((int8_t*)dst + aligned_bytes, bytes - aligned_bytes);
364 + }
365 +}
366 +
367 +void memcpy_to_shm(void* dst, void* src, const int64_t bytes) {
368 +#pragma GCC unroll 4
369 + for (int64_t i = 0; i < bytes; i += 64) {
370 + vec_op::INT8Vec64 data((int8_t*)src + i);
371 + data.nt_save((int8_t*)dst + i);
372 + }
373 +}
374 +
375 +void memcpy(void* dst, void* src, const int64_t bytes) {
376 + const int64_t aligned_bytes = ((bytes >> 6) << 6); // 64 bytes aligned
377 + int64_t i = 0;
378 +#pragma GCC unroll 4
379 + for (; i < aligned_bytes; i += 64) {
380 + vec_op::INT8Vec64 data((int8_t*)src + i);
381 + data.save((int8_t*)dst + i);
382 + }
383 + if (aligned_bytes < bytes) {
384 + vec_op::INT8Vec64 data((int8_t*)src + aligned_bytes);
385 + data.save((int8_t*)dst + aligned_bytes, bytes - aligned_bytes);
386 + }
387 +}
388 +
389 +template <typename scalar_t, int RANKS>
390 +void all_reduce_sum_impl(ThreadSHMContext* ctx, scalar_t* data,
391 + size_t elem_num) {
392 + CPU_KERNEL_GUARD_IN(all_reduce_sum_impl)
393 + using vec_t = typename KernelVecType<scalar_t>::scalar_vec_t;
394 + constexpr int64_t vec_elem_num = vec_t::get_elem_num();
395 + const int worldsize = ctx->group_size;
396 +
397 + shm_cc_ops::shm_cc_loop<scalar_t>(
398 + ctx, elem_num,
399 + [&](ThreadSHMContext* thread_ctx, int64_t data_offset,
400 + int64_t data_elem_num) {
401 + int rank = thread_ctx->rank;
402 + scalar_t* thread_shm_ptr =
403 + thread_ctx->get_thread_shm_ptr<scalar_t>(rank);
404 + scalar_t* thread_data_ptr = data + data_offset;
405 + int64_t thread_data_elem_num = data_elem_num * sizeof(scalar_t);
406 +
407 + scalar_t* remote_data_ptrs[RANKS - 1];
408 + vec_op::unroll_loop<int, RANKS - 1>([&](int idx) {
409 + remote_data_ptrs[idx] = thread_ctx->get_thread_shm_ptr<scalar_t>(
410 + thread_ctx->get_swizzled_rank(idx + 1));
411 + });
412 +
413 + thread_ctx->barrier(ThreadSHMStat::THREAD_READY);
414 +
415 + shm_cc_ops::memcpy_to_shm(thread_shm_ptr, thread_data_ptr,
416 + thread_data_elem_num);
417 +
418 + thread_ctx->barrier(ThreadSHMStat::SHM_DATA_READY);
419 +
420 + int64_t aligned_data_elem_num =
421 + (data_elem_num / vec_elem_num) * vec_elem_num;
422 + int64_t i = 0;
423 +#pragma GCC unroll 4
424 + for (; i < aligned_data_elem_num; i += vec_elem_num) {
425 + vec_t local_data(thread_data_ptr + i); // load from cache
426 + vec_op::FP32Vec16 local_data_fp32(local_data);
427 + vec_op::unroll_loop<int, RANKS - 1>([&](int idx) {
428 + vec_t remote_data(
429 + true, remote_data_ptrs[idx] + i); // stream load from shm
430 + vec_op::FP32Vec16 remote_data_fp32(remote_data);
431 + local_data_fp32 = local_data_fp32 + remote_data_fp32; // sum reduce
432 + });
433 + vec_t reduced_data(local_data_fp32);
434 + reduced_data.save(thread_data_ptr + i);
435 + }
436 +
437 + if (i < data_elem_num) {
438 + vec_t local_data(thread_data_ptr + i); // load from cache
439 + vec_op::FP32Vec16 local_data_fp32(local_data);
440 + vec_op::unroll_loop<int, RANKS - 1>([&](int idx) {
441 + vec_t remote_data(
442 + true, remote_data_ptrs[idx] + i); // stream load from shm
443 + vec_op::FP32Vec16 remote_data_fp32(remote_data);
444 + local_data_fp32 = local_data_fp32 + remote_data_fp32; // sum reduce
445 + });
446 + vec_t reduced_data(local_data_fp32);
447 + reduced_data.save(thread_data_ptr + i,
448 + data_elem_num - aligned_data_elem_num);
449 + }
450 +
451 + thread_ctx->barrier(ThreadSHMStat::DONE);
452 + });
453 +
454 + return;
455 +}
456 +}; // namespace shm_cc_ops
457 +
458 +std::vector<std::unique_ptr<SHMManager>> SHMManager::SingletonInstances = {};
459 +std::mutex SHMManager::SingletonInstancesLock = {};
460 +
461 +template <typename scalar_t>
462 +void shm_allreduce_sum(ThreadSHMContext* ctx, scalar_t* data, size_t elem_num) {
463 + switch (ctx->group_size) {
464 + case 2:
465 + shm_cc_ops::all_reduce_sum_impl<scalar_t, 2>(ctx, data, elem_num);
466 + break;
467 + case 3:
468 + shm_cc_ops::all_reduce_sum_impl<scalar_t, 3>(ctx, data, elem_num);
469 + break;
470 + case 4:
471 + shm_cc_ops::all_reduce_sum_impl<scalar_t, 4>(ctx, data, elem_num);
472 + break;
473 + case 8:
474 + shm_cc_ops::all_reduce_sum_impl<scalar_t, 8>(ctx, data, elem_num);
475 + break;
476 + default:
477 + TORCH_CHECK(false,
478 + "Invalid world size: " + std::to_string(ctx->group_size));
479 + }
480 +}
481 +
482 +template <typename scalar_t>
483 +void shm_gather_impl(ThreadSHMContext* ctx, scalar_t* data, size_t elem_num,
484 + scalar_t** outputs, const int dst) {
485 + CPU_KERNEL_GUARD_IN(shm_gather_impl)
486 + const int worldsize = ctx->group_size;
487 + TORCH_CHECK_LT(dst, worldsize);
488 + shm_cc_ops::shm_cc_loop<scalar_t>(
489 + ctx, elem_num,
490 + [&](ThreadSHMContext* thread_ctx, int64_t data_offset,
491 + int64_t data_elem_num) {
492 + int rank = thread_ctx->rank;
493 + scalar_t* thread_shm_ptr =
494 + thread_ctx->get_thread_shm_ptr<scalar_t>(rank);
495 +
496 + thread_ctx->barrier(ThreadSHMStat::THREAD_READY);
497 +
498 + shm_cc_ops::memcpy_to_shm(thread_shm_ptr, data + data_offset,
499 + data_elem_num * sizeof(scalar_t));
500 +
501 + thread_ctx->barrier(ThreadSHMStat::SHM_DATA_READY);
502 +
503 + if (rank == dst) {
504 + shm_cc_ops::memcpy(outputs[rank] + data_offset, data + data_offset,
505 + data_elem_num * sizeof(scalar_t));
506 + for (int i = 1; i < worldsize; ++i) {
507 + int src_rank = thread_ctx->get_swizzled_rank(i);
508 + scalar_t* src_ptr =
509 + thread_ctx->get_thread_shm_ptr<scalar_t>(src_rank); // shm
510 + scalar_t* dst_ptr = outputs[src_rank] + data_offset;
511 + shm_cc_ops::memcpy_from_shm(dst_ptr, src_ptr,
512 + data_elem_num * sizeof(scalar_t));
513 + }
514 + }
515 +
516 + thread_ctx->barrier(ThreadSHMStat::DONE);
517 + });
518 +
519 + return;
520 +}
521 +
522 +struct MemPiece {
523 + void* ptr;
524 + int64_t size;
525 +
526 + template <typename T>
527 + T* data_ptr() {
528 + return reinterpret_cast<T*>(ptr);
529 + }
530 +};
531 +
532 +struct TensorListMeta {
533 + int64_t tensor_bytes[MAX_P2P_SEND_TENSOR_NUM];
534 + torch::ScalarType tensor_types[MAX_P2P_SEND_TENSOR_NUM];
535 + int64_t tensor_num;
536 + int64_t total_bytes;
537 +
538 + TensorListMeta() : tensor_num(0), total_bytes(0) {
539 + static_assert(sizeof(TensorListMeta) % 64 == 0);
540 + static_assert(sizeof(TensorListMeta) <
541 + MIN_THREAD_PROCESS_SIZE); // To ensure the metadata always
542 + // hold by the thread 0
543 + for (int i = 0; i < MAX_P2P_SEND_TENSOR_NUM; ++i) {
544 + tensor_bytes[i] = 0;
545 + tensor_ptrs[i] = nullptr;
546 + tensor_types[i] = torch::ScalarType::Undefined;
547 + }
548 + }
549 +
550 + // For send and recv
551 + void bind_tensor_list(std::vector<torch::Tensor>& tensor_list) {
552 + TORCH_CHECK(tensor_types[0] == torch::ScalarType::Undefined,
553 + "Re-bind TensorListMeta is not allowed.")
554 + TORCH_CHECK_LE(tensor_list.size(), MAX_P2P_SEND_TENSOR_NUM);
555 + tensor_num = tensor_list.size();
556 + int64_t bytes_sum = 0;
557 + for (int i = 0; i < tensor_list.size(); ++i) {
558 + torch::Tensor& t = tensor_list[i];
559 + TORCH_CHECK(t.is_contiguous());
560 + tensor_bytes[i] = t.nbytes();
561 + tensor_types[i] = t.scalar_type();
562 + tensor_ptrs[i] = t.data_ptr();
563 + bytes_sum += t.nbytes();
564 + }
565 + total_bytes = bytes_sum;
566 + }
567 +
568 + // For recv
569 + std::vector<torch::Tensor> generate_tensor_list() {
570 + std::vector<torch::Tensor> tensor_list;
571 + tensor_list.reserve(tensor_num);
572 +
573 + for (int i = 0; i < tensor_num; ++i) {
574 + int64_t bytes = tensor_bytes[i];
575 + auto type = tensor_types[i];
576 + int64_t elem_bytes = torch::elementSize(type);
577 +
578 + TORCH_CHECK_EQ(bytes % elem_bytes, 0);
579 + int64_t elem_num = bytes / elem_bytes;
580 + auto options = torch::TensorOptions().dtype(type).device(torch::kCPU);
581 + tensor_list.emplace_back(torch::empty({elem_num}, options));
582 + }
583 + return tensor_list;
584 + }
585 +
586 + MemPiece get_data(int64_t offset) {
587 + for (int i = 0; i < tensor_num; ++i) {
588 + if (offset < tensor_bytes[i]) {
589 + return {reinterpret_cast<int8_t*>(tensor_ptrs[i]) + offset,
590 + tensor_bytes[i] - offset};
591 + }
592 + offset -= tensor_bytes[i];
593 + }
594 + return {nullptr, 0};
595 + }
596 +
597 + private:
598 + void* tensor_ptrs[MAX_P2P_SEND_TENSOR_NUM];
599 + int8_t _padding[40];
600 +};
601 +
602 +void shm_send_tensor_list_impl(ThreadSHMContext* ctx,
603 + const std::vector<torch::Tensor>& tensor_list) {
604 + CPU_KERNEL_GUARD_IN(shm_send_tensor_list_impl)
605 + std::vector<torch::Tensor> tensor_list_with_metadata;
606 + tensor_list_with_metadata.reserve(1 + tensor_list.size());
607 +
608 + auto options = torch::TensorOptions().dtype(torch::kInt8).device(torch::kCPU);
609 + tensor_list_with_metadata.emplace_back(
610 + torch::empty({sizeof(TensorListMeta)}, options));
611 + tensor_list_with_metadata.insert(tensor_list_with_metadata.end(),
612 + tensor_list.begin(), tensor_list.end());
613 +
614 + torch::Tensor& metadata_tensor = tensor_list_with_metadata[0];
615 + TORCH_CHECK_EQ(metadata_tensor.nbytes(), sizeof(TensorListMeta));
616 +
617 + TensorListMeta* metadata = new (metadata_tensor.data_ptr()) TensorListMeta();
618 + metadata->bind_tensor_list(tensor_list_with_metadata);
619 +
620 + shm_cc_ops::shm_cc_loop<int8_t>(
621 + ctx, metadata->total_bytes,
622 + [&](ThreadSHMContext* thread_ctx, int64_t data_offset,
623 + int64_t data_elem_num) {
624 + int rank = thread_ctx->rank;
625 + // Wait until the receiver set the stat to DONE
626 + thread_ctx->wait_for_one(rank, ThreadSHMStat::SHM_DATA_READY);
627 +
628 + int64_t curr_shm_offset = 0;
629 + while (curr_shm_offset < data_elem_num) {
630 + MemPiece frag = metadata->get_data(data_offset + curr_shm_offset);
631 + frag.size = std::min(frag.size, data_elem_num - curr_shm_offset);
632 + shm_cc_ops::memcpy(
633 + thread_ctx->get_thread_shm_ptr<int8_t>(rank) + curr_shm_offset,
634 + frag.ptr, frag.size);
635 + curr_shm_offset += frag.size;
636 + }
637 +
638 + thread_ctx->set_thread_stat(rank, ThreadSHMStat::SHM_DATA_READY);
639 + });
640 +}
641 +
642 +std::vector<torch::Tensor> shm_recv_tensor_list_impl(ThreadSHMContext* ctx,
643 + int64_t src) {
644 + CPU_KERNEL_GUARD_IN(shm_recv_tensor_list_impl)
645 + auto options = torch::TensorOptions().dtype(torch::kInt8).device(torch::kCPU);
646 + torch::Tensor metadata_tensor =
647 + torch::empty({sizeof(TensorListMeta)}, options);
648 +
649 + // Wait until the sender set the stat of the thread 0 to SHM_DATA_READY
650 + ctx->wait_for_one(src, ThreadSHMStat::DONE);
651 + shm_cc_ops::memcpy(metadata_tensor.data_ptr(),
652 + ctx->get_thread_shm_ptr<void>(src),
653 + sizeof(TensorListMeta));
654 + TensorListMeta* src_metadata =
655 + reinterpret_cast<TensorListMeta*>(metadata_tensor.data_ptr());
656 + std::vector<torch::Tensor> tensor_list_with_metadata =
657 + src_metadata->generate_tensor_list();
658 +
659 + TensorListMeta metadata;
660 + metadata.bind_tensor_list(tensor_list_with_metadata);
661 + TORCH_CHECK_EQ(metadata.tensor_num, src_metadata->tensor_num);
662 + TORCH_CHECK_EQ(metadata.total_bytes, src_metadata->total_bytes);
663 +
664 + shm_cc_ops::shm_cc_loop<int8_t>(
665 + ctx, metadata.total_bytes,
666 + [&](ThreadSHMContext* thread_ctx, int64_t data_offset,
667 + int64_t data_elem_num) {
668 + // Wait until the sender set the stat to SHM_DATA_READY
669 + thread_ctx->wait_for_one(src, ThreadSHMStat::DONE);
670 + int64_t curr_shm_offset = 0;
671 + while (curr_shm_offset < data_elem_num) {
672 + MemPiece frag = metadata.get_data(data_offset + curr_shm_offset);
673 + frag.size = std::min(frag.size, data_elem_num - curr_shm_offset);
674 + shm_cc_ops::memcpy(
675 + frag.ptr,
676 + thread_ctx->get_thread_shm_ptr<int8_t>(src) + curr_shm_offset,
677 + frag.size);
678 + curr_shm_offset += frag.size;
679 + }
680 +
681 + thread_ctx->set_thread_stat(src, ThreadSHMStat::DONE);
682 + });
683 +
684 + std::vector<torch::Tensor> tensor_list;
685 + tensor_list.reserve(metadata.tensor_num - 1);
686 + tensor_list.insert(tensor_list.begin(), tensor_list_with_metadata.begin() + 1,
687 + tensor_list_with_metadata.end());
688 +
689 + return tensor_list;
690 +}
691 +} // namespace
692 +
693 +void shm_gather(int64_t handle, torch::Tensor& data,
694 + const std::optional<std::vector<torch::Tensor>>& outputs,
695 + int64_t dst) {
696 + TORCH_CHECK(data.is_contiguous())
697 + VLLM_DISPATCH_FLOATING_TYPES(data.scalar_type(), "shm_gather_impl", [&] {
698 + CPU_KERNEL_GUARD_IN(shm_gather_impl)
699 +
700 + if (outputs.has_value()) {
701 + TORCH_CHECK_LE(outputs->size(), MAX_SHM_RANK_NUM);
702 + scalar_t* output_ptrs[MAX_SHM_RANK_NUM] = {nullptr};
703 + for (int i = 0; i < outputs->size(); ++i) {
704 + output_ptrs[i] = outputs->at(i).data_ptr<scalar_t>();
705 + }
706 + shm_gather_impl(SHMManager::get_singleton_instance(handle)->get_shm_ctx(),
707 + data.data_ptr<scalar_t>(), data.numel(), output_ptrs,
708 + dst);
709 + } else {
710 + shm_gather_impl(SHMManager::get_singleton_instance(handle)->get_shm_ctx(),
711 + data.data_ptr<scalar_t>(), data.numel(), (scalar_t**)(0),
712 + dst);
713 + }
714 +
715 + CPU_KERNEL_GUARD_OUT(shm_gather_impl)
716 + });
717 +}
718 +
719 +void shm_all_gather(int64_t handle, const torch::Tensor& data,
720 + torch::Tensor& output) {
721 + TORCH_CHECK(data.is_contiguous())
722 + TORCH_CHECK(output.is_contiguous())
723 +
724 + const int64_t input_elem_num = data.numel();
725 + const int64_t output_elem_num = output.numel();
726 + TORCH_CHECK_EQ(output_elem_num % input_elem_num, 0);
727 + const int world_size = output_elem_num / input_elem_num;
728 +
729 + VLLM_DISPATCH_FLOATING_TYPES(data.scalar_type(), "shm_all_gather_impl", [&] {
730 + CPU_KERNEL_GUARD_IN(shm_all_gather_impl)
731 + auto ctx = SHMManager::get_singleton_instance(handle)->get_shm_ctx();
732 + TORCH_CHECK_EQ(ctx->group_size, world_size);
733 +
734 + scalar_t* output_ptrs[MAX_SHM_RANK_NUM] = {nullptr};
735 + for (int i = 0; i < world_size; ++i) {
736 + output_ptrs[i] = output.data_ptr<scalar_t>() + i * input_elem_num;
737 + }
738 + shm_gather_impl(ctx, data.data_ptr<scalar_t>(), data.numel(), output_ptrs,
739 + ctx->rank);
740 + CPU_KERNEL_GUARD_OUT(shm_all_gather_impl)
741 + });
742 +}
743 +
744 +void shm_allreduce(int64_t handle, torch::Tensor& data) {
745 + TORCH_CHECK(data.is_contiguous())
746 + VLLM_DISPATCH_FLOATING_TYPES(data.scalar_type(), "shm_allreduce_sum", [&] {
747 + CPU_KERNEL_GUARD_IN(shm_allreduce_sum)
748 + shm_allreduce_sum(SHMManager::get_singleton_instance(handle)->get_shm_ctx(),
749 + data.data_ptr<scalar_t>(), data.numel());
750 + CPU_KERNEL_GUARD_OUT(shm_allreduce_sum)
751 + });
752 +}
753 +
754 +void shm_send_tensor_list(int64_t handle,
755 + const std::vector<torch::Tensor>& tensor_list,
756 + int64_t dst) {
757 + CPU_KERNEL_GUARD_IN(shm_send_tensor_list)
758 + shm_send_tensor_list_impl(
759 + SHMManager::get_singleton_instance(handle)->get_shm_ctx(), tensor_list);
760 + CPU_KERNEL_GUARD_OUT(shm_send_tensor_list)
761 +}
762 +
763 +std::vector<torch::Tensor> shm_recv_tensor_list(int64_t handle, int64_t src) {
764 + CPU_KERNEL_GUARD_IN(shm_recv_tensor_list)
765 + auto tensor_list = shm_recv_tensor_list_impl(
766 + SHMManager::get_singleton_instance(handle)->get_shm_ctx(), src);
767 + CPU_KERNEL_GUARD_OUT(shm_recv_tensor_list)
768 + return tensor_list;
769 +}
770 +
771 +int64_t init_shm_manager(const std::string& name, const int64_t group_size,
772 + const int64_t rank) {
773 + return SHMManager::create_singleton_instance(name, group_size, rank);
774 +}
775 +
776 +std::string join_shm_manager(int64_t handle, const std::string& name) {
777 + auto shm_manager = SHMManager::get_singleton_instance(handle);
778 + TORCH_CHECK(shm_manager);
779 + shm_manager->join(name);
780 + return shm_manager->get_shm_ctx()->to_string();
781 +}
\ No newline at end of file
csrc/cpu/torch_bindings.cpp
+43
@@ -22,6 +22,26 @@ void mla_decode_kvcache(torch::Tensor& out, torch::Tensor& query,
22 torch::Tensor& kv_cache, double scale,
23 torch::Tensor& block_tables, torch::Tensor& seq_lens);
24
25 +int64_t init_shm_manager(const std::string& name, const int64_t group_size,
26 + const int64_t rank);
27 +
28 +std::string join_shm_manager(int64_t handle, const std::string& name);
29 +
30 +void shm_allreduce(int64_t handle, torch::Tensor& data);
31 +
32 +void shm_gather(int64_t handle, torch::Tensor& data,
33 + const std::optional<std::vector<torch::Tensor>>& outputs,
34 + int64_t dst);
35 +
36 +void shm_all_gather(int64_t handle, const torch::Tensor& data,
37 + torch::Tensor& output);
38 +
39 +void shm_send_tensor_list(int64_t handle,
40 + const std::vector<torch::Tensor>& tensor_list,
41 + int64_t dst);
42 +
43 +std::vector<torch::Tensor> shm_recv_tensor_list(int64_t handle, int64_t src);
44 +
45 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
46 // vLLM custom ops
47
@@ -131,6 +151,29 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
151 " Tensor? azp, Tensor? bias) -> ()");
152 ops.impl("cutlass_scaled_mm_azp", torch::kCPU, &int8_scaled_mm_azp);
153 #endif
154 +
155 +// SHM CCL
156 +#ifdef __AVX512F__
157 + ops.def("init_shm_manager(str name, int group_size, int rank) -> int",
158 + &init_shm_manager);
159 + ops.def("join_shm_manager(int handle, str name) -> str", &join_shm_manager);
160 + ops.def("shm_allreduce(int handle, Tensor! data) -> ()");
161 + ops.impl("shm_allreduce", torch::kCPU, &shm_allreduce);
162 + ops.def(
163 + "shm_gather(int handle, Tensor data, Tensor[](a!)? outputs, int dst) -> "
164 + "()");
165 + ops.impl("shm_gather", torch::kCPU, &shm_gather);
166 + ops.def(
167 + "shm_all_gather(int handle, Tensor data, Tensor! output) -> "
168 + "()");
169 + ops.impl("shm_all_gather", torch::kCPU, &shm_all_gather);
170 + ops.def(
171 + "shm_send_tensor_list(int handle, Tensor[](a) tensor_list, int dst) -> "
172 + "()");
173 + ops.impl("shm_send_tensor_list", torch::kCPU, &shm_send_tensor_list);
174 + ops.def("shm_recv_tensor_list(int handle, int src) -> Tensor[](a)",
175 + &shm_recv_tensor_list);
176 +#endif
177 }
178
179 TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
csrc/cpu/utils.cpp
+1 -1
@@ -18,7 +18,7 @@ std::string init_cpu_threads_env(const std::string& cpu_ids) {
18
19 #ifndef VLLM_NUMA_DISABLED
20 std::string init_cpu_threads_env(const std::string& cpu_ids) {
21 - bitmask* omp_cpu_mask = numa_parse_cpustring(cpu_ids.c_str());
21 + bitmask* omp_cpu_mask = numa_parse_cpustring_all(cpu_ids.c_str());
22 TORCH_CHECK(omp_cpu_mask->size > 0);
23 std::vector<int> omp_cpu_ids;
24 omp_cpu_ids.reserve(omp_cpu_mask->size);
docs/source/getting_started/installation/cpu.md
+5 -3
@@ -272,12 +272,14 @@ $ python examples/offline_inference/basic/basic.py
272
273 - Decouple the HTTP serving components from the inference components. In a GPU backend configuration, the HTTP serving and tokenization tasks operate on the CPU, while inference runs on the GPU, which typically does not pose a problem. However, in a CPU-based setup, the HTTP serving and tokenization can cause significant context switching and reduced cache efficiency. Therefore, it is strongly recommended to segregate these two components for improved performance.
274
275 -- On CPU based setup with NUMA enabled, the memory access performance may be largely impacted by the [topology](https://github.com/intel/intel-extension-for-pytorch/blob/main/docs/tutorials/performance_tuning/tuning_guide.inc.md#non-uniform-memory-access-numa). For NUMA architecture, two optimizations are to recommended: Tensor Parallel or Data Parallel.
275 +- On CPU based setup with NUMA enabled, the memory access performance may be largely impacted by the [topology](https://github.com/intel/intel-extension-for-pytorch/blob/main/docs/tutorials/performance_tuning/tuning_guide.inc.md#non-uniform-memory-access-numa). For NUMA architecture, Tensor Parallel is a option for better performance.
276
277 - - Using Tensor Parallel for a latency constraints deployment: following GPU backend design, a Megatron-LM's parallel algorithm will be used to shard the model, based on the number of NUMA nodes (e.g. TP = 2 for a two NUMA node system). With [TP feature on CPU](gh-pr:6125) merged, Tensor Parallel is supported for serving and offline inferencing. In general each NUMA node is treated as one GPU card. Below is the example script to enable Tensor Parallel = 2 for serving:
277 + - Tensor Parallel is supported for serving and offline inferencing. In general each NUMA node is treated as one GPU card. Below is the example script to enable Tensor Parallel = 2 for serving:
278
279 ```console
280 VLLM_CPU_KVCACHE_SPACE=40 VLLM_CPU_OMP_THREADS_BIND="0-31|32-63" vllm serve meta-llama/Llama-2-7b-chat-hf -tp=2 --distributed-executor-backend mp
281 ```
282
283 - - Using Data Parallel for maximum throughput: to launch an LLM serving endpoint on each NUMA node along with one additional load balancer to dispatch the requests to those endpoints. Common solutions like [Nginx](#nginxloadbalancer) or HAProxy are recommended. Anyscale Ray project provides the feature on LLM [serving](https://docs.ray.io/en/latest/serve/index.html). Here is the example to setup a scalable LLM serving with [Ray Serve](https://github.com/intel/llm-on-ray/blob/main/docs/setup.inc.md).
283 + - For each thread id list in `VLLM_CPU_OMP_THREADS_BIND`, users should guarantee threads in the list belong to a same NUMA node.
284 +
285 + - Meanwhile, users should also take care of memory capacity of each NUMA node. The memory usage of each TP rank is the sum of `weight shard size` and `VLLM_CPU_KVCACHE_SPACE`, if it exceeds the capacity of a single NUMA node, TP worker will be killed due to out-of-memory.
vllm/distributed/device_communicators/cpu_communicator.py
+117 -12
@@ -1,10 +1,14 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 -from typing import Optional
3 +import os
4 +from typing import List, Optional
5
6 import torch
7 from torch.distributed import ProcessGroup
8
9 +from vllm.platforms import current_platform
10 +from vllm.platforms.interface import CpuArchEnum
11 +
12 from .base_device_communicator import DeviceCommunicatorBase
13
14
@@ -16,19 +20,120 @@ class CpuCommunicator(DeviceCommunicatorBase):
20 device_group: Optional[ProcessGroup] = None,
21 unique_name: str = ""):
22 super().__init__(cpu_group, device, device_group, unique_name)
19 - self.ipex_available = False
23 self.dist_module = torch.distributed
21 - try:
22 - import intel_extension_for_pytorch as ipex
23 - self.ipex_available = True
24 - self.dist_module = ipex.distributed
25 - except ImportError:
26 - """
27 - Intel IPEX not found. Falling back to PyTorch native
28 - all_reduce for CPU (e.g. MacOS)
29 - """
30 - pass
24 +
25 + if current_platform.get_cpu_architecture() == CpuArchEnum.X86:
26 + self.dist_module = _CPUSHMDistributed(self)
27
28 def all_reduce(self, input_):
29 self.dist_module.all_reduce(input_, group=self.device_group)
30 return input_
31 +
32 + def gather(self,
33 + input_: torch.Tensor,
34 + dst: int = 0,
35 + dim: int = -1) -> Optional[torch.Tensor]:
36 + """
37 + NOTE: We assume that the input tensor is on the same device across
38 + all the ranks.
39 + NOTE: `dst` is the local rank of the destination rank.
40 + """
41 + world_size = self.world_size
42 + assert -input_.dim() <= dim < input_.dim(), (
43 + f"Invalid dim ({dim}) for input tensor with shape {input_.size()}")
44 + if dim < 0:
45 + # Convert negative dim to positive.
46 + dim += input_.dim()
47 +
48 + # Allocate output tensor.
49 + if self.rank_in_group == dst:
50 + gather_list = [torch.empty_like(input_) for _ in range(world_size)]
51 + else:
52 + gather_list = None
53 +
54 + # Gather.
55 + self.dist_module.gather(input_,
56 + gather_list,
57 + dst=self.ranks[dst],
58 + group=self.device_group)
59 +
60 + if self.rank_in_group == dst:
61 + output_tensor = torch.cat(gather_list, dim=dim)
62 + else:
63 + output_tensor = None
64 + return output_tensor
65 +
66 + def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
67 + if dim < 0:
68 + # Convert negative dim to positive.
69 + dim += input_.dim()
70 + input_size = input_.size()
71 + # NOTE: we have to use concat-style all-gather here,
72 + # stack-style all-gather has compatibility issues with
73 + # torch.compile . see https://github.com/pytorch/pytorch/issues/138795
74 + output_size = (input_size[0] * self.world_size, ) + input_size[1:]
75 + # Allocate output tensor.
76 + output_tensor = torch.empty(output_size,
77 + dtype=input_.dtype,
78 + device=input_.device)
79 + # All-gather.
80 + self.dist_module.all_gather_into_tensor(output_tensor,
81 + input_,
82 + group=self.device_group)
83 +
84 + # Reshape
85 + output_tensor = output_tensor.reshape((self.world_size, ) + input_size)
86 + output_tensor = output_tensor.movedim(0, dim)
87 + output_tensor = output_tensor.reshape(input_size[:dim] +
88 + (self.world_size *
89 + input_size[dim], ) +
90 + input_size[dim + 1:])
91 + return output_tensor
92 +
93 +
94 +class _CPUSHMDistributed:
95 +
96 + def __init__(self, communicator: CpuCommunicator):
97 + instance_identifier = os.environ["VLLM_DIST_IDENT"]
98 + self.communicator = communicator
99 +
100 + group_ranks = [str(rank) for rank in self.communicator.ranks]
101 + shm_group_identifier = f"[{'-'.join(group_ranks)}]"
102 + self.group_name = f"{instance_identifier}-{shm_group_identifier}-cpushm"
103 +
104 + self.handle = self._init_cpu_shm()
105 +
106 + def _init_cpu_shm(self) -> int:
107 + handle = torch.ops._C.init_shm_manager(
108 + self.group_name,
109 + self.communicator.world_size,
110 + self.communicator.rank,
111 + )
112 + torch.distributed.barrier(self.communicator.device_group)
113 + torch.ops._C.join_shm_manager(
114 + handle,
115 + self.group_name,
116 + )
117 + torch.distributed.barrier(self.communicator.device_group)
118 +
119 + return handle
120 +
121 + def all_reduce(self,
122 + input: torch.Tensor,
123 + group: Optional[ProcessGroup] = None) -> None:
124 + torch.ops._C.shm_allreduce(self.handle, input)
125 +
126 + def gather(self,
127 + input: torch.Tensor,
128 + gather_list: Optional[List[torch.Tensor]],
129 + dst: int = -1,
130 + group: Optional[ProcessGroup] = None) -> None:
131 + # Note: different from the torch gather, here we use local dst rank.
132 + torch.ops._C.shm_gather(self.handle, input, gather_list,
133 + torch.distributed.get_group_rank(group, dst))
134 +
135 + def all_gather_into_tensor(self,
136 + output: torch.Tensor,
137 + input: torch.Tensor,
138 + group: Optional[ProcessGroup] = None) -> None:
139 + torch.ops._C.shm_all_gather(self.handle, input, output)
vllm/worker/cpu_worker.py
+7
@@ -1,5 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 """A CPU worker class."""
3 +import os
4 from typing import Dict, List, Optional, Set, Tuple, Type
5
6 import torch
@@ -139,6 +140,8 @@ class CPUWorker(LocalOrDistributedWorkerBase):
140
141 self.local_rank = local_rank
142 self.rank = rank
143 + vllm_config.parallel_config.rank = rank
144 +
145 self.distributed_init_method = distributed_init_method
146
147 self.is_driver_worker = is_driver_worker
@@ -217,6 +220,10 @@ class CPUWorker(LocalOrDistributedWorkerBase):
220 ret = torch.ops._C_utils.init_cpu_threads_env(self.local_omp_cpuid)
221 if ret:
222 logger.info(ret)
223 +
224 + # Note: unique identifier for creating allreduce shared memory
225 + os.environ["VLLM_DIST_IDENT"] = self.distributed_init_method.split(
226 + ":")[-1]
227 self.device = torch.device("cpu")
228 self.init_distributed_environment()
229 # Set random seed.