@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c7753a9809
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, 02:15 UTC

[Hardware][CPU] Vllm int8 quantization enablement for ARM CPU (#14129)

Signed-off-by: nishith-fujitsu <nishith.jaiswal@fujitsu.com>

nishith-fujitsu committed Jul 10, 2025 at 21:29 UTC c7753a9809344cd8eefbda6d472cb9ab348f2274
5 files changed +347 -30
cmake/cpu_extension.cmake
+24 -4
@@ -165,17 +165,32 @@ else()
165 endif()
166
167 #
168 -# Build oneDNN for W8A8 GEMM kernels (only for x86-AVX512 platforms)
169 -#
170 -if (AVX512_FOUND AND NOT AVX512_DISABLED)
168 +# Build oneDNN for W8A8 GEMM kernels (only for x86-AVX512 /ARM platforms)
169 +# Flag to enable ACL kernels for AARCH64 platforms
170 +if ( VLLM_BUILD_ACL STREQUAL "ON")
171 + set(USE_ACL ON)
172 +else()
173 + set(USE_ACL OFF)
174 +endif()
175 +
176 +if ((AVX512_FOUND AND NOT AVX512_DISABLED) OR ASIMD_FOUND)
177 FetchContent_Declare(
178 oneDNN
179 GIT_REPOSITORY https://github.com/oneapi-src/oneDNN.git
174 - GIT_TAG v3.7.1
180 + GIT_TAG v3.8.1
181 GIT_PROGRESS TRUE
182 GIT_SHALLOW TRUE
183 )
184
185 + if(USE_ACL)
186 + find_library(ARM_COMPUTE_LIBRARY NAMES arm_compute PATHS $ENV{ACL_ROOT_DIR}/build/)
187 + if(NOT ARM_COMPUTE_LIBRARY)
188 + message(FATAL_ERROR "Could not find ARM Compute Library: please set ACL_ROOT_DIR")
189 + endif()
190 + set(ONEDNN_AARCH64_USE_ACL "ON")
191 + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wl,-rpath,$ENV{ACL_ROOT_DIR}/build/")
192 + endif()
193 +
194 set(ONEDNN_LIBRARY_TYPE "STATIC")
195 set(ONEDNN_BUILD_DOC "OFF")
196 set(ONEDNN_BUILD_EXAMPLES "OFF")
@@ -264,6 +279,11 @@ elseif(POWER10_FOUND)
279 "csrc/cpu/quant.cpp"
280 ${VLLM_EXT_SRC})
281 endif()
282 +if (ASIMD_FOUND)
283 + set(VLLM_EXT_SRC
284 + "csrc/cpu/quant.cpp"
285 + ${VLLM_EXT_SRC})
286 +endif()
287
288 message(STATUS "CPU extension source files: ${VLLM_EXT_SRC}")
289
csrc/cpu/cpu_types_arm.hpp
+264 -3
@@ -33,6 +33,8 @@ namespace vec_op {
33 #endif
34
35 #define FORCE_INLINE __attribute__((always_inline)) inline
36 +// Number of elements in single ASIMD vector of given Datatype
37 +#define NUM_ELEMENTS_REG(vec) (sizeof(vec) / sizeof(vec[0]))
38
39 namespace {
40 template <typename T, T... indexes, typename F>
@@ -86,8 +88,8 @@ struct FP16Vec16 : public Vec<FP16Vec16> {
88 }
89
90 void save(void* ptr, const int elem_num) const {
89 - int full_blocks = elem_num / 8;
90 - int remainder = elem_num % 8;
91 + int full_blocks = elem_num / NUM_ELEMENTS_REG(reg.val[0]);
92 + int remainder = elem_num % NUM_ELEMENTS_REG(reg.val[0]);
93
94 if (full_blocks > 0) {
95 vst1q_f16(reinterpret_cast<__fp16*>(ptr), reg.val[0]);
@@ -197,6 +199,25 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
199 vcvtq_high_bf16_f32(vcvtq_low_bf16_f32(v.val[2]), v.val[3])}) {};
200
201 void save(void* ptr) const { *reinterpret_cast<bfloat16x8x2_t*>(ptr) = reg; };
202 + void save(void* ptr, const int elem_num) const {
203 + int full_blocks = elem_num / NUM_ELEMENTS_REG(reg.val[0]);
204 + int remainder = elem_num % NUM_ELEMENTS_REG(reg.val[0]);
205 + for (int i = 0; i < full_blocks; i++)
206 + vst1q_bf16(
207 + reinterpret_cast<__bf16*>(ptr) + NUM_ELEMENTS_REG(reg.val[0]) * i,
208 + reg.val[i]);
209 + if (remainder > 0) {
210 + bfloat16x8_t temp = reg.val[full_blocks];
211 + bfloat16_t* base = reinterpret_cast<bfloat16_t*>(ptr) + full_blocks * 8;
212 + if (remainder > 0) base[0] = vgetq_lane_bf16(temp, 0);
213 + if (remainder > 1) base[1] = vgetq_lane_bf16(temp, 1);
214 + if (remainder > 2) base[2] = vgetq_lane_bf16(temp, 2);
215 + if (remainder > 3) base[3] = vgetq_lane_bf16(temp, 3);
216 + if (remainder > 4) base[4] = vgetq_lane_bf16(temp, 4);
217 + if (remainder > 5) base[5] = vgetq_lane_bf16(temp, 5);
218 + if (remainder > 6) base[6] = vgetq_lane_bf16(temp, 6);
219 + }
220 + };
221 };
222
223 struct BF16Vec32 : public Vec<BF16Vec32> {
@@ -213,6 +234,25 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
234 : reg({vec8_data.reg, vec8_data.reg, vec8_data.reg, vec8_data.reg}) {};
235
236 void save(void* ptr) const { *reinterpret_cast<bfloat16x8x4_t*>(ptr) = reg; };
237 + void save(void* ptr, const int elem_num) const {
238 + int full_blocks = elem_num / NUM_ELEMENTS_REG(reg.val[0]);
239 + int remainder = elem_num % NUM_ELEMENTS_REG(reg.val[0]);
240 + for (int i = 0; i < full_blocks; i++)
241 + vst1q_bf16(
242 + reinterpret_cast<__bf16*>(ptr) + NUM_ELEMENTS_REG(reg.val[0]) * i,
243 + reg.val[i]);
244 + if (remainder > 0) {
245 + bfloat16x8_t temp = reg.val[full_blocks];
246 + bfloat16_t* base = reinterpret_cast<bfloat16_t*>(ptr) + full_blocks * 8;
247 + base[0] = vgetq_lane_bf16(temp, 0);
248 + if (remainder > 1) base[1] = vgetq_lane_bf16(temp, 1);
249 + if (remainder > 2) base[2] = vgetq_lane_bf16(temp, 2);
250 + if (remainder > 3) base[3] = vgetq_lane_bf16(temp, 3);
251 + if (remainder > 4) base[4] = vgetq_lane_bf16(temp, 4);
252 + if (remainder > 5) base[5] = vgetq_lane_bf16(temp, 5);
253 + if (remainder > 6) base[6] = vgetq_lane_bf16(temp, 6);
254 + }
255 + };
256 };
257 #endif
258
@@ -372,6 +412,48 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
412 }
413 };
414
415 +struct INT32Vec16 : public Vec<INT32Vec16> {
416 + constexpr static int VEC_ELEM_NUM = 16;
417 + union AliasReg {
418 + int32x4x4_t reg;
419 + int32_t values[VEC_ELEM_NUM];
420 + };
421 + int32x4x4_t reg;
422 +
423 + explicit INT32Vec16(const void* ptr) {
424 + reg.val[0] = vld1q_s32(reinterpret_cast<const int32_t*>(ptr));
425 + reg.val[1] = vld1q_s32(reinterpret_cast<const int32_t*>(ptr) + 4);
426 + reg.val[2] = vld1q_s32(reinterpret_cast<const int32_t*>(ptr) + 8);
427 + reg.val[3] = vld1q_s32(reinterpret_cast<const int32_t*>(ptr) + 12);
428 + }
429 +
430 + void save(int32_t* ptr) const {
431 + vst1q_s32(ptr, reg.val[0]);
432 + vst1q_s32(ptr + 4, reg.val[1]);
433 + vst1q_s32(ptr + 8, reg.val[2]);
434 + vst1q_s32(ptr + 12, reg.val[3]);
435 + };
436 +
437 + void save(int32_t* ptr, const int elem_num) const {
438 + int full_blocks = elem_num / NUM_ELEMENTS_REG(reg.val[0]);
439 + int remainder = elem_num % NUM_ELEMENTS_REG(reg.val[0]);
440 +
441 + for (int i = 0; i < full_blocks; i++)
442 + vst1q_s32(
443 + reinterpret_cast<__int32_t*>(ptr) + NUM_ELEMENTS_REG(reg.val[0]) * i,
444 + reg.val[i]);
445 +
446 + if (remainder > 0) {
447 + int32x4_t temp = reg.val[full_blocks];
448 + int32_t* base = reinterpret_cast<int32_t*>(ptr) + full_blocks * 4;
449 + if (remainder > 0) base[0] = vgetq_lane_s32(temp, 0);
450 + if (remainder > 1) base[1] = vgetq_lane_s32(temp, 1);
451 + if (remainder > 2) base[2] = vgetq_lane_s32(temp, 2);
452 + if (remainder > 3) base[3] = vgetq_lane_s32(temp, 3);
453 + }
454 + }
455 +};
456 +
457 struct FP32Vec16 : public Vec<FP32Vec16> {
458 constexpr static int VEC_ELEM_NUM = 16;
459 union AliasReg {
@@ -434,7 +516,12 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
516 reg.val[2] = vcvt_f32_f16(vget_low_f16(v.reg.val[1]));
517 reg.val[3] = vcvt_f32_f16(vget_high_f16(v.reg.val[1]));
518 };
437 -
519 + explicit FP32Vec16(const INT32Vec16& v) {
520 + reg.val[0] = vcvtq_f32_s32(v.reg.val[0]);
521 + reg.val[1] = vcvtq_f32_s32(v.reg.val[1]);
522 + reg.val[2] = vcvtq_f32_s32(v.reg.val[2]);
523 + reg.val[3] = vcvtq_f32_s32(v.reg.val[3]);
524 + };
525 FP32Vec16 operator+(const FP32Vec16& b) const {
526 return FP32Vec16(float32x4x4_t({vaddq_f32(reg.val[0], b.reg.val[0]),
527 vaddq_f32(reg.val[1], b.reg.val[1]),
@@ -463,6 +550,85 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
550 vdivq_f32(reg.val[3], b.reg.val[3])}));
551 };
552
553 + FP32Vec16 clamp(const FP32Vec16& min, const FP32Vec16& max) const {
554 + return FP32Vec16(float32x4x4_t(
555 + {vminq_f32(max.reg.val[0], vmaxq_f32(min.reg.val[0], reg.val[0])),
556 + vminq_f32(max.reg.val[1], vmaxq_f32(min.reg.val[1], reg.val[1])),
557 + vminq_f32(max.reg.val[2], vmaxq_f32(min.reg.val[2], reg.val[2])),
558 + vminq_f32(max.reg.val[3], vmaxq_f32(min.reg.val[3], reg.val[3]))}));
559 + };
560 +
561 + FP32Vec16 max(const FP32Vec16& b) const {
562 + return FP32Vec16(float32x4x4_t({vmaxq_f32(b.reg.val[0], reg.val[0]),
563 + vmaxq_f32(b.reg.val[1], reg.val[1]),
564 + vmaxq_f32(b.reg.val[2], reg.val[2]),
565 + vmaxq_f32(b.reg.val[3], reg.val[3])}));
566 + };
567 +
568 + FP32Vec16 max(const FP32Vec16& b, const int elem_num) const {
569 + int full_blocks = elem_num / NUM_ELEMENTS_REG(reg.val[0]);
570 + int remainder = elem_num % NUM_ELEMENTS_REG(reg.val[0]);
571 + float32x4x4_t temp;
572 +
573 + for (int i = 0; i < full_blocks; i++)
574 + temp.val[i] = vmaxq_f32(b.reg.val[i], reg.val[i]);
575 +
576 + if (remainder > 0) {
577 + float max_v = std::max(vgetq_lane_f32(reg.val[full_blocks], 0),
578 + vgetq_lane_f32(b.reg.val[full_blocks], 0));
579 + temp.val[full_blocks] = vsetq_lane_f32(max_v, temp.val[full_blocks], 0);
580 + }
581 + if (remainder > 1) {
582 + float max_v = std::max(vgetq_lane_f32(reg.val[full_blocks], 1),
583 + vgetq_lane_f32(b.reg.val[full_blocks], 1));
584 + temp.val[full_blocks] = vsetq_lane_f32(max_v, temp.val[full_blocks], 1);
585 + }
586 + if (remainder > 2) {
587 + float max_v = std::max(vgetq_lane_f32(reg.val[full_blocks], 2),
588 + vgetq_lane_f32(b.reg.val[full_blocks], 2));
589 + temp.val[full_blocks] = vsetq_lane_f32(max_v, temp.val[full_blocks], 2);
590 + }
591 + return FP32Vec16(temp);
592 + };
593 +
594 + FP32Vec16 min(const FP32Vec16& b) const {
595 + return FP32Vec16(float32x4x4_t({
596 + vminq_f32(b.reg.val[0], reg.val[0]),
597 + vminq_f32(b.reg.val[1], reg.val[1]),
598 + vminq_f32(b.reg.val[2], reg.val[2]),
599 + vminq_f32(b.reg.val[3], reg.val[3]),
600 + }));
601 + };
602 + FP32Vec16 min(const FP32Vec16& b, const int elem_num) const {
603 + int full_blocks = elem_num / NUM_ELEMENTS_REG(reg.val[0]);
604 + const int remainder = elem_num % NUM_ELEMENTS_REG(reg.val[0]);
605 + float32x4x4_t temp;
606 + for (int i = 0; i < full_blocks; i++)
607 + temp.val[i] = vminq_f32(b.reg.val[i], reg.val[i]);
608 +
609 + if (remainder > 0) {
610 + float min_v = std::min(vgetq_lane_f32(reg.val[full_blocks], 0),
611 + vgetq_lane_f32(b.reg.val[full_blocks], 0));
612 + temp.val[full_blocks] = vsetq_lane_f32(min_v, temp.val[full_blocks], 0);
613 + }
614 + if (remainder > 1) {
615 + float min_v = std::min(vgetq_lane_f32(reg.val[full_blocks], 1),
616 + vgetq_lane_f32(b.reg.val[full_blocks], 1));
617 + temp.val[full_blocks] = vsetq_lane_f32(min_v, temp.val[full_blocks], 1);
618 + }
619 + if (remainder > 2) {
620 + float min_v = std::min(vgetq_lane_f32(reg.val[full_blocks], 2),
621 + vgetq_lane_f32(b.reg.val[full_blocks], 2));
622 + temp.val[full_blocks] = vsetq_lane_f32(min_v, temp.val[full_blocks], 2);
623 + }
624 +
625 + return FP32Vec16(temp);
626 + };
627 + FP32Vec16 abs() const {
628 + return FP32Vec16(
629 + float32x4x4_t({vabsq_f32(reg.val[0]), vabsq_f32(reg.val[1]),
630 + vabsq_f32(reg.val[2]), vabsq_f32(reg.val[3])}));
631 + }
632 float reduce_sum() const {
633 AliasReg ar;
634 ar.reg = reg;
@@ -473,6 +639,24 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
639 return answer;
640 };
641
642 + float reduce_max() const {
643 + AliasReg ar;
644 + ar.reg = reg;
645 + float max_v = std::numeric_limits<float>::lowest();
646 + unroll_loop<int, VEC_ELEM_NUM>(
647 + [&max_v, &ar](int i) { max_v = std::max(max_v, ar.values[i]); });
648 + return max_v;
649 + }
650 +
651 + float reduce_min() const {
652 + AliasReg ar;
653 + ar.reg = reg;
654 + float min_v = std::numeric_limits<float>::max();
655 + unroll_loop<int, VEC_ELEM_NUM>(
656 + [&min_v, &ar](int i) { min_v = std::min(min_v, ar.values[i]); });
657 + return min_v;
658 + }
659 +
660 template <int group_size>
661 float reduce_sub_sum(int idx) {
662 static_assert(VEC_ELEM_NUM % group_size == 0);
@@ -493,6 +677,83 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
677 vst1q_f32(ptr + 8, reg.val[2]);
678 vst1q_f32(ptr + 12, reg.val[3]);
679 };
680 +
681 + void save(float* ptr, const int elem_num) const {
682 + int full_blocks = elem_num / NUM_ELEMENTS_REG(reg.val[0]);
683 + int remainder = elem_num % NUM_ELEMENTS_REG(reg.val[0]);
684 +
685 + for (int i = 0; i < full_blocks; i++)
686 + vst1q_f32(
687 + reinterpret_cast<float32_t*>(ptr) + NUM_ELEMENTS_REG(reg.val[0]) * i,
688 + reg.val[i]);
689 +
690 + if (remainder > 0) {
691 + float32x4_t temp = reg.val[full_blocks];
692 + float* base = reinterpret_cast<float32_t*>(ptr) +
693 + full_blocks * NUM_ELEMENTS_REG(reg.val[0]);
694 + if (remainder > 0) base[0] = vgetq_lane_f32(temp, 0);
695 + if (remainder > 1) base[1] = vgetq_lane_f32(temp, 1);
696 + if (remainder > 2) base[2] = vgetq_lane_f32(temp, 2);
697 + }
698 + }
699 +};
700 +
701 +struct INT8Vec16 : public Vec<INT8Vec16> {
702 + constexpr static int VEC_ELEM_NUM = 16;
703 + union AliasReg {
704 + int8x16_t reg;
705 + int8_t values[VEC_ELEM_NUM];
706 + };
707 + int8x16_t reg;
708 +
709 + explicit INT8Vec16(const FP32Vec16& vec) {
710 + // Convert each 128-bit float32 vector to int32
711 + int32x4_t part0 =
712 + vcvtq_s32_f32(vec.reg.val[0]); // Convert first 128-bit block
713 + int32x4_t part1 =
714 + vcvtq_s32_f32(vec.reg.val[1]); // Convert second 128-bit block
715 + int32x4_t part2 =
716 + vcvtq_s32_f32(vec.reg.val[2]); // Convert third 128-bit block
717 + int32x4_t part3 =
718 + vcvtq_s32_f32(vec.reg.val[3]); // Convert fourth 128-bit block
719 +
720 + // Narrow each 32-bit vector to 8 bits and combine
721 + int8x8_t lower =
722 + vqmovn_s16(vcombine_s16(vqmovn_s32(part0), vqmovn_s32(part1)));
723 + int8x8_t upper =
724 + vqmovn_s16(vcombine_s16(vqmovn_s32(part2), vqmovn_s32(part3)));
725 + reg = vcombine_s8(lower, upper); // Combine to form a single 128-bit vector
726 + }
727 +
728 + void save(int8_t* ptr) const { vst1q_s8(ptr, reg); };
729 +
730 + void save(int8_t* ptr, const int elem_num) const {
731 + int full_blocks = elem_num / NUM_ELEMENTS_REG(reg);
732 + int remainder = elem_num % NUM_ELEMENTS_REG(reg);
733 +
734 + for (int i = 0; i < full_blocks; i++)
735 + vst1q_s8(reinterpret_cast<int8_t*>(ptr) + NUM_ELEMENTS_REG(reg) * i, reg);
736 + if (remainder > 0) {
737 + int8x16_t temp = reg;
738 + int8_t* base =
739 + reinterpret_cast<int8_t*>(ptr) + full_blocks * NUM_ELEMENTS_REG(reg);
740 + if (remainder > 0) base[0] = vgetq_lane_s8(temp, 0);
741 + if (remainder > 1) base[1] = vgetq_lane_s8(temp, 1);
742 + if (remainder > 2) base[2] = vgetq_lane_s8(temp, 2);
743 + if (remainder > 3) base[3] = vgetq_lane_s8(temp, 3);
744 + if (remainder > 4) base[4] = vgetq_lane_s8(temp, 4);
745 + if (remainder > 5) base[5] = vgetq_lane_s8(temp, 5);
746 + if (remainder > 6) base[6] = vgetq_lane_s8(temp, 6);
747 + if (remainder > 7) base[7] = vgetq_lane_s8(temp, 7);
748 + if (remainder > 8) base[8] = vgetq_lane_s8(temp, 8);
749 + if (remainder > 9) base[9] = vgetq_lane_s8(temp, 9);
750 + if (remainder > 10) base[10] = vgetq_lane_s8(temp, 10);
751 + if (remainder > 11) base[11] = vgetq_lane_s8(temp, 11);
752 + if (remainder > 12) base[12] = vgetq_lane_s8(temp, 12);
753 + if (remainder > 13) base[13] = vgetq_lane_s8(temp, 13);
754 + if (remainder > 14) base[14] = vgetq_lane_s8(temp, 14);
755 + }
756 + };
757 };
758
759 template <typename T>
csrc/cpu/dnnl_helper.hpp
+45 -13
@@ -57,6 +57,7 @@ class DNNLPrimitiveHelper {
57 // Note: Due to the limitation of oneDNN
58 // (https://github.com/oneapi-src/oneDNN/issues/1636), the quantized bias is
59 // not supported.
60 +
61 template <typename OutputT, typename BiasT>
62 static void gemm_s8s8_jit(const int8_t* a, const int8_t* b, OutputT* c,
63 const BiasT* bias, dnnl_dim_t M, dnnl_dim_t N,
@@ -90,6 +91,27 @@ class DNNLPrimitiveHelper {
91 }
92
93 dnnl::matmul::primitive_desc matmul_pd;
94 +// Create memory descriptors with format_tag::any for the primitive. This
95 +// enables the matmul primitive to choose memory layouts for an
96 +// optimized primitive implementation, and these layouts may differ from the
97 +// ones provided by the user.
98 +#ifdef __aarch64__
99 + auto mat_src_md = dnnl::memory::desc({M, K}, dnnl::memory::data_type::s8,
100 + dnnl::memory::format_tag::any);
101 + auto mat_weights_md = dnnl::memory::desc(
102 + {K, N}, dnnl::memory::data_type::s8, dnnl::memory::format_tag::any);
103 + auto mat_dst_md =
104 + dnnl::memory::desc({M, N}, OutputType, dnnl::memory::format_tag::any);
105 + if (bias) {
106 + dnnl::memory::desc bias_md({1, N}, BiasType, {N, 1});
107 + matmul_pd = dnnl::matmul::primitive_desc(default_engine(), mat_src_md,
108 + mat_weights_md, bias_md,
109 + mat_dst_md, attr);
110 + } else {
111 + matmul_pd = dnnl::matmul::primitive_desc(
112 + default_engine(), mat_src_md, mat_weights_md, mat_dst_md, attr);
113 + }
114 +#else
115 if (bias) {
116 dnnl::memory::desc bias_md({1, N}, BiasType, {N, 1});
117 matmul_pd = dnnl::matmul::primitive_desc(default_engine(), a_md, b_md,
@@ -98,6 +120,7 @@ class DNNLPrimitiveHelper {
120 matmul_pd = dnnl::matmul::primitive_desc(default_engine(), a_md, b_md,
121 c_md, attr);
122 }
123 +#endif
124 dnnl::matmul matmul(matmul_pd);
125
126 auto& engine = default_engine();
@@ -111,24 +134,34 @@ class DNNLPrimitiveHelper {
134 (void*)b_scales);
135
136 auto& stream = default_stream();
137 +
138 + auto mat_src_mem = a_m;
139 + auto mat_weights_mem = b_m;
140 + auto mat_dst_mem = c_m;
141 +#ifdef __aarch64__
142 + if (matmul_pd.weights_desc() != b_m.get_desc()) {
143 + mat_weights_mem = dnnl::memory(matmul_pd.weights_desc(), engine);
144 + dnnl::reorder(b_m, mat_weights_mem).execute(stream, b_m, mat_weights_mem);
145 + }
146 +#endif
147 if constexpr (InputNoScale) {
148 if (bias) {
149 dnnl::memory::desc bias_md({N}, BiasType, {1});
150 dnnl::memory bias_m(bias_md, engine, (void*)bias);
151 matmul.execute(
152 stream, {
120 - {DNNL_ARG_SRC, a_m},
121 - {DNNL_ARG_WEIGHTS, b_m},
153 + {DNNL_ARG_SRC, mat_src_mem},
154 + {DNNL_ARG_WEIGHTS, mat_weights_mem},
155 {DNNL_ARG_BIAS, bias_m},
123 - {DNNL_ARG_DST, c_m},
156 + {DNNL_ARG_DST, mat_dst_mem},
157 {DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS, b_scales_m},
158 });
159 } else {
160 matmul.execute(
161 stream, {
129 - {DNNL_ARG_SRC, a_m},
130 - {DNNL_ARG_WEIGHTS, b_m},
131 - {DNNL_ARG_DST, c_m},
162 + {DNNL_ARG_SRC, mat_src_mem},
163 + {DNNL_ARG_WEIGHTS, mat_weights_mem},
164 + {DNNL_ARG_DST, mat_dst_mem},
165 {DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS, b_scales_m},
166 });
167 }
@@ -138,19 +171,19 @@ class DNNLPrimitiveHelper {
171 dnnl::memory bias_m(bias_md, engine, (void*)bias);
172 matmul.execute(
173 stream, {
141 - {DNNL_ARG_SRC, a_m},
142 - {DNNL_ARG_WEIGHTS, b_m},
174 + {DNNL_ARG_SRC, mat_src_mem},
175 + {DNNL_ARG_WEIGHTS, mat_weights_mem},
176 {DNNL_ARG_BIAS, bias_m},
144 - {DNNL_ARG_DST, c_m},
177 + {DNNL_ARG_DST, mat_dst_mem},
178 {DNNL_ARG_ATTR_SCALES | DNNL_ARG_SRC, a_scales_m},
179 {DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS, b_scales_m},
180 });
181 } else {
182 matmul.execute(
183 stream, {
151 - {DNNL_ARG_SRC, a_m},
152 - {DNNL_ARG_WEIGHTS, b_m},
153 - {DNNL_ARG_DST, c_m},
184 + {DNNL_ARG_SRC, mat_src_mem},
185 + {DNNL_ARG_WEIGHTS, mat_weights_mem},
186 + {DNNL_ARG_DST, mat_dst_mem},
187 {DNNL_ARG_ATTR_SCALES | DNNL_ARG_SRC, a_scales_m},
188 {DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS, b_scales_m},
189 });
@@ -170,5 +203,4 @@ class DNNLPrimitiveHelper {
203 return stream;
204 }
205 };
173 -
206 #endif
csrc/cpu/quant.cpp
+12 -9
@@ -36,7 +36,7 @@ struct KernelVecType<c10::Half> {
36 using cvt_vec_type = vec_op::FP32Vec16;
37 };
38
39 -#ifdef __AVX512F__
39 +#if defined(__AVX512F__) || defined(__aarch64__)
40 template <bool AZP, typename scalar_t>
41 void static_scaled_int8_quant_impl(const scalar_t* input, int8_t* output,
42 const float* scale, const int32_t* azp,
@@ -598,8 +598,9 @@ void static_scaled_int8_quant_impl(const scalar_t* input, int8_t* output,
598 const float* scale, const int32_t* azp,
599 const int num_tokens,
600 const int hidden_size) {
601 - TORCH_CHECK(
602 - false, "static_scaled_int8_quant_impl requires AVX512/powerpc64 support.")
601 + TORCH_CHECK(false,
602 + "static_scaled_int8_quant_impl requires AVX512/powerpc64/AArch64 "
603 + "support.")
604 }
605
606 template <typename scalar_t>
@@ -607,9 +608,9 @@ void dynamic_scaled_int8_quant_impl(const scalar_t* input, int8_t* output,
608 float* scale, int32_t* azp,
609 const int num_tokens,
610 const int hidden_size) {
610 - TORCH_CHECK(
611 - false,
612 - "dynamic_scaled_int8_quant_impl requires AVX512/powerpc64 support.")
611 + TORCH_CHECK(false,
612 + "dynamic_scaled_int8_quant_impl requires "
613 + "AVX512/powerpc64/AArch64 support.")
614 }
615
616 template <bool PerChannel, typename scalar_t>
@@ -617,7 +618,8 @@ void static_quant_epilogue(const float* input, scalar_t* output,
618 const float a_scale, const float* b_scale,
619 const int32_t* azp_with_adj, const int num_tokens,
620 const int hidden_size) {
620 - TORCH_CHECK(false, "static_quant_epilogue requires AVX512/powerpc64 support.")
621 + TORCH_CHECK(
622 + false, "static_quant_epilogue requires AVX512/powerpc64/AArch64 support.")
623 }
624
625 template <typename scalar_t>
@@ -626,8 +628,9 @@ void dynamic_quant_epilogue(const float* input, scalar_t* output,
628 const int32_t* azp, const int32_t* azp_with_adj,
629 const scalar_t* bias, const int num_tokens,
630 const int hidden_size) {
629 - TORCH_CHECK(false,
630 - "dynamic_quant_epilogue requires AVX512/powerpc64 support.")
631 + TORCH_CHECK(
632 + false,
633 + "dynamic_quant_epilogue requires AVX512/powerpc64/AArch64 support.")
634 }
635 #endif
636 } // namespace
csrc/cpu/torch_bindings.cpp
+2 -1
@@ -151,8 +151,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
151 ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding);
152
153 // Quantization
154 -#ifdef __AVX512F__
154 +#if defined(__AVX512F__) || defined(__aarch64__)
155 at::Tag stride_tag = at::Tag::needs_fixed_stride_order;
156 +
157 // Compute int8 quantized tensor for given scaling factor.
158 ops.def(
159 "static_scaled_int8_quant(Tensor! out, Tensor input, Tensor scale,"