@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 061043eaca

[CPU][Perf] Accelerate unquantized MoE for AArch64 (#46353)

Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com>

Fadi Arafeh committed Jun 24, 2026 at 15:14 UTC 061043eacac0473769d86a1babb9adeb8402e833
15 files changed +672 -32
.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
+2 -2
@@ -8,7 +8,7 @@ set -ex
8 CORE_RANGE=${CORE_RANGE:-0-31}
9 OMP_CORE_RANGE=${OMP_CORE_RANGE:-0-31}
10
11 -export CMAKE_BUILD_PARALLEL_LEVEL=16
11 +export CMAKE_BUILD_PARALLEL_LEVEL=32
12
13 # Setup cleanup
14 remove_docker_container() {
@@ -37,7 +37,7 @@ function cpu_tests() {
37 pytest -x -v -s tests/kernels/test_onednn.py
38 pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
39 pytest -x -v -s tests/kernels/core/test_cpu_activation.py
40 - pytest -x -v -s tests/kernels/moe/test_moe.py -k test_cpu_fused_moe_basic
40 + pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
41 pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
42
43 # skip tests requiring model downloads if HF_TOKEN is not set
benchmarks/kernels/cpu/benchmark_cpu_fused_moe.py
+7 -8
@@ -7,6 +7,7 @@ import time
7 import numpy as np
8 import torch
9
10 +from vllm.platforms import CpuArchEnum, current_platform
11 from vllm.utils.argparse_utils import FlexibleArgumentParser
12 from vllm.utils.torch_utils import set_random_seed
13
@@ -14,17 +15,15 @@ from vllm.utils.torch_utils import set_random_seed
15 try:
16 from vllm._custom_ops import cpu_fused_moe, cpu_prepack_moe_weight
17 except (ImportError, AttributeError) as e:
17 - print("ERROR: CPU fused MoE operations are not available on this platform.")
18 - print("This benchmark requires x86 CPU with proper vLLM CPU extensions compiled.")
19 - print(
20 - "The cpu_fused_moe kernel is typically available on Linux x86_64 "
21 - "with AVX2/AVX512."
22 - )
18 print(f"Import error: {e}")
19 sys.exit(1)
20
21 # ISA selection following test_cpu_fused_moe.py pattern
27 -ISA_CHOICES = ["amx", "vec"] if torch.cpu._is_amx_tile_supported() else ["vec"]
22 +ISA_CHOICES = ["vec"]
23 +if torch.cpu._is_amx_tile_supported():
24 + ISA_CHOICES.append("amx")
25 +if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
26 + ISA_CHOICES.append("neon")
27
28
29 @torch.inference_mode()
@@ -145,7 +144,7 @@ if __name__ == "__main__":
144 "--isa",
145 type=str,
146 choices=ISA_CHOICES,
148 - default=ISA_CHOICES[0],
147 + default="vec",
148 help=f"ISA to use (available: {ISA_CHOICES})",
149 )
150 parser.add_argument("--seed", type=int, default=0)
cmake/cpu_extension.cmake
+1
@@ -423,6 +423,7 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
423 set(VLLM_EXT_SRC
424 "csrc/cpu/shm.cpp"
425 "csrc/cpu/activation_lut_bf16.cpp"
426 + "csrc/cpu/cpu_fused_moe.cpp"
427 ${VLLM_EXT_SRC})
428 endif()
429
csrc/cpu/cpu_fused_moe.cpp
+76 -14
@@ -1,5 +1,3 @@
1 -#include <sleef.h>
2 -
1 #include "cpu/cpu_types.hpp"
2 #include "cpu/utils.hpp"
3 #include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp"
@@ -16,6 +14,18 @@
14 #define AMX_DISPATCH(...) case cpu_utils::ISA::AMX:
15 #endif
16
17 +#if defined(ARM_BF16_SUPPORT)
18 + #include "cpu/micro_gemm/cpu_micro_gemm_neon.hpp"
19 + #define NEON_DISPATCH(...) \
20 + case cpu_utils::ISA::NEON: { \
21 + using gemm_t = \
22 + cpu_micro_gemm::MicroGemm<cpu_utils::ISA::NEON, scalar_t>; \
23 + return __VA_ARGS__(); \
24 + }
25 +#else
26 + #define NEON_DISPATCH(...) case cpu_utils::ISA::NEON:
27 +#endif
28 +
29 #define CPU_ISA_DISPATCH_IMPL(ISA_TYPE, ...) \
30 [&] { \
31 switch (ISA_TYPE) { \
@@ -25,6 +35,7 @@
35 cpu_micro_gemm::MicroGemm<cpu_utils::ISA::VEC, scalar_t>; \
36 return __VA_ARGS__(); \
37 } \
38 + NEON_DISPATCH(__VA_ARGS__) \
39 default: { \
40 TORCH_CHECK(false, "Invalid CPU ISA type."); \
41 } \
@@ -59,10 +70,12 @@ void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
70 const int32_t input_stride,
71 const int32_t output_stride) {
72 using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
73 +#if !defined(__aarch64__)
74 // For GPT-OSS interleaved gate-up weights
75 alignas(64) static int32_t index[16] = {0, 2, 4, 6, 8, 10, 12, 14,
76 16, 18, 20, 22, 24, 26, 28, 30};
77 vec_op::INT32Vec16 index_vec(index);
78 +#endif
79 vec_op::FP32Vec16 gate_up_max_vec(7.0);
80 vec_op::FP32Vec16 up_min_vec(-7.0);
81 vec_op::FP32Vec16 alpha_vec(1.702);
@@ -72,8 +85,15 @@ void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
85
86 for (int32_t m = 0; m < m_size; ++m) {
87 for (int32_t n = 0; n < n_size; n += 32) {
88 + // Note: AdvSIMD does not support gather loads
89 +#if defined(__aarch64__)
90 + vec_op::FP32Vec16 gate_vec(vec_op::uninit);
91 + vec_op::FP32Vec16 up_vec(vec_op::uninit);
92 + vec_op::FP32Vec16::load_even_odd(input + n, gate_vec, up_vec);
93 +#else
94 vec_op::FP32Vec16 gate_vec(input + n, index_vec);
95 vec_op::FP32Vec16 up_vec(input + n + 1, index_vec);
96 +#endif
97 gate_vec = gate_vec.min(gate_up_max_vec);
98 up_vec = up_vec.clamp(up_min_vec, gate_up_max_vec);
99 auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec * alpha_vec));
@@ -174,7 +194,7 @@ void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
194 auto inner_vec = w1_vec * (gate_vec + w3_vec * gate_pow3_vec);
195 // Note: can't use fast_exp form because diffusiongemma will generate
196 // wrong results
177 - vec_op::FP32Vec16 tanh_vec(Sleef_tanhf16_u10(inner_vec.reg));
197 + auto tanh_vec = inner_vec.tanh();
198 auto gelu_tanh = gate_vec * w2_vec * (one_vec + tanh_vec);
199 auto gated_output_fp32 = up_vec * gelu_tanh;
200 scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
@@ -240,6 +260,7 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
260 constexpr int32_t gemm_n_tile_size = gemm_t::NSize;
261 constexpr int32_t gemm_m_tile_size = gemm_t::MaxMSize;
262 constexpr int32_t min_w13_n_tile_size = 2 * gemm_n_tile_size;
263 + constexpr bool pack_a = gemm_t::PackA;
264 static_assert(gemm_n_tile_size % 16 == 0);
265
266 TORCH_CHECK_EQ(output_size_13 % min_w13_n_tile_size, 0);
@@ -266,12 +287,18 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
287
288 const int32_t w2_input_tile_size = cpu_utils::round_up<64>(
289 gemm_m_tile_size * input_size_2 * sizeof(scalar_t));
290 + // use w2 input buffer only when we need to pack input
291 + const int32_t w2_input_buffer_size =
292 + pack_a ? cpu_utils::round_up<64>(gemm_m_tile_size * input_size_2 *
293 + sizeof(scalar_t))
294 + : 0;
295
296 const int32_t w2_n_tile_size = [&]() {
297 const int64_t cache_size = cpu_utils::get_available_l2_size();
272 - // input tile + weight
298 + // input tile + optional packed input + weight
299 const int32_t n_size_cache_limit =
274 - (cache_size - w2_input_tile_size) / (input_size_2 * sizeof(scalar_t));
300 + (cache_size - (pack_a ? w2_input_buffer_size : w2_input_tile_size)) /
301 + (input_size_2 * sizeof(scalar_t));
302 const int32_t n_size_thread_limit =
303 output_size_2 / std::max(1, thread_num / topk_num);
304 const int32_t n_size = cpu_utils::round_down<gemm_n_tile_size>(
@@ -324,6 +351,9 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
351 const int32_t w13_output_buffer_offset = w13_thread_buffer_offset;
352 w13_thread_buffer_offset += w13_output_buffer_size;
353
354 + const int32_t w2_input_buffer_offset = w13_thread_buffer_offset;
355 + w13_thread_buffer_offset += w2_input_buffer_size;
356 +
357 // Weighted sum thread buffer
358 const int32_t ws_output_buffer_size =
359 cpu_utils::round_up<64>(output_size_2 * sizeof(float));
@@ -403,7 +433,8 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
433 gemm_t gemm;
434
435 const int32_t input_size_13_bytes = input_size_13 * sizeof(scalar_t);
406 - const int32_t w13_n_group_stride = 16 * input_size_13;
436 + const int32_t w13_n_group_stride =
437 + gemm_t::WeightOCGroupSize * input_size_13;
438 const int32_t w13_n_tile_stride = gemm_n_tile_size * input_size_13;
439
440 for (;;) {
@@ -466,8 +497,23 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
497 token_idx += gemm_m_tile_size) {
498 const int32_t actual_token_num =
499 std::min(gemm_m_tile_size, curr_token_num - token_idx);
469 - // copy inputs
470 - {
500 +
501 + scalar_t* __restrict__ curr_w13_gemm_input_buffer = nullptr;
502 + if constexpr (pack_a) {
503 + // copy and pack inputs
504 + curr_w13_gemm_input_buffer = w13_input_buffer;
505 + const scalar_t* w13_input_rows[gemm_m_tile_size];
506 + for (int32_t i = 0; i < actual_token_num; ++i) {
507 + w13_input_rows[i] =
508 + input + curr_expand_token_id_buffer[i] * input_size_13;
509 + }
510 + gemm_t::pack_input_from_rows(w13_input_rows,
511 + curr_w13_gemm_input_buffer,
512 + actual_token_num, input_size_13);
513 + curr_expand_token_id_buffer += actual_token_num;
514 + } else {
515 + // copy inputs
516 + curr_w13_gemm_input_buffer = curr_w13_input_buffer;
517 scalar_t* __restrict__ curr_w13_input_buffer_iter =
518 curr_w13_input_buffer;
519 for (int32_t i = 0; i < actual_token_num; ++i) {
@@ -499,14 +545,12 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
545 scalar_t* __restrict__ w13_weight_ptr_1_iter = w13_weight_ptr_1;
546 scalar_t* __restrict__ w13_bias_ptr_0_iter = w13_bias_ptr_0;
547 scalar_t* __restrict__ w13_bias_ptr_1_iter = w13_bias_ptr_1;
502 - scalar_t* __restrict__ curr_w13_input_buffer_iter =
503 - curr_w13_input_buffer;
548 float* __restrict__ w13_output_buffer_0_iter = w13_output_buffer;
549 float* __restrict__ w13_output_buffer_1_iter =
550 w13_output_buffer + actual_n_tile_size / 2;
551 for (int32_t i = 0; i < actual_n_tile_size;
552 i += min_w13_n_tile_size) {
509 - gemm.gemm(curr_w13_input_buffer_iter, w13_weight_ptr_0_iter,
553 + gemm.gemm(curr_w13_gemm_input_buffer, w13_weight_ptr_0_iter,
554 w13_output_buffer_0_iter, actual_token_num,
555 input_size_13, input_size_13, w13_n_group_stride,
556 actual_n_tile_size, false);
@@ -519,7 +563,7 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
563 w13_bias_ptr_0_iter += gemm_n_tile_size;
564 }
565
522 - gemm.gemm(curr_w13_input_buffer_iter, w13_weight_ptr_1_iter,
566 + gemm.gemm(curr_w13_gemm_input_buffer, w13_weight_ptr_1_iter,
567 w13_output_buffer_1_iter, actual_token_num,
568 input_size_13, input_size_13, w13_n_group_stride,
569 actual_n_tile_size, false);
@@ -572,7 +616,8 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
616 gemm_t gemm;
617
618 const int32_t w2_n_tile_stride = gemm_n_tile_size * input_size_2;
575 - const int32_t w2_n_group_stride = 16 * input_size_2;
619 + const int32_t w2_n_group_stride =
620 + gemm_t::WeightOCGroupSize * input_size_2;
621
622 for (;;) {
623 int32_t task_id = counter_ptr->acquire_counter();
@@ -611,13 +656,30 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
656 token_idx += gemm_m_tile_size) {
657 const int32_t actual_token_num =
658 std::min(gemm_m_tile_size, curr_token_num - token_idx);
659 + scalar_t* __restrict__ curr_w2_gemm_input_buffer =
660 + curr_w13_gemm_output_buffer;
661 + if constexpr (pack_a) {
662 + uint8_t* __restrict__ thread_buffer =
663 + thread_buffer_start + thread_id * w13_thread_buffer_offset;
664 + scalar_t* __restrict__ w2_input_buffer =
665 + reinterpret_cast<scalar_t*>(thread_buffer +
666 + w2_input_buffer_offset);
667 + curr_w2_gemm_input_buffer = w2_input_buffer;
668 + const scalar_t* w2_input_rows[gemm_m_tile_size];
669 + for (int32_t i = 0; i < actual_token_num; ++i) {
670 + w2_input_rows[i] = curr_w13_gemm_output_buffer + i * input_size_2;
671 + }
672 + gemm_t::pack_input_from_rows(w2_input_rows,
673 + curr_w2_gemm_input_buffer,
674 + actual_token_num, input_size_2);
675 + }
676
677 scalar_t* __restrict__ w2_weight_ptr_iter = w2_weight_ptr;
678 scalar_t* __restrict__ w2_bias_ptr_iter = w2_bias_ptr;
679 float* __restrict__ curr_w2_gemm_output_buffer_iter =
680 curr_w2_gemm_output_buffer;
681 for (int32_t i = 0; i < actual_n_tile_size; i += gemm_n_tile_size) {
620 - gemm.gemm(curr_w13_gemm_output_buffer, w2_weight_ptr_iter,
682 + gemm.gemm(curr_w2_gemm_input_buffer, w2_weight_ptr_iter,
683 curr_w2_gemm_output_buffer_iter, actual_token_num,
684 input_size_2, input_size_2, w2_n_group_stride,
685 output_size_2, false);
csrc/cpu/cpu_types_arm.hpp
+30 -1
@@ -497,6 +497,26 @@ struct FP32Vec16 : public VectorizedRegWrapper<FP32Vec16, 4, float> {
497 reg.val[3] = Vectorized<float>(vcvt_f32_f16(vget_high_f16(v.reg.val[1])));
498 };
499
500 + static FORCE_INLINE void load_even_odd(const float* ptr, FP32Vec16& even,
501 + FP32Vec16& odd) noexcept {
502 + const float32x4x2_t x01 = vuzpq_f32(vld1q_f32(ptr), vld1q_f32(ptr + 4));
503 + const float32x4x2_t x23 =
504 + vuzpq_f32(vld1q_f32(ptr + 8), vld1q_f32(ptr + 12));
505 + const float32x4x2_t x45 =
506 + vuzpq_f32(vld1q_f32(ptr + 16), vld1q_f32(ptr + 20));
507 + const float32x4x2_t x67 =
508 + vuzpq_f32(vld1q_f32(ptr + 24), vld1q_f32(ptr + 28));
509 +
510 + even.reg.val[0] = VectorizedT(x01.val[0]);
511 + even.reg.val[1] = VectorizedT(x23.val[0]);
512 + even.reg.val[2] = VectorizedT(x45.val[0]);
513 + even.reg.val[3] = VectorizedT(x67.val[0]);
514 + odd.reg.val[0] = VectorizedT(x01.val[1]);
515 + odd.reg.val[1] = VectorizedT(x23.val[1]);
516 + odd.reg.val[2] = VectorizedT(x45.val[1]);
517 + odd.reg.val[3] = VectorizedT(x67.val[1]);
518 + }
519 +
520 FORCE_INLINE FP32Vec16 operator+(const FP32Vec16& b) const noexcept {
521 FP32Vec16 r(uninit);
522 r.reg.val[0] = reg.val[0] + b.reg.val[0];
@@ -515,6 +535,15 @@ struct FP32Vec16 : public VectorizedRegWrapper<FP32Vec16, 4, float> {
535 return r;
536 }
537
538 + FORCE_INLINE FP32Vec16 operator-() const noexcept {
539 + FP32Vec16 r(uninit);
540 + r.reg.val[0] = reg.val[0].neg();
541 + r.reg.val[1] = reg.val[1].neg();
542 + r.reg.val[2] = reg.val[2].neg();
543 + r.reg.val[3] = reg.val[3].neg();
544 + return r;
545 + }
546 +
547 FORCE_INLINE FP32Vec16 operator*(const FP32Vec16& b) const noexcept {
548 FP32Vec16 r(uninit);
549 r.reg.val[0] = reg.val[0] * b.reg.val[0];
@@ -933,4 +962,4 @@ inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
962
963 inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 1); };
964
936 -}; // namespace vec_op
\ No newline at end of file
965 +}; // namespace vec_op
csrc/cpu/cpu_types_scalar.hpp
+7
@@ -363,6 +363,13 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
363 return FP32Vec16(ret);
364 }
365
366 + FP32Vec16 tanh() const {
367 + f32x16_t ret;
368 + unroll_loop<int, VEC_ELEM_NUM>(
369 + [&ret, this](int i) { ret.val[i] = std::tanh(reg.val[i]); });
370 + return FP32Vec16(ret);
371 + }
372 +
373 float reduce_sum() const {
374 float result = 0.0f;
375 unroll_loop<int, VEC_ELEM_NUM>(
csrc/cpu/cpu_types_x86.hpp
+9
@@ -3,6 +3,7 @@
3 #define CPU_TYPES_X86_HPP
4
5 #include <immintrin.h>
6 +#include <sleef.h>
7 #include <torch/all.h>
8
9 #ifndef __AVX2__
@@ -592,6 +593,8 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
593
594 FP32Vec16 abs() const { return FP32Vec16(_mm512_abs_ps(reg)); }
595
596 + FP32Vec16 tanh() const { return FP32Vec16(Sleef_tanhf16_u10(reg)); }
597 +
598 float reduce_sum() const { return _mm512_reduce_add_ps(reg); }
599
600 float reduce_max() const { return _mm512_reduce_max_ps(reg); }
@@ -789,6 +792,12 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
792 _mm256_andnot_ps(sign_mask, reg_high));
793 }
794
795 + FP32Vec16 tanh() const {
796 + FP32Vec8 low(reg_low);
797 + FP32Vec8 high(reg_high);
798 + return FP32Vec16(low.tanh().reg, high.tanh().reg);
799 + }
800 +
801 FP32Vec16 min(const FP32Vec16& b) const {
802 return FP32Vec16(_mm256_min_ps(reg_low, b.reg_low),
803 _mm256_min_ps(reg_high, b.reg_high));
csrc/cpu/micro_gemm/cpu_micro_gemm_amx.hpp
+2
@@ -213,6 +213,8 @@ class MicroGemm<cpu_utils::ISA::AMX, scalar_t> {
213 public:
214 static constexpr int32_t MaxMSize = 32;
215 static constexpr int32_t NSize = 32;
216 + static constexpr int32_t WeightOCGroupSize = 16;
217 + static constexpr bool PackA = false;
218
219 public:
220 MicroGemm() : curr_m_(-1) {
csrc/cpu/micro_gemm/cpu_micro_gemm_impl.hpp
+3
@@ -21,6 +21,9 @@ class MicroGemm {
21 public:
22 static constexpr int32_t MaxMSize = 16;
23 static constexpr int32_t NSize = 16;
24 + static constexpr int32_t WeightOCGroupSize = 16;
25 + // callers must pack A matrix before GEMM
26 + static constexpr bool PackA = false;
27
28 public:
29 void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
csrc/cpu/micro_gemm/cpu_micro_gemm_neon.hpp new
+503
@@ -0,0 +1,503 @@
1 +#ifndef CPU_MICRO_GEMM_NEON_HPP
2 +#define CPU_MICRO_GEMM_NEON_HPP
3 +
4 +#include <algorithm>
5 +#include <cstdint>
6 +
7 +#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
8 +
9 +#include <arm_bf16.h>
10 +#include <arm_neon.h>
11 +
12 +namespace cpu_micro_gemm {
13 +
14 +namespace {
15 +
16 +constexpr int32_t K = 4;
17 +constexpr int32_t Cols = 2;
18 +constexpr int32_t TileSize = K * Cols;
19 +constexpr int32_t Mr = 8;
20 +constexpr int32_t Nr = 8;
21 +constexpr int32_t Nr_gemv = 16;
22 +
23 +// a = [a0, a1, a2, a3], b = [b0, b1, b2, b3] -> [a0, a1, b0, b1]
24 +FORCE_INLINE float32x4_t zip1_f32x4(const float32x4_t a, const float32x4_t b) {
25 + return vreinterpretq_f32_f64(
26 + vzip1q_f64(vreinterpretq_f64_f32(a), vreinterpretq_f64_f32(b)));
27 +}
28 +
29 +// a = [a0, a1, a2, a3], b = [b0, b1, b2, b3] -> [a2, a3, b2, b3]
30 +FORCE_INLINE float32x4_t zip2_f32x4(const float32x4_t a, const float32x4_t b) {
31 + return vreinterpretq_f32_f64(
32 + vzip2q_f64(vreinterpretq_f64_f32(a), vreinterpretq_f64_f32(b)));
33 +}
34 +
35 +FORCE_INLINE void init_acc_rowpair(float32x4_t& acc01, float32x4_t& acc23,
36 + float32x4_t& acc45, float32x4_t& acc67,
37 + const float* __restrict__ c_ptr,
38 + const int64_t ldc, const int32_t m_rows,
39 + const bool accum_c) {
40 + if (!accum_c || m_rows == 0) {
41 + acc01 = vdupq_n_f32(0.0f);
42 + acc23 = vdupq_n_f32(0.0f);
43 + acc45 = vdupq_n_f32(0.0f);
44 + acc67 = vdupq_n_f32(0.0f);
45 + return;
46 + }
47 +
48 + const float32x4_t row0_0123 = vld1q_f32(c_ptr);
49 + const float32x4_t row0_4567 = vld1q_f32(c_ptr + 4);
50 + const float32x4_t row1_0123 =
51 + (m_rows == 2) ? vld1q_f32(c_ptr + ldc) : vdupq_n_f32(0.0f);
52 + const float32x4_t row1_4567 =
53 + (m_rows == 2) ? vld1q_f32(c_ptr + ldc + 4) : vdupq_n_f32(0.0f);
54 +
55 + acc01 = zip1_f32x4(row0_0123, row1_0123);
56 + acc23 = zip2_f32x4(row0_0123, row1_0123);
57 + acc45 = zip1_f32x4(row0_4567, row1_4567);
58 + acc67 = zip2_f32x4(row0_4567, row1_4567);
59 +}
60 +
61 +FORCE_INLINE void store_acc_rowpair(const float32x4_t acc01,
62 + const float32x4_t acc23,
63 + const float32x4_t acc45,
64 + const float32x4_t acc67,
65 + float* __restrict__ c_ptr,
66 + const int64_t ldc, const int32_t m_rows) {
67 + if (m_rows == 0) {
68 + return;
69 + }
70 +
71 + vst1q_f32(c_ptr, zip1_f32x4(acc01, acc23));
72 + vst1q_f32(c_ptr + 4, zip1_f32x4(acc45, acc67));
73 +
74 + if (m_rows == 2) {
75 + vst1q_f32(c_ptr + ldc, zip2_f32x4(acc01, acc23));
76 + vst1q_f32(c_ptr + ldc + 4, zip2_f32x4(acc45, acc67));
77 + }
78 +}
79 +
80 +FORCE_INLINE void gemm_micro_bfmmla_8x8_packed_a(
81 + const bfloat16_t* __restrict__ a_packed,
82 + const bfloat16_t* __restrict__ b_packed, float* __restrict__ c_ptr,
83 + const int32_t m, const int32_t k_size, const int64_t ldc,
84 + const bool accum_c) {
85 + float32x4_t acc0101, acc0123, acc0145, acc0167;
86 + float32x4_t acc2301, acc2323, acc2345, acc2367;
87 + float32x4_t acc4501, acc4523, acc4545, acc4567;
88 + float32x4_t acc6701, acc6723, acc6745, acc6767;
89 +
90 + init_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc,
91 + std::min(2, m), accum_c);
92 + init_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc,
93 + std::min(2, std::max(0, m - 2)), accum_c);
94 + init_acc_rowpair(acc4501, acc4523, acc4545, acc4567, c_ptr + 4 * ldc, ldc,
95 + std::min(2, std::max(0, m - 4)), accum_c);
96 + init_acc_rowpair(acc6701, acc6723, acc6745, acc6767, c_ptr + 6 * ldc, ldc,
97 + std::min(2, std::max(0, m - 6)), accum_c);
98 +
99 + const bfloat16_t* __restrict__ a_tile = a_packed;
100 + const bfloat16_t* __restrict__ b_tile = b_packed;
101 +
102 +#pragma GCC unroll 8
103 + for (int32_t k_idx = 0; k_idx < k_size; k_idx += K) {
104 + const bfloat16x8_t a_tile01 = vld1q_bf16(a_tile);
105 + const bfloat16x8_t a_tile23 = vld1q_bf16(a_tile + TileSize);
106 + const bfloat16x8_t a_tile45 = vld1q_bf16(a_tile + 2 * TileSize);
107 + const bfloat16x8_t a_tile67 = vld1q_bf16(a_tile + 3 * TileSize);
108 +
109 + const bfloat16x8_t b_tile01 = vld1q_bf16(b_tile);
110 + const bfloat16x8_t b_tile23 = vld1q_bf16(b_tile + TileSize);
111 + const bfloat16x8_t b_tile45 = vld1q_bf16(b_tile + 2 * TileSize);
112 + const bfloat16x8_t b_tile67 = vld1q_bf16(b_tile + 3 * TileSize);
113 +
114 + acc0101 = vbfmmlaq_f32(acc0101, a_tile01, b_tile01);
115 + acc2301 = vbfmmlaq_f32(acc2301, a_tile23, b_tile01);
116 + acc4501 = vbfmmlaq_f32(acc4501, a_tile45, b_tile01);
117 + acc6701 = vbfmmlaq_f32(acc6701, a_tile67, b_tile01);
118 +
119 + acc0123 = vbfmmlaq_f32(acc0123, a_tile01, b_tile23);
120 + acc2323 = vbfmmlaq_f32(acc2323, a_tile23, b_tile23);
121 + acc4523 = vbfmmlaq_f32(acc4523, a_tile45, b_tile23);
122 + acc6723 = vbfmmlaq_f32(acc6723, a_tile67, b_tile23);
123 +
124 + acc0145 = vbfmmlaq_f32(acc0145, a_tile01, b_tile45);
125 + acc2345 = vbfmmlaq_f32(acc2345, a_tile23, b_tile45);
126 + acc4545 = vbfmmlaq_f32(acc4545, a_tile45, b_tile45);
127 + acc6745 = vbfmmlaq_f32(acc6745, a_tile67, b_tile45);
128 +
129 + acc0167 = vbfmmlaq_f32(acc0167, a_tile01, b_tile67);
130 + acc2367 = vbfmmlaq_f32(acc2367, a_tile23, b_tile67);
131 + acc4567 = vbfmmlaq_f32(acc4567, a_tile45, b_tile67);
132 + acc6767 = vbfmmlaq_f32(acc6767, a_tile67, b_tile67);
133 +
134 + a_tile += 4 * TileSize;
135 + b_tile += Nr * K;
136 + }
137 +
138 + store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc,
139 + std::min(2, m));
140 + store_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc,
141 + std::min(2, std::max(0, m - 2)));
142 + store_acc_rowpair(acc4501, acc4523, acc4545, acc4567, c_ptr + 4 * ldc, ldc,
143 + std::min(2, std::max(0, m - 4)));
144 + store_acc_rowpair(acc6701, acc6723, acc6745, acc6767, c_ptr + 6 * ldc, ldc,
145 + std::min(2, std::max(0, m - 6)));
146 +}
147 +
148 +FORCE_INLINE void gemm_micro_bfmmla_4x16_packed_a(
149 + const bfloat16_t* __restrict__ a_packed,
150 + const bfloat16_t* __restrict__ b_packed, float* __restrict__ c_ptr,
151 + const int32_t m, const int32_t k_size, const int64_t b_n_group_stride,
152 + const int64_t ldc, const bool accum_c) {
153 + const int32_t m_rows_01 = std::min(2, m);
154 + const int32_t m_rows_23 = std::min(2, std::max(0, m - 2));
155 +
156 + float32x4_t acc0101, acc0123, acc0145, acc0167;
157 + float32x4_t acc2301, acc2323, acc2345, acc2367;
158 + float32x4_t acc0189, acc011011, acc011213, acc011415;
159 + float32x4_t acc2389, acc231011, acc231213, acc231415;
160 +
161 + init_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, m_rows_01,
162 + accum_c);
163 + init_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc,
164 + m_rows_23, accum_c);
165 + init_acc_rowpair(acc0189, acc011011, acc011213, acc011415, c_ptr + 8, ldc,
166 + m_rows_01, accum_c);
167 + init_acc_rowpair(acc2389, acc231011, acc231213, acc231415,
168 + c_ptr + 2 * ldc + 8, ldc, m_rows_23, accum_c);
169 +
170 + const bfloat16_t* __restrict__ a_tile = a_packed;
171 + const bfloat16_t* __restrict__ b_tile0 = b_packed;
172 + const bfloat16_t* __restrict__ b_tile1 = b_packed + b_n_group_stride;
173 +
174 +#pragma GCC unroll 8
175 + for (int32_t k_idx = 0; k_idx < k_size; k_idx += K) {
176 + const bfloat16x8_t a_tile01 = vld1q_bf16(a_tile);
177 + const bfloat16x8_t a_tile23 = vld1q_bf16(a_tile + TileSize);
178 + const bfloat16x8_t b_tile01 = vld1q_bf16(b_tile0);
179 + const bfloat16x8_t b_tile23 = vld1q_bf16(b_tile0 + TileSize);
180 + const bfloat16x8_t b_tile45 = vld1q_bf16(b_tile0 + 2 * TileSize);
181 + const bfloat16x8_t b_tile67 = vld1q_bf16(b_tile0 + 3 * TileSize);
182 + const bfloat16x8_t b_tile89 = vld1q_bf16(b_tile1);
183 + const bfloat16x8_t b_tile1011 = vld1q_bf16(b_tile1 + TileSize);
184 + const bfloat16x8_t b_tile1213 = vld1q_bf16(b_tile1 + 2 * TileSize);
185 + const bfloat16x8_t b_tile1415 = vld1q_bf16(b_tile1 + 3 * TileSize);
186 +
187 + acc0101 = vbfmmlaq_f32(acc0101, a_tile01, b_tile01);
188 + acc2301 = vbfmmlaq_f32(acc2301, a_tile23, b_tile01);
189 + acc0123 = vbfmmlaq_f32(acc0123, a_tile01, b_tile23);
190 + acc2323 = vbfmmlaq_f32(acc2323, a_tile23, b_tile23);
191 +
192 + acc0145 = vbfmmlaq_f32(acc0145, a_tile01, b_tile45);
193 + acc2345 = vbfmmlaq_f32(acc2345, a_tile23, b_tile45);
194 + acc0167 = vbfmmlaq_f32(acc0167, a_tile01, b_tile67);
195 + acc2367 = vbfmmlaq_f32(acc2367, a_tile23, b_tile67);
196 +
197 + acc0189 = vbfmmlaq_f32(acc0189, a_tile01, b_tile89);
198 + acc2389 = vbfmmlaq_f32(acc2389, a_tile23, b_tile89);
199 + acc011011 = vbfmmlaq_f32(acc011011, a_tile01, b_tile1011);
200 + acc231011 = vbfmmlaq_f32(acc231011, a_tile23, b_tile1011);
201 +
202 + acc011213 = vbfmmlaq_f32(acc011213, a_tile01, b_tile1213);
203 + acc231213 = vbfmmlaq_f32(acc231213, a_tile23, b_tile1213);
204 + acc011415 = vbfmmlaq_f32(acc011415, a_tile01, b_tile1415);
205 + acc231415 = vbfmmlaq_f32(acc231415, a_tile23, b_tile1415);
206 +
207 + a_tile += 2 * TileSize;
208 + b_tile0 += Nr * K;
209 + b_tile1 += Nr * K;
210 + }
211 +
212 + store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, m_rows_01);
213 + store_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc,
214 + m_rows_23);
215 + store_acc_rowpair(acc0189, acc011011, acc011213, acc011415, c_ptr + 8, ldc,
216 + m_rows_01);
217 + store_acc_rowpair(acc2389, acc231011, acc231213, acc231415,
218 + c_ptr + 2 * ldc + 8, ldc, m_rows_23);
219 +}
220 +
221 +} // namespace
222 +
223 +template <typename scalar_t>
224 +class MicroGemm<cpu_utils::ISA::NEON, scalar_t> {
225 + public:
226 + static constexpr int32_t MaxMSize = 8;
227 + static constexpr int32_t NSize = 32;
228 + static constexpr int32_t WeightOCGroupSize = Nr;
229 + static constexpr bool PackA = false;
230 +
231 + public:
232 + void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
233 + TORCH_CHECK(false, "NEON BFMMLA MicroGemm only supports bfloat16.");
234 + }
235 +
236 + static void pack_weight(const scalar_t* __restrict__ /*weight*/,
237 + scalar_t* __restrict__ /*packed_weight*/,
238 + const int32_t /*output_size*/,
239 + const int32_t /*input_size*/) {
240 + TORCH_CHECK(false, "NEON BFMMLA MicroGemm only supports bfloat16.");
241 + }
242 +};
243 +
244 +template <>
245 +class MicroGemm<cpu_utils::ISA::NEON, c10::BFloat16> {
246 + public:
247 + using scalar_t = c10::BFloat16;
248 +
249 + static constexpr int32_t MaxMSize = 8;
250 + static constexpr int32_t NSize = 32;
251 + static constexpr int32_t WeightOCGroupSize = Nr;
252 + static constexpr bool PackA = true;
253 +
254 + public:
255 + // physical layout [
256 + // M / 8; Mr is 8
257 + // K / 4; K for bfmmla is 4
258 + // 4, ; 4 row-pairs for each 8 rows
259 + // 2, ; row-pair is 2 rows
260 + // 4 ; 4 elements per row
261 + // ]
262 +
263 + static void pack_input_from_rows(const scalar_t* const* __restrict__ rows,
264 + scalar_t* __restrict__ a_packed,
265 + const int32_t m, const int32_t k) {
266 + TORCH_CHECK(m > 0 && m <= MaxMSize);
267 + TORCH_CHECK_EQ(k % K, 0);
268 +
269 + auto* __restrict__ out = reinterpret_cast<bfloat16_t*>(a_packed);
270 + const bfloat16x8_t zero_q = vdupq_n_bf16(bfloat16_t{});
271 + const bfloat16x4_t zero = vget_low_bf16(zero_q);
272 +
273 + for (int32_t row_base = 0; row_base < m; row_base += Mr) {
274 + const int32_t actual_m = std::min(Mr, m - row_base);
275 + const bfloat16_t* __restrict__ row[Mr];
276 + for (int32_t i = 0; i < actual_m; ++i) {
277 + row[i] = reinterpret_cast<const bfloat16_t*>(rows[row_base + i]);
278 + }
279 +
280 + if (actual_m == 8) {
281 + int32_t k_idx = 0;
282 + for (; k_idx + 8 <= k; k_idx += 8) {
283 + bfloat16_t* __restrict__ block0 = out;
284 + bfloat16_t* __restrict__ block1 = out + 4 * TileSize;
285 +
286 + bfloat16x8_t a0 = vld1q_bf16(row[0] + k_idx);
287 + bfloat16x8_t a1 = vld1q_bf16(row[1] + k_idx);
288 + vst1q_bf16(block0,
289 + vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
290 + vst1q_bf16(block1,
291 + vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
292 +
293 + a0 = vld1q_bf16(row[2] + k_idx);
294 + a1 = vld1q_bf16(row[3] + k_idx);
295 + vst1q_bf16(block0 + TileSize,
296 + vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
297 + vst1q_bf16(block1 + TileSize,
298 + vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
299 +
300 + a0 = vld1q_bf16(row[4] + k_idx);
301 + a1 = vld1q_bf16(row[5] + k_idx);
302 + vst1q_bf16(block0 + 2 * TileSize,
303 + vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
304 + vst1q_bf16(block1 + 2 * TileSize,
305 + vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
306 +
307 + a0 = vld1q_bf16(row[6] + k_idx);
308 + a1 = vld1q_bf16(row[7] + k_idx);
309 + vst1q_bf16(block0 + 3 * TileSize,
310 + vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
311 + vst1q_bf16(block1 + 3 * TileSize,
312 + vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
313 +
314 + out += 8 * TileSize;
315 + }
316 +
317 + for (; k_idx < k; k_idx += K) {
318 + bfloat16x4_t a0 = vld1_bf16(row[0] + k_idx);
319 + bfloat16x4_t a1 = vld1_bf16(row[1] + k_idx);
320 + vst1q_bf16(out, vcombine_bf16(a0, a1));
321 +
322 + a0 = vld1_bf16(row[2] + k_idx);
323 + a1 = vld1_bf16(row[3] + k_idx);
324 + vst1q_bf16(out + TileSize, vcombine_bf16(a0, a1));
325 +
326 + a0 = vld1_bf16(row[4] + k_idx);
327 + a1 = vld1_bf16(row[5] + k_idx);
328 + vst1q_bf16(out + 2 * TileSize, vcombine_bf16(a0, a1));
329 +
330 + a0 = vld1_bf16(row[6] + k_idx);
331 + a1 = vld1_bf16(row[7] + k_idx);
332 + vst1q_bf16(out + 3 * TileSize, vcombine_bf16(a0, a1));
333 +
334 + out += 4 * TileSize;
335 + }
336 + continue;
337 + }
338 +
339 + if (actual_m == 4) {
340 + int32_t k_idx = 0;
341 + for (; k_idx + 8 <= k; k_idx += 8) {
342 + bfloat16_t* __restrict__ block0 = out;
343 + bfloat16_t* __restrict__ block1 = out + 2 * TileSize;
344 +
345 + bfloat16x8_t a0 = vld1q_bf16(row[0] + k_idx);
346 + bfloat16x8_t a1 = vld1q_bf16(row[1] + k_idx);
347 + vst1q_bf16(block0,
348 + vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
349 + vst1q_bf16(block1,
350 + vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
351 +
352 + a0 = vld1q_bf16(row[2] + k_idx);
353 + a1 = vld1q_bf16(row[3] + k_idx);
354 + vst1q_bf16(block0 + TileSize,
355 + vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
356 + vst1q_bf16(block1 + TileSize,
357 + vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
358 +
359 + out += 4 * TileSize;
360 + }
361 +
362 + for (; k_idx < k; k_idx += K) {
363 + bfloat16x4_t a0 = vld1_bf16(row[0] + k_idx);
364 + bfloat16x4_t a1 = vld1_bf16(row[1] + k_idx);
365 + vst1q_bf16(out, vcombine_bf16(a0, a1));
366 +
367 + a0 = vld1_bf16(row[2] + k_idx);
368 + a1 = vld1_bf16(row[3] + k_idx);
369 + vst1q_bf16(out + TileSize, vcombine_bf16(a0, a1));
370 +
371 + out += 2 * TileSize;
372 + }
373 + continue;
374 + }
375 +
376 + const int32_t row_pair_count = (actual_m <= 4) ? 2 : Mr / 2;
377 +
378 + int32_t k_idx = 0;
379 + for (; k_idx + 8 <= k; k_idx += 8) {
380 + bfloat16_t* __restrict__ block0 = out;
381 + bfloat16_t* __restrict__ block1 = out + row_pair_count * TileSize;
382 +
383 + bfloat16x8_t a0 = vld1q_bf16(row[0] + k_idx);
384 + bfloat16x8_t a1 = (actual_m > 1) ? vld1q_bf16(row[1] + k_idx) : zero_q;
385 + vst1q_bf16(block0, vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
386 + vst1q_bf16(block1,
387 + vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
388 +
389 + a0 = (actual_m > 2) ? vld1q_bf16(row[2] + k_idx) : zero_q;
390 + a1 = (actual_m > 3) ? vld1q_bf16(row[3] + k_idx) : zero_q;
391 + vst1q_bf16(block0 + TileSize,
392 + vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
393 + vst1q_bf16(block1 + TileSize,
394 + vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
395 +
396 + if (actual_m > 4) {
397 + a0 = vld1q_bf16(row[4] + k_idx);
398 + a1 = (actual_m > 5) ? vld1q_bf16(row[5] + k_idx) : zero_q;
399 + vst1q_bf16(block0 + 2 * TileSize,
400 + vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
401 + vst1q_bf16(block1 + 2 * TileSize,
402 + vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
403 +
404 + a0 = (actual_m > 6) ? vld1q_bf16(row[6] + k_idx) : zero_q;
405 + a1 = (actual_m > 7) ? vld1q_bf16(row[7] + k_idx) : zero_q;
406 + vst1q_bf16(block0 + 3 * TileSize,
407 + vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
408 + vst1q_bf16(block1 + 3 * TileSize,
409 + vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
410 + }
411 +
412 + out += 2 * row_pair_count * TileSize;
413 + }
414 +
415 + for (; k_idx < k; k_idx += K) {
416 + bfloat16x4_t a0 = vld1_bf16(row[0] + k_idx);
417 + bfloat16x4_t a1 = (actual_m > 1) ? vld1_bf16(row[1] + k_idx) : zero;
418 + vst1q_bf16(out, vcombine_bf16(a0, a1));
419 +
420 + a0 = (actual_m > 2) ? vld1_bf16(row[2] + k_idx) : zero;
421 + a1 = (actual_m > 3) ? vld1_bf16(row[3] + k_idx) : zero;
422 + vst1q_bf16(out + TileSize, vcombine_bf16(a0, a1));
423 +
424 + if (actual_m > 4) {
425 + a0 = vld1_bf16(row[4] + k_idx);
426 + a1 = (actual_m > 5) ? vld1_bf16(row[5] + k_idx) : zero;
427 + vst1q_bf16(out + 2 * TileSize, vcombine_bf16(a0, a1));
428 +
429 + a0 = (actual_m > 6) ? vld1_bf16(row[6] + k_idx) : zero;
430 + a1 = (actual_m > 7) ? vld1_bf16(row[7] + k_idx) : zero;
431 + vst1q_bf16(out + 3 * TileSize, vcombine_bf16(a0, a1));
432 + }
433 + out += row_pair_count * TileSize;
434 + }
435 + }
436 + }
437 +
438 + void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
439 + (void)lda; // A is packed, so lda is not needed
440 + TORCH_CHECK_EQ(k % K, 0);
441 +
442 + for (int32_t n_idx = 0; n_idx < NSize; n_idx += Nr_gemv) {
443 + const bfloat16_t* __restrict__ b_panel =
444 + reinterpret_cast<const bfloat16_t*>(b_ptr) + n_idx * k;
445 +
446 + for (int32_t row_base = 0; row_base < m; row_base += Mr) {
447 + const int32_t panel_m = std::min(Mr, m - row_base);
448 + const bfloat16_t* __restrict__ a_panel =
449 + reinterpret_cast<const bfloat16_t*>(a_ptr) + row_base * k;
450 + float* __restrict__ c_panel = c_ptr + row_base * ldc + n_idx;
451 +
452 + if (panel_m <= 4) {
453 + gemm_micro_bfmmla_4x16_packed_a(a_panel, b_panel, c_panel, panel_m, k,
454 + b_n_group_stride, ldc, accum_c);
455 + } else {
456 + gemm_micro_bfmmla_8x8_packed_a(a_panel, b_panel, c_panel, panel_m, k,
457 + ldc, accum_c);
458 + gemm_micro_bfmmla_8x8_packed_a(a_panel, b_panel + b_n_group_stride,
459 + c_panel + Nr, panel_m, k, ldc,
460 + accum_c);
461 + }
462 + }
463 + }
464 + }
465 +
466 + // physical layout [
467 + // N / 8; Nr is 8
468 + // K / 4; K for bfmmla is 4
469 + // 4, ; 4 col-pairs for each 8 cols
470 + // 2, ; col-pair is 2 cols
471 + // 4 ; 4 elements per col
472 + // ]
473 + static void pack_weight(const c10::BFloat16* __restrict__ weight,
474 + c10::BFloat16* __restrict__ packed_weight,
475 + const int32_t output_size, const int32_t input_size) {
476 + TORCH_CHECK_EQ(output_size % NSize, 0);
477 + TORCH_CHECK_EQ(input_size % K, 0);
478 +
479 + for (int32_t o_idx = 0; o_idx < output_size; o_idx += Nr) {
480 + c10::BFloat16* __restrict__ dst = packed_weight + o_idx * input_size;
481 + for (int32_t k_idx = 0; k_idx < input_size; k_idx += K) {
482 + for (int32_t pair_idx = 0; pair_idx < Nr; pair_idx += Cols) {
483 + const c10::BFloat16* __restrict__ row0 =
484 + weight + (o_idx + pair_idx) * input_size;
485 + const c10::BFloat16* __restrict__ row1 = row0 + input_size;
486 + dst[0] = row0[k_idx + 0];
487 + dst[1] = row0[k_idx + 1];
488 + dst[2] = row0[k_idx + 2];
489 + dst[3] = row0[k_idx + 3];
490 + dst[4] = row1[k_idx + 0];
491 + dst[5] = row1[k_idx + 1];
492 + dst[6] = row1[k_idx + 2];
493 + dst[7] = row1[k_idx + 3];
494 + dst += TileSize;
495 + }
496 + }
497 + }
498 + }
499 +};
500 +
501 +} // namespace cpu_micro_gemm
502 +
503 +#endif
csrc/cpu/micro_gemm/cpu_micro_gemm_vec.hpp
+2
@@ -104,6 +104,8 @@ class MicroGemm<cpu_utils::ISA::VEC, scalar_t> {
104 public:
105 static constexpr int32_t MaxMSize = 8;
106 static constexpr int32_t NSize = 32;
107 + static constexpr int32_t WeightOCGroupSize = 16;
108 + static constexpr bool PackA = false;
109
110 public:
111 void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
csrc/cpu/torch_bindings.cpp
+2 -2
@@ -538,7 +538,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
538 #endif
539
540 // fused moe
541 -#if defined(__AVX512F__)
541 +#if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT))
542 ops.def(
543 "prepack_moe_weight(Tensor weight, Tensor(a1!) packed_weight, str isa) "
544 "-> ()");
@@ -549,7 +549,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
549 "bool skip_weighted, "
550 "str act, str isa) -> ()");
551 ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe);
552 -#endif
552 +#endif // #if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT))
553 ops.def(
554 "mla_decode_kvcache("
555 " Tensor! out, Tensor query, Tensor kv_cache,"
csrc/cpu/utils.hpp
+4 -1
@@ -2,13 +2,14 @@
2 #define UTILS_HPP
3
4 #include <atomic>
5 +#include <string>
6 #include <unistd.h>
7 #include <ATen/cpu/Utils.h>
8
9 #include "cpu/cpu_types.hpp"
10
11 namespace cpu_utils {
11 -enum class ISA { AMX, VEC, RVV };
12 +enum class ISA { AMX, VEC, RVV, NEON };
13
14 inline ISA get_isa(const std::string& isa) {
15 if (isa == "amx") {
@@ -17,6 +18,8 @@ inline ISA get_isa(const std::string& isa) {
18 return ISA::VEC;
19 } else if (isa == "rvv") {
20 return ISA::RVV;
21 + } else if (isa == "neon") {
22 + return ISA::NEON;
23 } else {
24 TORCH_CHECK(false, "Invalid isa type: " + isa);
25 }
tests/kernels/moe/test_cpu_fused_moe.py
+12 -4
@@ -5,10 +5,13 @@ import pytest
5 import torch
6
7 from tests.kernels.allclose_default import get_default_atol, get_default_rtol
8 -from vllm._custom_ops import cpu_fused_moe, cpu_prepack_moe_weight
8 +from vllm._custom_ops import (
9 + cpu_fused_moe,
10 + cpu_prepack_moe_weight,
11 +)
12 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
13 from vllm.model_executor.layers.fused_moe.cpu_fused_moe import _CPU_MOE_ACT_FN
11 -from vllm.platforms import current_platform
14 +from vllm.platforms import CpuArchEnum, current_platform
15 from vllm.utils.torch_utils import set_random_seed
16
17 if not current_platform.is_cpu():
@@ -26,8 +29,13 @@ ACT = [
29 MoEActivation.GELU,
30 MoEActivation.GELU_TANH,
31 ]
29 -USE_BIAS = [True, False]
30 -ISA = ["amx", "vec"] if torch.cpu._is_amx_tile_supported() else ["vec"]
32 +USE_BIAS = [False, True]
33 +ISA = ["vec"]
34 +if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
35 + ISA.append("neon")
36 +if torch.cpu._is_amx_tile_supported():
37 + ISA.append("amx")
38 +
39 DTYPE = [torch.bfloat16]
40
41
vllm/model_executor/layers/fused_moe/cpu_fused_moe.py
+12
@@ -16,6 +16,7 @@ from vllm._custom_ops import (
16 from vllm.model_executor.layers.activation import SiluAndMul
17 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
18 from vllm.model_executor.layers.quantization.utils.layer_utils import replace_parameter
19 +from vllm.platforms import CpuArchEnum, current_platform
20 from vllm.utils.torch_utils import direct_register_custom_op
21
22 _CPU_MOE_LAYER_CACHE = {}
@@ -312,6 +313,17 @@ class CPUFusedMOE:
313 if supports_amx:
314 return False, "none"
315
316 + supports_neon = current_platform.get_cpu_architecture() == CpuArchEnum.ARM
317 + if supports_neon:
318 + if (
319 + dtype == torch.bfloat16
320 + and w13_input_size % 4 == 0
321 + and w2_input_size % 4 == 0
322 + ):
323 + return True, "neon"
324 + else:
325 + return False, "none"
326 +
327 return True, "vec"
328
329 def init_moe_grouped_gemm(