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 2, 06:00 UTC
[HARDWARE][POWER] Enable SHM communicator support for PowerPC (#43754)
Signed-off-by: Rukhaiya <rukhaiya@c643n08aix1-lp1.pok.stglabs.ibm.com> Signed-off-by: Rukhaiya <bibirukhaiya123@gmail.com> Co-authored-by: Rukhaiya <rukhaiya@c643n08aix1-lp1.pok.stglabs.ibm.com> Co-authored-by: Akash kaothalkar <61960177+Akashcodes732@users.noreply.github.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
Rukhaiya2004 committed
Jun 2, 2026 at 15:36 UTC
689b0eeb9e014648f7226e1fb1adbadbd481f27c
6 files changed
+125
-10
=4.5.1
cmake/cpu_extension.cmake
+6
@@ -423,6 +423,12 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
423
${VLLM_EXT_SRC})
424
endif()
425
426
+if (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
427
+ set(VLLM_EXT_SRC
428
+ "csrc/cpu/shm.cpp"
429
+ ${VLLM_EXT_SRC})
430
+endif()
431
+
432
if(USE_ONEDNN)
433
set(VLLM_EXT_SRC
434
"csrc/cpu/dnnl_kernels.cpp"
csrc/cpu/cpu_types_vsx.hpp
+106
-1
@@ -89,6 +89,35 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
89
}
90
};
91
92
+struct FP16Vec16 : public Vec<FP16Vec16> {
93
+ constexpr static int VEC_ELEM_NUM = 16;
94
+ ss16x8x2_t reg;
95
+
96
+ explicit FP16Vec16(const void* ptr) {
97
+ reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)ptr);
98
+ reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)ptr);
99
+ }
100
+
101
+ explicit FP16Vec16(bool, const void* ptr) : FP16Vec16(ptr) {}
102
+
103
+ explicit FP16Vec16(const FP32Vec16&);
104
+
105
+ void save(void* ptr) const {
106
+ vec_xst(reg.val[0], 0, (signed short*)ptr);
107
+ vec_xst(reg.val[1], 16, (signed short*)ptr);
108
+ }
109
+
110
+ void save(void* ptr, int elem_num) const {
111
+ int num = std::max(0, std::min(elem_num, VEC_ELEM_NUM));
112
+ if (num <= 8) {
113
+ vec_xst_len(reg.val[0], (signed short*)ptr, num * 2);
114
+ } else {
115
+ vec_xst(reg.val[0], 0, (signed short*)ptr);
116
+ vec_xst_len(reg.val[1], (signed short*)ptr + 8, (num - 8) * 2);
117
+ }
118
+ }
119
+};
120
+
121
struct BF16Vec16 : public Vec<BF16Vec16> {
122
constexpr static int VEC_ELEM_NUM = 16;
123
@@ -100,6 +129,8 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
129
reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)ptr);
130
}
131
132
+ explicit BF16Vec16(bool, const void* ptr) : BF16Vec16(ptr) {}
133
+
134
explicit BF16Vec16(const FP32Vec16&);
135
136
void save(void* ptr) const {
@@ -379,6 +410,8 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
410
reg.val[3] = vec_xl(48, ptr);
411
}
412
413
+ explicit FP32Vec16(bool, const float* ptr) : FP32Vec16(ptr) {}
414
+
415
explicit FP32Vec16(f32x4x4_t data) : reg(data) {}
416
417
explicit FP32Vec16(const FP32Vec16& data) {
@@ -402,6 +435,7 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
435
reg.val[3] = data.reg.val[1];
436
}
437
438
+ explicit FP32Vec16(const FP16Vec16& v);
439
explicit FP32Vec16(const BF16Vec16& v) {
440
reg.val[0] = (__vector float)vec_mergeh(zero, v.reg.val[0]);
441
reg.val[1] = (__vector float)vec_mergel(zero, v.reg.val[0]);
@@ -735,6 +769,40 @@ inline BF16Vec8::BF16Vec8(const FP32Vec8& v) {
769
#endif
770
}
771
772
+inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
773
+ alignas(16) float temp_fp32[16];
774
+ alignas(16) c10::Half temp_fp16[16];
775
+
776
+ vec_xst(v.reg.val[0], 0, temp_fp32);
777
+ vec_xst(v.reg.val[1], 16, temp_fp32);
778
+ vec_xst(v.reg.val[2], 32, temp_fp32);
779
+ vec_xst(v.reg.val[3], 48, temp_fp32);
780
+
781
+ for (int i = 0; i < 16; i++) {
782
+ temp_fp16[i] = c10::Half(temp_fp32[i]);
783
+ }
784
+
785
+ reg.val[0] = (__vector signed short)vec_xl(0, (signed short*)temp_fp16);
786
+ reg.val[1] = (__vector signed short)vec_xl(16, (signed short*)temp_fp16);
787
+}
788
+
789
+inline FP32Vec16::FP32Vec16(const FP16Vec16& v) {
790
+ alignas(16) c10::Half temp_fp16[16];
791
+ alignas(16) float temp_fp32[16];
792
+
793
+ vec_xst(v.reg.val[0], 0, (signed short*)temp_fp16);
794
+ vec_xst(v.reg.val[1], 16, (signed short*)temp_fp16);
795
+
796
+ for (int i = 0; i < 16; i++) {
797
+ temp_fp32[i] = float(temp_fp16[i]);
798
+ }
799
+
800
+ reg.val[0] = vec_xl(0, temp_fp32);
801
+ reg.val[1] = vec_xl(16, temp_fp32);
802
+ reg.val[2] = vec_xl(32, temp_fp32);
803
+ reg.val[3] = vec_xl(48, temp_fp32);
804
+}
805
+
806
inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
807
#ifdef _ARCH_PWR10
808
__vector signed short ret[4];
@@ -794,6 +862,43 @@ inline void prefetch(const void* addr) {
862
__asm__ __volatile__("dcbt 0, %0" : : "r"(addr) : "memory");
863
}
864
797
-}; // namespace vec_op
865
+struct INT8Vec64 {
866
+ __vector signed char data[4];
867
+
868
+ INT8Vec64() = default;
869
+
870
+ explicit INT8Vec64(const int8_t* ptr) {
871
+ data[0] = vec_xl(0, ptr);
872
+ data[1] = vec_xl(16, ptr);
873
+ data[2] = vec_xl(32, ptr);
874
+ data[3] = vec_xl(48, ptr);
875
+ }
876
+
877
+ explicit INT8Vec64(bool, const int8_t* ptr) : INT8Vec64(ptr) {}
878
+
879
+ void save(int8_t* ptr) const {
880
+ vec_xst(data[0], 0, ptr);
881
+ vec_xst(data[1], 16, ptr);
882
+ vec_xst(data[2], 32, ptr);
883
+ vec_xst(data[3], 48, ptr);
884
+ }
885
+
886
+ void save(int8_t* ptr, int elem_num) const {
887
+ if (elem_num <= 0) return;
888
+
889
+ int full_vecs = elem_num / 16;
890
+ for (int i = 0; i < full_vecs && i < 4; i++) {
891
+ vec_xst(data[i], i * 16, ptr);
892
+ }
893
+
894
+ int remaining = elem_num % 16;
895
+ if (remaining > 0 && full_vecs < 4) {
896
+ vec_xst_len(data[full_vecs], ptr + full_vecs * 16, remaining);
897
+ }
898
+ }
899
+
900
+ void nt_save(int8_t* ptr) const { save(ptr); }
901
+};
902
+} // namespace vec_op
903
904
#endif
csrc/cpu/shm.cpp
+10
-8
@@ -5,7 +5,7 @@
5
#include <sys/stat.h>
6
#include <unistd.h>
7
8
-#ifdef __aarch64__
8
+#if defined(__aarch64__) || defined(__powerpc64__)
9
#include <atomic>
10
#endif
11
@@ -38,7 +38,7 @@ struct KernelVecType<c10::Half> {
38
};
39
40
struct ThreadSHMContext {
41
-#ifdef __aarch64__
41
+#if defined(__aarch64__) || defined(__powerpc64__)
42
// memory model is weaker on AArch64, so we use atomic variables for
43
// consumer (load-acquire) and producer (store-release) to make sure
44
// that a stamp cannot be ready before the corresponding data is ready.
@@ -75,7 +75,7 @@ struct ThreadSHMContext {
75
TORCH_CHECK(group_size <= MAX_SHM_RANK_NUM);
76
TORCH_CHECK((size_t)this % 64 == 0);
77
TORCH_CHECK((size_t)thread_shm_ptr % 64 == 0);
78
-#ifdef __aarch64__
78
+#if defined(__aarch64__) || defined(__powerpc64__)
79
_curr_thread_stamp[0].store(1, std::memory_order_relaxed);
80
_curr_thread_stamp[1].store(1, std::memory_order_relaxed);
81
_ready_thread_stamp[0].store(0, std::memory_order_relaxed);
@@ -124,7 +124,7 @@ struct ThreadSHMContext {
124
}
125
126
char get_curr_stamp(int idx) const {
127
-#ifdef __aarch64__
127
+#if defined(__aarch64__) || defined(__powerpc64__)
128
return _curr_thread_stamp[idx].load(std::memory_order_acquire);
129
#else
130
return _curr_thread_stamp[idx];
@@ -132,7 +132,7 @@ struct ThreadSHMContext {
132
}
133
134
char get_ready_stamp(int idx) const {
135
-#ifdef __aarch64__
135
+#if defined(__aarch64__) || defined(__powerpc64__)
136
return _ready_thread_stamp[idx].load(std::memory_order_acquire);
137
#else
138
return _ready_thread_stamp[idx];
@@ -140,7 +140,7 @@ struct ThreadSHMContext {
140
}
141
142
void next_stamp() {
143
-#ifdef __aarch64__
143
+#if defined(__aarch64__) || defined(__powerpc64__)
144
_curr_thread_stamp[local_stamp_buffer_idx].fetch_add(
145
1, std::memory_order_release);
146
#else
@@ -150,7 +150,7 @@ struct ThreadSHMContext {
150
}
151
152
void commit_ready_stamp() {
153
-#ifdef __aarch64__
153
+#if defined(__aarch64__) || defined(__powerpc64__)
154
_ready_thread_stamp[local_stamp_buffer_idx].store(
155
_curr_thread_stamp[local_stamp_buffer_idx].load(
156
std::memory_order_relaxed),
@@ -186,8 +186,10 @@ struct ThreadSHMContext {
186
break;
187
}
188
++_spinning_count;
189
-#ifdef __aarch64__
189
+#if defined(__aarch64__)
190
__asm__ __volatile__("yield");
191
+#elif defined(__powerpc64__)
192
+ __asm__ __volatile__("or 1,1,1");
193
#else
194
_mm_pause();
195
#endif // __aarch64__
csrc/cpu/torch_bindings.cpp
+2
-1
@@ -378,7 +378,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
378
#endif
379
380
// SHM CCL
381
-#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__))
381
+#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__)) || \
382
+ defined(__powerpc64__)
383
ops.def(
384
"init_shm_manager(str name, int group_size, int rank, int thread_num) -> "
385
"int",
vllm/distributed/device_communicators/cpu_communicator.py
+1
@@ -32,6 +32,7 @@ class CpuCommunicator(DeviceCommunicatorBase):
32
(
33
current_platform.get_cpu_architecture() == CpuArchEnum.X86
34
or current_platform.get_cpu_architecture() == CpuArchEnum.ARM
35
+ or current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC
36
)
37
and hasattr(torch.ops._C, "init_shm_manager")
38
and (unique_name.startswith("tp") or unique_name.startswith("pp"))