Fix: Build error seen on Power Architecture (vllm-project#10421)

Signed-off-by: Manjul Mohan <[email protected]> Signed-off-by: B-201 <[email protected]> Signed-off-by: Isotr0py <[email protected]> Signed-off-by: youkaichao <[email protected]> Signed-off-by: ismael-dm <[email protected]> Signed-off-by: Andrew Nesbitt <[email protected]> Signed-off-by: mgoin <[email protected]> Signed-off-by: yan ma <[email protected]> Signed-off-by: Angus Wang <[email protected]> Signed-off-by: Lucas Wilkinson <[email protected]> Signed-off-by: rickyx <[email protected]> Signed-off-by: Jee Jee Li <[email protected]> Signed-off-by: Mengqing Cao <[email protected]> Signed-off-by: Travis Johnson <[email protected]> Co-authored-by: Manjul Mohan [email protected] <[email protected]> Co-authored-by: B-201 <[email protected]> Co-authored-by: Isotr0py <[email protected]> Co-authored-by: youkaichao <[email protected]> Co-authored-by: ismael-dm <[email protected]> Co-authored-by: Andrew Nesbitt <[email protected]> Co-authored-by: Michael Goin <[email protected]> Co-authored-by: Yan Ma <[email protected]> Co-authored-by: Angus Wang <[email protected]> Co-authored-by: Lucas Wilkinson <[email protected]> Co-authored-by: Ricky Xu <[email protected]> Co-authored-by: Kevin H. Luu <[email protected]> Co-authored-by: Jee Jee Li <[email protected]> Co-authored-by: Mengqing Cao <[email protected]> Co-authored-by: Travis Johnson <[email protected]> Co-authored-by: Russell Bryant <[email protected]> Signed-off-by: rickyx <[email protected]>
rickyyx · Nov 20, 2024 · 61b0962 · 61b0962
1 parent 6472906
commit 61b0962
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 6 deletions.
diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake
@@ -16,10 +16,16 @@ include_directories("${CMAKE_SOURCE_DIR}/csrc")
 #
 # Check the compile flags
 #
-list(APPEND CXX_COMPILE_FLAGS
-    "-fopenmp"
-    "-mf16c"
-    "-DVLLM_CPU_EXTENSION")
+if (CMAKE_SYSTEM_PROCESSOR STREQUAL "ppc64le")
+    list(APPEND CXX_COMPILE_FLAGS
+        "-fopenmp"
+        "-DVLLM_CPU_EXTENSION")
+else()
+    list(APPEND CXX_COMPILE_FLAGS
+        "-fopenmp"
+        "-mf16c"
+        "-DVLLM_CPU_EXTENSION")
+endif()
 
 execute_process(COMMAND cat /proc/cpuinfo
                 RESULT_VARIABLE CPUINFO_RET

diff --git a/csrc/cpu/attention.cpp b/csrc/cpu/attention.cpp
@@ -24,12 +24,20 @@ struct KernelVecType<float> {
 
 template <>
 struct KernelVecType<c10::Half> {
+#ifdef __powerpc64__
+  // Power architecture-specific vector types
+  using q_load_vec_type = vec_op::FP32Vec8;
+  using k_load_vec_type = vec_op::FP32Vec16;
+  using v_load_vec_type = vec_op::FP32Vec16;
+#else
+  // Fallback for other architectures, including x86
   using q_load_vec_type = vec_op::FP16Vec8;
-  using q_vec_type = vec_op::FP32Vec16;
   using k_load_vec_type = vec_op::FP16Vec16;
+  using v_load_vec_type = vec_op::FP16Vec16;
+#endif
+  using q_vec_type = vec_op::FP32Vec16;
   using k_vec_type = vec_op::FP32Vec16;
   using qk_acc_vec_type = vec_op::FP32Vec16;
-  using v_load_vec_type = vec_op::FP16Vec16;
 };
 
 #ifdef __AVX512BF16__

diff --git a/csrc/cpu/quant.cpp b/csrc/cpu/quant.cpp
@@ -25,7 +25,13 @@ struct KernelVecType<c10::BFloat16> {
 
 template <>
 struct KernelVecType<c10::Half> {
+#ifdef __powerpc64__
+  // Power architecture-specific vector type
+  using load_vec_type = vec_op::FP32Vec16;
+#else
+  // Fallback for other architectures
   using load_vec_type = vec_op::FP16Vec16;
+#endif
   using azp_adj_load_vec_type = vec_op::INT32Vec16;
   using cvt_vec_type = vec_op::FP32Vec16;
 };