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, 18:15 UTC
[CPU Backend] [Perf] Accelerate tensor-parallel/data-parallel inference across NUMA domains on Arm (#32792)
Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com>
Fadi Arafeh committed
Jan 22, 2026 at 18:55 UTC
744ef30484169694ef4aaba9420a03ed5b861b66
6 files changed
+164
-6
cmake/cpu_extension.cmake
+6
@@ -379,6 +379,12 @@ if (AVX512_FOUND AND NOT AVX512_DISABLED)
379
endif()
380
endif()
381
382
+if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
383
+ set(VLLM_EXT_SRC
384
+ "csrc/cpu/shm.cpp"
385
+ ${VLLM_EXT_SRC})
386
+endif()
387
+
388
if(USE_ONEDNN)
389
set(VLLM_EXT_SRC
390
"csrc/cpu/dnnl_kernels.cpp"
csrc/cpu/cpu_types_arm.hpp
+99
-1
@@ -80,8 +80,10 @@ struct FP16Vec16 : public Vec<FP16Vec16> {
80
reg.val[1] = vld1q_f16(reinterpret_cast<const __fp16*>(ptr) + 8);
81
}
82
83
- explicit FP16Vec16(const FP32Vec16& vec);
83
+ // ASIMD does not support non-temporal loads
84
+ explicit FP16Vec16(bool, const void* ptr) : FP16Vec16(ptr) {}
85
86
+ explicit FP16Vec16(const FP32Vec16& vec);
87
void save(void* ptr) const {
88
vst1q_f16(reinterpret_cast<__fp16*>(ptr), reg.val[0]);
89
vst1q_f16(reinterpret_cast<__fp16*>(ptr) + 8, reg.val[1]);
@@ -190,6 +192,9 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
192
explicit BF16Vec16(const void* ptr)
193
: reg(*reinterpret_cast<const bfloat16x8x2_t*>(ptr)) {};
194
195
+ // ASIMD does not support non-temporal loads
196
+ explicit BF16Vec16(bool, const void* ptr) : BF16Vec16(ptr) {}
197
+
198
explicit BF16Vec16(bfloat16x8x2_t data) : reg(data) {};
199
200
explicit BF16Vec16(const FP32Vec16&);
@@ -474,6 +479,9 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
479
: reg({vld1q_f32(ptr), vld1q_f32(ptr + 4), vld1q_f32(ptr + 8),
480
vld1q_f32(ptr + 12)}) {}
481
482
+ // ASIMD does not support non-temporal loads
483
+ explicit FP32Vec16(bool, const float* ptr) : FP32Vec16(ptr) {}
484
+
485
explicit FP32Vec16(float32x4x4_t data) : reg(data) {}
486
487
explicit FP32Vec16(const FP32Vec8& data) {
@@ -756,6 +764,96 @@ struct INT8Vec16 : public Vec<INT8Vec16> {
764
};
765
};
766
767
+struct INT8Vec64 : public Vec<INT8Vec64> {
768
+ constexpr static int VEC_ELEM_NUM = 64;
769
+ union AliasReg {
770
+ int8x16x4_t reg;
771
+ int8_t values[VEC_ELEM_NUM];
772
+ };
773
+ int8x16x4_t reg;
774
+
775
+ explicit INT8Vec64(const int8_t* ptr) { reg = vld1q_s8_x4(ptr); }
776
+
777
+ // ASIMD does not support non-temporal loads
778
+ explicit INT8Vec64(bool, const int8_t* ptr) : INT8Vec64(ptr) {}
779
+
780
+ void save(int8_t* ptr) const { vst1q_s8_x4(ptr, reg); }
781
+
782
+ // masked store
783
+ void save(int8_t* p, int elem_num) const {
784
+ TORCH_CHECK(elem_num <= VEC_ELEM_NUM && elem_num > 0);
785
+
786
+ if (elem_num == VEC_ELEM_NUM) {
787
+ vst1q_s8_x4(p, reg);
788
+ return;
789
+ }
790
+
791
+ const int full_quadwords = elem_num / 16;
792
+ const int remaining_bytes = elem_num % 16;
793
+
794
+ for (int i = 0; i < full_quadwords; ++i) {
795
+ vst1q_s8(p + 16 * i, reg.val[i]);
796
+ }
797
+
798
+ if (remaining_bytes) {
799
+ const int8x16_t v = reg.val[full_quadwords];
800
+ int8_t* tail = p + 16 * full_quadwords;
801
+ switch (remaining_bytes) {
802
+ case 15:
803
+ tail[14] = vgetq_lane_s8(v, 14);
804
+ [[fallthrough]];
805
+ case 14:
806
+ tail[13] = vgetq_lane_s8(v, 13);
807
+ [[fallthrough]];
808
+ case 13:
809
+ tail[12] = vgetq_lane_s8(v, 12);
810
+ [[fallthrough]];
811
+ case 12:
812
+ tail[11] = vgetq_lane_s8(v, 11);
813
+ [[fallthrough]];
814
+ case 11:
815
+ tail[10] = vgetq_lane_s8(v, 10);
816
+ [[fallthrough]];
817
+ case 10:
818
+ tail[9] = vgetq_lane_s8(v, 9);
819
+ [[fallthrough]];
820
+ case 9:
821
+ tail[8] = vgetq_lane_s8(v, 8);
822
+ [[fallthrough]];
823
+ case 8:
824
+ tail[7] = vgetq_lane_s8(v, 7);
825
+ [[fallthrough]];
826
+ case 7:
827
+ tail[6] = vgetq_lane_s8(v, 6);
828
+ [[fallthrough]];
829
+ case 6:
830
+ tail[5] = vgetq_lane_s8(v, 5);
831
+ [[fallthrough]];
832
+ case 5:
833
+ tail[4] = vgetq_lane_s8(v, 4);
834
+ [[fallthrough]];
835
+ case 4:
836
+ tail[3] = vgetq_lane_s8(v, 3);
837
+ [[fallthrough]];
838
+ case 3:
839
+ tail[2] = vgetq_lane_s8(v, 2);
840
+ [[fallthrough]];
841
+ case 2:
842
+ tail[1] = vgetq_lane_s8(v, 1);
843
+ [[fallthrough]];
844
+ case 1:
845
+ tail[0] = vgetq_lane_s8(v, 0);
846
+ break;
847
+ default:
848
+ break;
849
+ }
850
+ }
851
+ }
852
+
853
+ // ASIMD does not support non-temporal stores
854
+ void nt_save(int8_t* ptr) const { save(ptr); }
855
+}; // INT8Vec64
856
+
857
template <typename T>
858
struct VecType {
859
using vec_type = void;
csrc/cpu/shm.cpp
+50
-2
@@ -5,6 +5,10 @@
5
#include <sys/stat.h>
6
#include <unistd.h>
7
8
+#ifdef __aarch64__
9
+ #include <atomic>
10
+#endif
11
+
12
namespace {
13
#define MAX_SHM_RANK_NUM 8
14
#define PER_THREAD_SHM_BUFFER_BYTES (4 * 1024 * 1024)
@@ -34,8 +38,17 @@ struct KernelVecType<c10::Half> {
38
};
39
40
struct ThreadSHMContext {
41
+#ifdef __aarch64__
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.
45
+ std::atomic<char> _curr_thread_stamp[2];
46
+ std::atomic<char> _ready_thread_stamp[2];
47
+ static_assert(std::atomic<char>::is_always_lock_free);
48
+#else
49
volatile char _curr_thread_stamp[2];
50
volatile char _ready_thread_stamp[2];
51
+#endif // __aarch64__
52
int local_stamp_buffer_idx;
53
int remote_stamp_buffer_idx;
54
int thread_id;
@@ -62,10 +75,17 @@ 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__
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);
82
+ _ready_thread_stamp[1].store(0, std::memory_order_relaxed);
83
+#else
84
_curr_thread_stamp[0] = 1;
85
_curr_thread_stamp[1] = 1;
86
_ready_thread_stamp[0] = 0;
87
_ready_thread_stamp[1] = 0;
88
+#endif // __aarch64__
89
_thread_buffer_mask[0] = 0;
90
_thread_buffer_mask[1] = 0;
91
for (int i = 0; i < MAX_SHM_RANK_NUM; ++i) {
@@ -103,19 +123,43 @@ struct ThreadSHMContext {
123
_thread_buffer_mask[local_stamp_buffer_idx] ^= 0xFFFFFFFFFFFFFFFF;
124
}
125
106
- char get_curr_stamp(int idx) const { return _curr_thread_stamp[idx]; }
126
+ char get_curr_stamp(int idx) const {
127
+#ifdef __aarch64__
128
+ return _curr_thread_stamp[idx].load(std::memory_order_acquire);
129
+#else
130
+ return _curr_thread_stamp[idx];
131
+#endif // __aarch64__
132
+ }
133
108
- char get_ready_stamp(int idx) const { return _ready_thread_stamp[idx]; }
134
+ char get_ready_stamp(int idx) const {
135
+#ifdef __aarch64__
136
+ return _ready_thread_stamp[idx].load(std::memory_order_acquire);
137
+#else
138
+ return _ready_thread_stamp[idx];
139
+#endif // __aarch64__
140
+ }
141
142
void next_stamp() {
143
+#ifdef __aarch64__
144
+ _curr_thread_stamp[local_stamp_buffer_idx].fetch_add(
145
+ 1, std::memory_order_release);
146
+#else
147
_mm_mfence();
148
_curr_thread_stamp[local_stamp_buffer_idx] += 1;
149
+#endif // __aarch64__
150
}
151
152
void commit_ready_stamp() {
153
+#ifdef __aarch64__
154
+ _ready_thread_stamp[local_stamp_buffer_idx].store(
155
+ _curr_thread_stamp[local_stamp_buffer_idx].load(
156
+ std::memory_order_relaxed),
157
+ std::memory_order_release);
158
+#else
159
_mm_mfence();
160
_ready_thread_stamp[local_stamp_buffer_idx] =
161
_curr_thread_stamp[local_stamp_buffer_idx];
162
+#endif // __aarch64__
163
}
164
165
int get_swizzled_rank(int idx) { return swizzled_ranks[idx]; }
@@ -142,7 +186,11 @@ struct ThreadSHMContext {
186
break;
187
}
188
++_spinning_count;
189
+#ifdef __aarch64__
190
+ __asm__ __volatile__("yield");
191
+#else
192
_mm_pause();
193
+#endif // __aarch64__
194
}
195
}
196
csrc/cpu/torch_bindings.cpp
+2
-2
@@ -230,7 +230,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
230
#endif
231
232
// SHM CCL
233
-#ifdef __AVX512F__
233
+#if defined(__AVX512F__) || defined(__aarch64__)
234
ops.def("init_shm_manager(str name, int group_size, int rank) -> int",
235
&init_shm_manager);
236
ops.def("join_shm_manager(int handle, str name) -> str", &join_shm_manager);
@@ -250,7 +250,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
250
ops.impl("shm_send_tensor_list", torch::kCPU, &shm_send_tensor_list);
251
ops.def("shm_recv_tensor_list(int handle, int src) -> Tensor[](a)",
252
&shm_recv_tensor_list);
253
-#endif
253
+#endif // #if defined(__AVX512F__) || defined(__aarch64__)
254
255
// sgl-kernels
256
#if defined(__AVX512BF16__) && defined(__AVX512F__) && defined(__AVX512VNNI__)
vllm/distributed/device_communicators/cpu_communicator.py
+4
-1
@@ -29,7 +29,10 @@ class CpuCommunicator(DeviceCommunicatorBase):
29
self.dist_module = torch.distributed
30
31
if (
32
- (current_platform.get_cpu_architecture() == CpuArchEnum.X86)
32
+ (
33
+ current_platform.get_cpu_architecture() == CpuArchEnum.X86
34
+ or current_platform.get_cpu_architecture() == CpuArchEnum.ARM
35
+ )
36
and hasattr(torch.ops._C, "init_shm_manager")
37
and (unique_name.startswith("tp") or unique_name.startswith("pp"))
38
):
vllm/v1/worker/cpu_worker.py
+3
@@ -66,6 +66,9 @@ class CPUWorker(Worker):
66
self.local_omp_cpuid = self._get_autobind_cpu_ids(
67
lambda cpus: cpus[-1:]
68
)
69
+ elif cpu_arch == CpuArchEnum.ARM:
70
+ # For AArch64, no SMT
71
+ self.local_omp_cpuid = self._get_autobind_cpu_ids(lambda cpus: cpus)
72
else:
73
self.local_omp_cpuid = "nobind"
74
elif omp_cpuids == "nobind":