@tduncan / Nvidia-Optimization-Stack-VLLM / commits / c314af1abf

[CPU][Perf] INT8 Fused MoE Kernel for Arm CPUs (#48637)

Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>

Fadi Arafeh committed Jul 27, 2026 at 06:53 UTC c314af1abfddff7b6cce9af578be72c496c6a5e4
12 files changed +1742 -216
cmake/cpu_extension.cmake
+19 -1
@@ -15,6 +15,7 @@ endif()
15 #
16 set(ENABLE_X86_ISA $ENV{VLLM_CPU_X86})
17 set(ENABLE_ARM_BF16 $ENV{VLLM_CPU_ARM_BF16})
18 +set(ENABLE_ARM_I8MM $ENV{VLLM_CPU_ARM_I8MM})
19 set(ENABLE_RVV_BF16 $ENV{VLLM_CPU_RVV_BF16})
20
21 include_directories("${CMAKE_SOURCE_DIR}/csrc")
@@ -96,12 +97,14 @@ if (MACOSX_FOUND AND CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
97 set(ENABLE_NUMA OFF)
98 check_sysctl(hw.optional.neon ASIMD_FOUND)
99 check_sysctl(hw.optional.arm.FEAT_BF16 ARM_BF16_FOUND)
100 + check_sysctl(hw.optional.arm.FEAT_I8MM ARM_I8MM_FOUND)
101 else()
102 find_isa(${CPUINFO} "Power11" POWER11_FOUND)
103 find_isa(${CPUINFO} "POWER10" POWER10_FOUND)
104 find_isa(${CPUINFO} "POWER9" POWER9_FOUND)
105 find_isa(${CPUINFO} "asimd" ASIMD_FOUND) # Check for ARM NEON support
106 find_isa(${CPUINFO} "bf16" ARM_BF16_FOUND) # Check for ARM BF16 support
107 + find_isa(${CPUINFO} "i8mm" ARM_I8MM_FOUND) # Check for ARM I8MM support
108 find_isa(${CPUINFO} "S390" S390_FOUND)
109 find_isa(${CPUINFO} "zvfhmin" RVV_FP16_FOUND) # Check for RISC-V Vector FP16 support
110 find_isa(${CPUINFO} "zvfbfmin" RVV_BF16_FOUND) # Check for RISC-V Vector BF16 support
@@ -111,6 +114,11 @@ else()
114 set(ARM_BF16_FOUND ON)
115 message(STATUS "ARM BF16 support enabled via VLLM_CPU_ARM_BF16 environment variable")
116 endif()
117 + if (ENABLE_ARM_I8MM)
118 + set(ARM_I8MM_FOUND ON)
119 + message(STATUS
120 + "ARM I8MM support enabled via VLLM_CPU_ARM_I8MM environment variable")
121 + endif()
122 # Some kernels (e.g. Bianbu on Spacemit X100) do not report zvfbfmin
123 # in /proc/cpuinfo despite hardware support. VLLM_CPU_RVV_BF16=1
124 # overrides the detection result.
@@ -166,6 +174,11 @@ elseif (ASIMD_FOUND)
174 message(WARNING "BF16 functionality is not available")
175 set(MARCH_FLAGS "-march=armv8.2-a+dotprod+fp16")
176 endif()
177 + if(ARM_I8MM_FOUND)
178 + message(STATUS "I8MM extension detected")
179 + string(APPEND MARCH_FLAGS "+i8mm")
180 + add_compile_definitions(ARM_I8MM_SUPPORT)
181 + endif()
182 list(APPEND CXX_COMPILE_FLAGS ${MARCH_FLAGS})
183 elseif (S390_FOUND)
184 message(STATUS "S390 detected")
@@ -447,8 +460,13 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
460 "csrc/cpu/shm.cpp"
461 "csrc/cpu/activation_lut_bf16.cpp"
462 "csrc/cpu/cpu_tanhf_neon.hpp"
450 - "csrc/cpu/cpu_fused_moe.cpp"
463 ${VLLM_EXT_SRC})
464 + if (ARM_BF16_FOUND)
465 + set(VLLM_EXT_SRC "csrc/cpu/cpu_fused_moe.cpp" ${VLLM_EXT_SRC})
466 + if (ARM_I8MM_FOUND)
467 + set(VLLM_EXT_SRC "csrc/cpu/cpu_fused_moe_int8.cpp" ${VLLM_EXT_SRC})
468 + endif()
469 + endif()
470 endif()
471
472 if (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
csrc/cpu/cpu_fused_moe.cpp
+5 -187
@@ -1,5 +1,6 @@
1 #include "cpu/cpu_types.hpp"
2 #include "cpu/utils.hpp"
3 +#include "cpu/cpu_fused_moe_activations.hpp"
4 #include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp"
5 #include "cpu/cpu_arch_macros.h"
6
@@ -43,193 +44,9 @@
44 }()
45
46 namespace {
46 -enum class FusedMOEAct {
47 - SiluAndMul,
48 - SwigluOAIAndMul,
49 - GeluAndMul,
50 - GeluTanhAndMul,
51 -};
52 -
53 -FusedMOEAct get_act_type(const std::string& act) {
54 - if (act == "silu") {
55 - return FusedMOEAct::SiluAndMul;
56 - } else if (act == "swigluoai") {
57 - return FusedMOEAct::SwigluOAIAndMul;
58 - } else if (act == "gelu") {
59 - return FusedMOEAct::GeluAndMul;
60 - } else if (act == "gelu_tanh") {
61 - return FusedMOEAct::GeluTanhAndMul;
62 - } else {
63 - TORCH_CHECK(false, "Invalid act type: " + act);
64 - }
65 -}
47
67 -template <typename scalar_t>
68 -void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
69 - const int32_t m_size, const int32_t n_size,
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);
82 - vec_op::FP32Vec16 one_vec(1.0);
83 -
84 - DEFINE_FAST_EXP
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));
100 - auto glu = gate_vec * sigmoid_vec;
101 - auto gated_output_fp32 = (one_vec + up_vec) * glu;
102 - scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
103 - gated_output.save(output + n / 2);
104 - }
105 - input += input_stride;
106 - output += output_stride;
107 - }
108 -}
109 -
110 -template <typename scalar_t>
111 -void silu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
112 - const int32_t m_size, const int32_t n_size,
113 - const int32_t input_stride, const int32_t output_stride) {
114 - using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
115 - const int32_t dim = n_size / 2;
116 - float* __restrict__ gate = input;
117 - float* __restrict__ up = input + dim;
118 - vec_op::FP32Vec16 one_vec(1.0);
119 -
120 - DEFINE_FAST_EXP
121 -
122 - for (int32_t m = 0; m < m_size; ++m) {
123 - for (int32_t n = 0; n < dim; n += 16) {
124 - vec_op::FP32Vec16 gate_vec(gate + n);
125 - vec_op::FP32Vec16 up_vec(up + n);
126 - auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec));
127 - auto silu = gate_vec * sigmoid_vec;
128 - auto gated_output_fp32 = up_vec * silu;
129 - scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
130 - gated_output.save(output + n);
131 - }
132 - gate += input_stride;
133 - up += input_stride;
134 - output += output_stride;
135 - }
136 -}
137 -
138 -template <typename scalar_t>
139 -void gelu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
140 - const int32_t m_size, const int32_t n_size,
141 - const int32_t input_stride, const int32_t output_stride) {
142 - using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
143 - const int32_t dim = n_size / 2;
144 - float* __restrict__ gate = input;
145 - float* __restrict__ up = input + dim;
146 - vec_op::FP32Vec16 one_vec(1.0);
147 - vec_op::FP32Vec16 w1_vec(M_SQRT1_2);
148 - vec_op::FP32Vec16 w2_vec(0.5);
149 - alignas(64) float temp[16];
150 -
151 - DEFINE_FAST_EXP
152 -
153 - for (int32_t m = 0; m < m_size; ++m) {
154 - for (int32_t n = 0; n < dim; n += 16) {
155 - vec_op::FP32Vec16 gate_vec(gate + n);
156 - vec_op::FP32Vec16 up_vec(up + n);
157 - auto er_input_vec = gate_vec * w1_vec;
158 -
159 - er_input_vec.save(temp);
160 - for (int32_t i = 0; i < 16; ++i) {
161 - temp[i] = std::erf(temp[i]);
162 - }
163 - vec_op::FP32Vec16 er_vec(temp);
164 - auto gelu = gate_vec * w2_vec * (one_vec + er_vec);
165 - auto gated_output_fp32 = up_vec * gelu;
166 - scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
167 - gated_output.save(output + n);
168 - }
169 - gate += input_stride;
170 - up += input_stride;
171 - output += output_stride;
172 - }
173 -}
174 -
175 -template <typename scalar_t>
176 -void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
177 - const int32_t m_size, const int32_t n_size,
178 - const int32_t input_stride,
179 - const int32_t output_stride) {
180 - using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
181 - const int32_t dim = n_size / 2;
182 - float* __restrict__ gate = input;
183 - float* __restrict__ up = input + dim;
184 - vec_op::FP32Vec16 one_vec(1.0);
185 - vec_op::FP32Vec16 w1_vec(0.7978845608028654);
186 - vec_op::FP32Vec16 w2_vec(0.5);
187 - vec_op::FP32Vec16 w3_vec(0.044715);
188 -
189 - for (int32_t m = 0; m < m_size; ++m) {
190 - for (int32_t n = 0; n < dim; n += 16) {
191 - vec_op::FP32Vec16 gate_vec(gate + n);
192 - vec_op::FP32Vec16 up_vec(up + n);
193 - auto gate_pow3_vec = gate_vec * gate_vec * gate_vec;
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
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);
201 - gated_output.save(output + n);
202 - }
203 - gate += input_stride;
204 - up += input_stride;
205 - output += output_stride;
206 - }
207 -}
208 -
209 -template <typename scalar_t>
210 -FORCE_INLINE void apply_gated_act(const FusedMOEAct act,
211 - float* __restrict__ input,
212 - scalar_t* __restrict__ output,
213 - const int32_t m, const int32_t n,
214 - const int32_t input_stride,
215 - const int32_t output_stride) {
216 - switch (act) {
217 - case FusedMOEAct::SwigluOAIAndMul:
218 - swigluoai_and_mul(input, output, m, n, input_stride, output_stride);
219 - return;
220 - case FusedMOEAct::SiluAndMul:
221 - silu_and_mul(input, output, m, n, input_stride, output_stride);
222 - return;
223 - case FusedMOEAct::GeluAndMul:
224 - gelu_and_mul(input, output, m, n, input_stride, output_stride);
225 - return;
226 - case FusedMOEAct::GeluTanhAndMul:
227 - gelu_tanh_and_mul(input, output, m, n, input_stride, output_stride);
228 - return;
229 - default:
230 - TORCH_CHECK(false, "Unsupported act type.");
231 - }
232 -}
48 +using cpu_fused_moe_utils::apply_gated_act;
49 +using cpu_fused_moe_utils::FusedMOEAct;
50
51 template <typename scalar_t, typename gemm_t>
52 void prepack_moe_weight_impl(scalar_t* __restrict__ weight_ptr,
@@ -817,6 +634,7 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
634 }
635 }
636 }
637 +
638 } // namespace
639
640 void prepack_moe_weight(
@@ -864,7 +682,7 @@ void cpu_fused_moe(
682 const int32_t input_size_2 = w2.size(2);
683 const int32_t output_size_2 = w2.size(1);
684 const int32_t topk_num = topk_id.size(1);
867 - const FusedMOEAct act_type = get_act_type(act);
685 + const FusedMOEAct act_type = cpu_fused_moe_utils::get_act_type(act);
686 cpu_utils::ISA isa_type = cpu_utils::get_isa(isa);
687 TORCH_CHECK(!skip_weighted || topk_num == 1,
688 "skip_weighted is only supported for topk=1 on CPU");
csrc/cpu/cpu_fused_moe_activations.hpp new
+204
@@ -0,0 +1,204 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +#ifndef CPU_FUSED_MOE_ACTIVATIONS_HPP
5 +#define CPU_FUSED_MOE_ACTIVATIONS_HPP
6 +
7 +#include <cmath>
8 +#include <cstdint>
9 +#include <string>
10 +
11 +#include "cpu/cpu_arch_macros.h"
12 +#include "cpu/utils.hpp"
13 +
14 +namespace cpu_fused_moe_utils {
15 +enum class FusedMOEAct {
16 + SiluAndMul,
17 + SwigluOAIAndMul,
18 + GeluAndMul,
19 + GeluTanhAndMul,
20 +};
21 +
22 +inline FusedMOEAct get_act_type(const std::string& act) {
23 + if (act == "silu") {
24 + return FusedMOEAct::SiluAndMul;
25 + } else if (act == "swigluoai") {
26 + return FusedMOEAct::SwigluOAIAndMul;
27 + } else if (act == "gelu") {
28 + return FusedMOEAct::GeluAndMul;
29 + } else if (act == "gelu_tanh") {
30 + return FusedMOEAct::GeluTanhAndMul;
31 + } else {
32 + TORCH_CHECK(false, "Invalid act type: " + act);
33 + }
34 +}
35 +
36 +template <typename scalar_t>
37 +void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
38 + const int32_t m_size, const int32_t n_size,
39 + const int32_t input_stride,
40 + const int32_t output_stride) {
41 + using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
42 +#if !defined(__aarch64__)
43 + // For GPT-OSS interleaved gate-up weights
44 + alignas(64) static int32_t index[16] = {0, 2, 4, 6, 8, 10, 12, 14,
45 + 16, 18, 20, 22, 24, 26, 28, 30};
46 + vec_op::INT32Vec16 index_vec(index);
47 +#endif
48 + vec_op::FP32Vec16 gate_up_max_vec(7.0);
49 + vec_op::FP32Vec16 up_min_vec(-7.0);
50 + vec_op::FP32Vec16 alpha_vec(1.702);
51 + vec_op::FP32Vec16 one_vec(1.0);
52 +
53 + DEFINE_FAST_EXP
54 +
55 + for (int32_t m = 0; m < m_size; ++m) {
56 + for (int32_t n = 0; n < n_size; n += 32) {
57 + // Note: AdvSIMD does not support gather loads
58 +#if defined(__aarch64__)
59 + vec_op::FP32Vec16 gate_vec(vec_op::uninit);
60 + vec_op::FP32Vec16 up_vec(vec_op::uninit);
61 + vec_op::FP32Vec16::load_even_odd(input + n, gate_vec, up_vec);
62 +#else
63 + vec_op::FP32Vec16 gate_vec(input + n, index_vec);
64 + vec_op::FP32Vec16 up_vec(input + n + 1, index_vec);
65 +#endif
66 + gate_vec = gate_vec.min(gate_up_max_vec);
67 + up_vec = up_vec.clamp(up_min_vec, gate_up_max_vec);
68 + auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec * alpha_vec));
69 + auto glu = gate_vec * sigmoid_vec;
70 + auto gated_output_fp32 = (one_vec + up_vec) * glu;
71 + scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
72 + gated_output.save(output + n / 2);
73 + }
74 + input += input_stride;
75 + output += output_stride;
76 + }
77 +}
78 +
79 +template <typename scalar_t>
80 +void silu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
81 + const int32_t m_size, const int32_t n_size,
82 + const int32_t input_stride, const int32_t output_stride) {
83 + using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
84 + const int32_t dim = n_size / 2;
85 + float* __restrict__ gate = input;
86 + float* __restrict__ up = input + dim;
87 + vec_op::FP32Vec16 one_vec(1.0);
88 +
89 + DEFINE_FAST_EXP
90 +
91 + for (int32_t m = 0; m < m_size; ++m) {
92 + for (int32_t n = 0; n < dim; n += 16) {
93 + vec_op::FP32Vec16 gate_vec(gate + n);
94 + vec_op::FP32Vec16 up_vec(up + n);
95 + auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec));
96 + auto silu = gate_vec * sigmoid_vec;
97 + auto gated_output_fp32 = up_vec * silu;
98 + scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
99 + gated_output.save(output + n);
100 + }
101 + gate += input_stride;
102 + up += input_stride;
103 + output += output_stride;
104 + }
105 +}
106 +
107 +template <typename scalar_t>
108 +void gelu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
109 + const int32_t m_size, const int32_t n_size,
110 + const int32_t input_stride, const int32_t output_stride) {
111 + using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
112 + const int32_t dim = n_size / 2;
113 + float* __restrict__ gate = input;
114 + float* __restrict__ up = input + dim;
115 + vec_op::FP32Vec16 one_vec(1.0);
116 + vec_op::FP32Vec16 w1_vec(M_SQRT1_2);
117 + vec_op::FP32Vec16 w2_vec(0.5);
118 + alignas(64) float temp[16];
119 +
120 + DEFINE_FAST_EXP
121 +
122 + for (int32_t m = 0; m < m_size; ++m) {
123 + for (int32_t n = 0; n < dim; n += 16) {
124 + vec_op::FP32Vec16 gate_vec(gate + n);
125 + vec_op::FP32Vec16 up_vec(up + n);
126 + auto er_input_vec = gate_vec * w1_vec;
127 +
128 + er_input_vec.save(temp);
129 + for (int32_t i = 0; i < 16; ++i) {
130 + temp[i] = std::erf(temp[i]);
131 + }
132 + vec_op::FP32Vec16 er_vec(temp);
133 + auto gelu = gate_vec * w2_vec * (one_vec + er_vec);
134 + auto gated_output_fp32 = up_vec * gelu;
135 + scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
136 + gated_output.save(output + n);
137 + }
138 + gate += input_stride;
139 + up += input_stride;
140 + output += output_stride;
141 + }
142 +}
143 +
144 +template <typename scalar_t>
145 +void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
146 + const int32_t m_size, const int32_t n_size,
147 + const int32_t input_stride,
148 + const int32_t output_stride) {
149 + using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
150 + const int32_t dim = n_size / 2;
151 + float* __restrict__ gate = input;
152 + float* __restrict__ up = input + dim;
153 + vec_op::FP32Vec16 one_vec(1.0);
154 + vec_op::FP32Vec16 w1_vec(0.7978845608028654);
155 + vec_op::FP32Vec16 w2_vec(0.5);
156 + vec_op::FP32Vec16 w3_vec(0.044715);
157 +
158 + for (int32_t m = 0; m < m_size; ++m) {
159 + for (int32_t n = 0; n < dim; n += 16) {
160 + vec_op::FP32Vec16 gate_vec(gate + n);
161 + vec_op::FP32Vec16 up_vec(up + n);
162 + auto gate_pow3_vec = gate_vec * gate_vec * gate_vec;
163 + auto inner_vec = w1_vec * (gate_vec + w3_vec * gate_pow3_vec);
164 + // Note: can't use fast_exp form because diffusiongemma will generate
165 + // wrong results
166 + auto tanh_vec = inner_vec.tanh();
167 + auto gelu_tanh = gate_vec * w2_vec * (one_vec + tanh_vec);
168 + auto gated_output_fp32 = up_vec * gelu_tanh;
169 + scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
170 + gated_output.save(output + n);
171 + }
172 + gate += input_stride;
173 + up += input_stride;
174 + output += output_stride;
175 + }
176 +}
177 +
178 +template <typename scalar_t>
179 +FORCE_INLINE void apply_gated_act(const FusedMOEAct act,
180 + float* __restrict__ input,
181 + scalar_t* __restrict__ output,
182 + const int32_t m, const int32_t n,
183 + const int32_t input_stride,
184 + const int32_t output_stride) {
185 + switch (act) {
186 + case FusedMOEAct::SwigluOAIAndMul:
187 + swigluoai_and_mul(input, output, m, n, input_stride, output_stride);
188 + return;
189 + case FusedMOEAct::SiluAndMul:
190 + silu_and_mul(input, output, m, n, input_stride, output_stride);
191 + return;
192 + case FusedMOEAct::GeluAndMul:
193 + gelu_and_mul(input, output, m, n, input_stride, output_stride);
194 + return;
195 + case FusedMOEAct::GeluTanhAndMul:
196 + gelu_tanh_and_mul(input, output, m, n, input_stride, output_stride);
197 + return;
198 + default:
199 + TORCH_CHECK(false, "Unsupported act type.");
200 + }
201 +}
202 +} // namespace cpu_fused_moe_utils
203 +
204 +#endif
csrc/cpu/cpu_fused_moe_int8.cpp new
+647
@@ -0,0 +1,647 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +#include "cpu/cpu_arch_macros.h"
5 +
6 +#include <algorithm>
7 +#include <cstdint>
8 +#include <cstring>
9 +#include <optional>
10 +#include <string>
11 +
12 +#include "cpu/cpu_fused_moe_activations.hpp"
13 +#include "cpu/cpu_types.hpp"
14 +#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
15 +#include "cpu/utils.hpp"
16 +
17 +#if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT)
18 + #include "cpu/micro_gemm/cpu_micro_gemm_int8_neon.hpp"
19 + #define NEON_DISPATCH(SCALAR_TYPE, ...) \
20 + case cpu_utils::ISA::NEON: { \
21 + using gemm_t = \
22 + cpu_micro_gemm::MicroGemmINT8<cpu_utils::ISA::NEON, SCALAR_TYPE>; \
23 + return __VA_ARGS__(); \
24 + }
25 +#else
26 + #define NEON_DISPATCH(SCALAR_TYPE, ...) case cpu_utils::ISA::NEON:
27 +#endif
28 +
29 +#define CPU_INT8_ISA_DISPATCH_IMPL(ISA_TYPE, SCALAR_TYPE, ...) \
30 + [&] { \
31 + switch (ISA_TYPE) { \
32 + NEON_DISPATCH(SCALAR_TYPE, __VA_ARGS__) \
33 + default: { \
34 + TORCH_CHECK(false, "Invalid CPU ISA type."); \
35 + } \
36 + } \
37 + }()
38 +
39 +namespace {
40 +using cpu_fused_moe_utils::apply_gated_act;
41 +using cpu_fused_moe_utils::FusedMOEAct;
42 +
43 +template <typename gemm_t>
44 +void prepack_moe_weight_int8_impl(const int8_t* __restrict__ weight_ptr,
45 + int8_t* __restrict__ packed_weight_ptr,
46 + const int32_t expert_num,
47 + const int32_t output_size,
48 + const int32_t input_size,
49 + const int64_t expert_stride) {
50 +#pragma omp parallel for
51 + for (int32_t e_idx = 0; e_idx < expert_num; ++e_idx) {
52 + gemm_t::pack_weight(weight_ptr + expert_stride * e_idx,
53 + packed_weight_ptr + expert_stride * e_idx, output_size,
54 + input_size);
55 + }
56 +}
57 +
58 +// INT8 MoE kernel, based on the original BF16 kernel in cpu_fused_moe.cpp
59 +template <typename scalar_t, typename gemm_t>
60 +void fused_moe_int8_impl(
61 + scalar_t* __restrict__ output, const scalar_t* __restrict__ input,
62 + const int8_t* __restrict__ w13, const int8_t* __restrict__ w2,
63 + const float* __restrict__ w13_scales, const float* __restrict__ w2_scales,
64 + scalar_t* __restrict__ w13_bias, scalar_t* __restrict__ w2_bias,
65 + const float* __restrict__ topk_weights, const int32_t* __restrict__ topk_id,
66 + const FusedMOEAct act_type, const int32_t token_num,
67 + const int32_t expert_num, const int32_t topk_num,
68 + const int32_t input_size_13, const int32_t output_size_13,
69 + const int32_t input_size_2, const int32_t output_size_2,
70 + const bool skip_weighted) {
71 + using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
72 + constexpr int32_t gemm_n_tile_size = gemm_t::NSize;
73 + constexpr int32_t gemm_m_tile_size = gemm_t::MaxMSize;
74 + constexpr int32_t min_w13_n_tile_size = 2 * gemm_n_tile_size;
75 +
76 + TORCH_CHECK_EQ(input_size_13 % gemm_t::K, 0);
77 + TORCH_CHECK_EQ(input_size_2 % gemm_t::K, 0);
78 + TORCH_CHECK_EQ(output_size_13 % min_w13_n_tile_size, 0);
79 + TORCH_CHECK_EQ(output_size_2 % gemm_n_tile_size, 0);
80 + TORCH_CHECK_EQ(output_size_13 / 2, input_size_2);
81 +
82 + const int32_t thread_num = cpu_utils::get_max_threads();
83 + const int32_t w13_input_buffer_size = cpu_utils::round_up<64>(
84 + gemm_m_tile_size * input_size_13 * sizeof(int8_t));
85 + const int32_t w2_input_buffer_size =
86 + cpu_utils::round_up<64>(gemm_m_tile_size * input_size_2 * sizeof(int8_t));
87 +
88 + const int32_t w13_n_tile_size = [&]() {
89 + const int64_t cache_size = cpu_utils::get_available_l2_size();
90 + const int32_t n_size_cache_limit =
91 + (cache_size - w13_input_buffer_size) /
92 + (gemm_m_tile_size * sizeof(float) + input_size_13 * sizeof(int8_t));
93 + const int32_t n_size_thread_limit =
94 + output_size_13 / std::max(1, thread_num / topk_num);
95 + const int32_t n_size = cpu_utils::round_down<min_w13_n_tile_size>(
96 + std::min(n_size_cache_limit, n_size_thread_limit));
97 + return std::max(n_size, min_w13_n_tile_size);
98 + }();
99 +
100 + const int32_t w2_n_tile_size = [&]() {
101 + const int64_t cache_size = cpu_utils::get_available_l2_size();
102 + const int32_t n_size_cache_limit =
103 + (cache_size - w2_input_buffer_size) / (input_size_2 * sizeof(int8_t));
104 + const int32_t n_size_thread_limit =
105 + output_size_2 / std::max(1, thread_num / topk_num);
106 + const int32_t n_size = cpu_utils::round_down<gemm_n_tile_size>(
107 + std::min(n_size_cache_limit, n_size_thread_limit));
108 + return std::max(n_size, gemm_n_tile_size);
109 + }();
110 +
111 + int32_t common_buffer_offset = 0;
112 + const int32_t token_num_per_group_buffer_offset = common_buffer_offset;
113 + common_buffer_offset += cpu_utils::round_up<64>(expert_num * sizeof(int32_t));
114 + const int32_t cu_token_num_per_group_buffer_offset = common_buffer_offset;
115 + common_buffer_offset +=
116 + cpu_utils::round_up<64>((expert_num + 1) * sizeof(int32_t));
117 + const int32_t expanded_token_num = token_num * topk_num;
118 + const int32_t expand_token_id_buffer_offset = common_buffer_offset;
119 + common_buffer_offset +=
120 + cpu_utils::round_up<64>(expanded_token_num * sizeof(int32_t));
121 + const int32_t expand_token_id_index_buffer_offset = common_buffer_offset;
122 + common_buffer_offset +=
123 + cpu_utils::round_up<64>(expanded_token_num * sizeof(int32_t));
124 + const int32_t input_quant_buffer_offset = common_buffer_offset;
125 + common_buffer_offset +=
126 + cpu_utils::round_up<64>(token_num * input_size_13 * sizeof(int8_t));
127 + const int32_t input_scale_buffer_offset = common_buffer_offset;
128 + common_buffer_offset += cpu_utils::round_up<64>(token_num * sizeof(float));
129 + const int32_t w13_gemm_output_buffer_offset = common_buffer_offset;
130 + common_buffer_offset += cpu_utils::round_up<64>(
131 + expanded_token_num * input_size_2 * sizeof(scalar_t));
132 + const int32_t w13_output_scale_buffer_offset = common_buffer_offset;
133 + common_buffer_offset +=
134 + cpu_utils::round_up<64>(expanded_token_num * sizeof(float));
135 + const int32_t w2_gemm_output_buffer_offset = common_buffer_offset;
136 + common_buffer_offset += cpu_utils::round_up<64>(
137 + expanded_token_num * output_size_2 * sizeof(float));
138 +
139 + int32_t gemm_thread_buffer_offset = 0;
140 + const int32_t gemm_input_buffer_offset = gemm_thread_buffer_offset;
141 + gemm_thread_buffer_offset +=
142 + std::max(w13_input_buffer_size, w2_input_buffer_size);
143 + const int32_t gemm_output_buffer_offset = gemm_thread_buffer_offset;
144 + gemm_thread_buffer_offset += cpu_utils::round_up<64>(
145 + gemm_m_tile_size * std::max(w13_n_tile_size, w2_n_tile_size) *
146 + sizeof(int32_t));
147 +
148 + const int32_t ws_output_buffer_offset = 0;
149 + const int32_t ws_thread_buffer_size =
150 + cpu_utils::round_up<64>(output_size_2 * sizeof(float));
151 + const int32_t thread_buffer_size =
152 + std::max(gemm_thread_buffer_offset, ws_thread_buffer_size);
153 + const int32_t buffer_size =
154 + common_buffer_offset + thread_buffer_size * thread_num;
155 + cpu_utils::ScratchPadManager::get_scratchpad_manager()->realloc(buffer_size);
156 + uint8_t* common_buffer_start =
157 + cpu_utils::ScratchPadManager::get_scratchpad_manager()
158 + ->get_data<uint8_t>();
159 + uint8_t* thread_buffer_start = common_buffer_start + common_buffer_offset;
160 +
161 + int32_t* __restrict__ token_num_per_group_buffer = reinterpret_cast<int32_t*>(
162 + common_buffer_start + token_num_per_group_buffer_offset);
163 + int32_t* __restrict__ cu_token_num_per_group_buffer =
164 + reinterpret_cast<int32_t*>(common_buffer_start +
165 + cu_token_num_per_group_buffer_offset);
166 + int32_t* __restrict__ expand_token_id_buffer = reinterpret_cast<int32_t*>(
167 + common_buffer_start + expand_token_id_buffer_offset);
168 + int32_t* __restrict__ expand_token_id_index_buffer =
169 + reinterpret_cast<int32_t*>(common_buffer_start +
170 + expand_token_id_index_buffer_offset);
171 + int8_t* __restrict__ input_quant_buffer = reinterpret_cast<int8_t*>(
172 + common_buffer_start + input_quant_buffer_offset);
173 + float* __restrict__ input_scale_buffer =
174 + reinterpret_cast<float*>(common_buffer_start + input_scale_buffer_offset);
175 +
176 + std::memset(token_num_per_group_buffer, 0, expert_num * sizeof(int32_t));
177 + for (int32_t i = 0; i < expanded_token_num; ++i) {
178 + ++token_num_per_group_buffer[topk_id[i]];
179 + }
180 +
181 + int32_t token_num_sum = 0;
182 + cu_token_num_per_group_buffer[0] = 0;
183 + int32_t* token_index_buffer = cu_token_num_per_group_buffer + 1;
184 + for (int32_t i = 0; i < expert_num; ++i) {
185 + token_index_buffer[i] = token_num_sum;
186 + token_num_sum += token_num_per_group_buffer[i];
187 + }
188 +
189 + for (int32_t i = 0; i < token_num; ++i) {
190 + const int32_t* curr_topk_id = topk_id + i * topk_num;
191 + int32_t* curr_index_buffer = expand_token_id_index_buffer + i * topk_num;
192 + for (int32_t j = 0; j < topk_num; ++j) {
193 + const int32_t curr_expert_id = curr_topk_id[j];
194 + const int32_t curr_index = token_index_buffer[curr_expert_id]++;
195 + expand_token_id_buffer[curr_index] = i;
196 + curr_index_buffer[j] = curr_index;
197 + }
198 + }
199 +
200 +// quantize inputs
201 +#pragma omp parallel for
202 + for (int32_t token_idx = 0; token_idx < token_num; ++token_idx) {
203 + gemm_t::quantize_row(input + token_idx * input_size_13,
204 + input_quant_buffer + token_idx * input_size_13,
205 + input_scale_buffer[token_idx], input_size_13);
206 + }
207 +
208 + {
209 + alignas(64) cpu_utils::Counter counter;
210 + cpu_utils::Counter* counter_ptr = &counter;
211 +
212 +// w13 GEMM + act
213 +#pragma omp parallel for schedule(static, 1)
214 + for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) {
215 + const int32_t task_num_per_expert =
216 + (output_size_13 + w13_n_tile_size - 1) / w13_n_tile_size;
217 + const int32_t task_num = task_num_per_expert * expert_num;
218 + uint8_t* __restrict__ thread_buffer =
219 + thread_buffer_start + thread_id * thread_buffer_size;
220 + int8_t* __restrict__ gemm_input_buffer =
221 + reinterpret_cast<int8_t*>(thread_buffer + gemm_input_buffer_offset);
222 + float* __restrict__ gemm_output_buffer =
223 + reinterpret_cast<float*>(thread_buffer + gemm_output_buffer_offset);
224 + auto* __restrict__ w13_gemm_output_buffer = reinterpret_cast<scalar_t*>(
225 + common_buffer_start + w13_gemm_output_buffer_offset);
226 + gemm_t gemm;
227 +
228 + const int32_t w13_n_group_stride =
229 + gemm_t::WeightOCGroupSize * input_size_13;
230 + const int32_t w13_n_tile_stride = gemm_n_tile_size * input_size_13;
231 +
232 + for (;;) {
233 + const int32_t task_id = counter_ptr->acquire_counter();
234 + if (task_id >= task_num) {
235 + break;
236 + }
237 + const int32_t curr_expert_id = task_id / task_num_per_expert;
238 + const int32_t curr_output_group_id = task_id % task_num_per_expert;
239 + const int32_t curr_token_num =
240 + token_num_per_group_buffer[curr_expert_id];
241 + if (curr_token_num == 0) {
242 + continue;
243 + }
244 +
245 + const int32_t actual_n_tile_size =
246 + std::min(w13_n_tile_size,
247 + output_size_13 - curr_output_group_id * w13_n_tile_size);
248 + const int32_t* __restrict__ curr_expand_token_id_buffer =
249 + expand_token_id_buffer +
250 + cu_token_num_per_group_buffer[curr_expert_id];
251 + scalar_t* __restrict__ curr_w13_gemm_output_buffer =
252 + w13_gemm_output_buffer +
253 + cu_token_num_per_group_buffer[curr_expert_id] * input_size_2 +
254 + curr_output_group_id * w13_n_tile_size / 2;
255 +
256 + const int8_t* w13_weight_ptr_0 = nullptr;
257 + const int8_t* w13_weight_ptr_1 = nullptr;
258 + const float* w13_scale_ptr_0 = nullptr;
259 + const float* w13_scale_ptr_1 = nullptr;
260 + scalar_t* w13_bias_ptr_0 = nullptr;
261 + scalar_t* w13_bias_ptr_1 = nullptr;
262 + if (act_type == FusedMOEAct::SwigluOAIAndMul) {
263 + const int32_t output_offset = curr_output_group_id * w13_n_tile_size;
264 + w13_weight_ptr_0 = w13 +
265 + curr_expert_id * input_size_13 * output_size_13 +
266 + output_offset * input_size_13;
267 + w13_weight_ptr_1 =
268 + w13_weight_ptr_0 + actual_n_tile_size / 2 * input_size_13;
269 + w13_scale_ptr_0 =
270 + w13_scales + curr_expert_id * output_size_13 + output_offset;
271 + w13_scale_ptr_1 = w13_scale_ptr_0 + actual_n_tile_size / 2;
272 + if (w13_bias != nullptr) {
273 + w13_bias_ptr_0 =
274 + w13_bias + curr_expert_id * output_size_13 + output_offset;
275 + w13_bias_ptr_1 = w13_bias_ptr_0 + actual_n_tile_size / 2;
276 + }
277 + } else {
278 + const int32_t output_offset =
279 + curr_output_group_id * (w13_n_tile_size / 2);
280 + w13_weight_ptr_0 = w13 +
281 + curr_expert_id * input_size_13 * output_size_13 +
282 + output_offset * input_size_13;
283 + w13_weight_ptr_1 =
284 + w13_weight_ptr_0 + output_size_13 / 2 * input_size_13;
285 + w13_scale_ptr_0 =
286 + w13_scales + curr_expert_id * output_size_13 + output_offset;
287 + w13_scale_ptr_1 = w13_scale_ptr_0 + output_size_13 / 2;
288 + if (w13_bias != nullptr) {
289 + w13_bias_ptr_0 =
290 + w13_bias + curr_expert_id * output_size_13 + output_offset;
291 + w13_bias_ptr_1 = w13_bias_ptr_0 + output_size_13 / 2;
292 + }
293 + }
294 +
295 + for (int32_t token_idx = 0; token_idx < curr_token_num;
296 + token_idx += gemm_m_tile_size) {
297 + const int32_t actual_token_num =
298 + std::min(gemm_m_tile_size, curr_token_num - token_idx);
299 + const int8_t* input_rows[gemm_m_tile_size];
300 + alignas(64) float input_scales[gemm_m_tile_size];
301 + // gather and pack
302 + for (int32_t i = 0; i < actual_token_num; ++i) {
303 + const int32_t curr_token_id = curr_expand_token_id_buffer[i];
304 + input_rows[i] = input_quant_buffer + curr_token_id * input_size_13;
305 + input_scales[i] = input_scale_buffer[curr_token_id];
306 + }
307 + gemm_t::pack_input_from_rows(input_rows, gemm_input_buffer,
308 + actual_token_num, input_size_13);
309 + curr_expand_token_id_buffer += actual_token_num;
310 +
311 + const int8_t* w13_weight_ptr_0_iter = w13_weight_ptr_0;
312 + const int8_t* w13_weight_ptr_1_iter = w13_weight_ptr_1;
313 + const float* w13_scale_ptr_0_iter = w13_scale_ptr_0;
314 + const float* w13_scale_ptr_1_iter = w13_scale_ptr_1;
315 + scalar_t* w13_bias_ptr_0_iter = w13_bias_ptr_0;
316 + scalar_t* w13_bias_ptr_1_iter = w13_bias_ptr_1;
317 + float* w13_output_buffer_0_iter = gemm_output_buffer;
318 + float* w13_output_buffer_1_iter =
319 + gemm_output_buffer + actual_n_tile_size / 2;
320 +
321 + for (int32_t i = 0; i < actual_n_tile_size;
322 + i += min_w13_n_tile_size) {
323 + auto* output_0_int32 =
324 + reinterpret_cast<int32_t*>(w13_output_buffer_0_iter);
325 + gemm.gemm(gemm_input_buffer, w13_weight_ptr_0_iter, output_0_int32,
326 + actual_token_num, input_size_13, w13_n_group_stride,
327 + actual_n_tile_size);
328 + gemm_t::dequantize_tile(output_0_int32, w13_output_buffer_0_iter,
329 + input_scales, w13_scale_ptr_0_iter,
330 + actual_token_num, gemm_n_tile_size,
331 + actual_n_tile_size);
332 + if (w13_bias != nullptr) {
333 + cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>(
334 + w13_output_buffer_0_iter, w13_output_buffer_0_iter,
335 + w13_bias_ptr_0_iter, actual_token_num, actual_n_tile_size,
336 + actual_n_tile_size);
337 + w13_bias_ptr_0_iter += gemm_n_tile_size;
338 + }
339 +
340 + auto* output_1_int32 =
341 + reinterpret_cast<int32_t*>(w13_output_buffer_1_iter);
342 + gemm.gemm(gemm_input_buffer, w13_weight_ptr_1_iter, output_1_int32,
343 + actual_token_num, input_size_13, w13_n_group_stride,
344 + actual_n_tile_size);
345 + gemm_t::dequantize_tile(output_1_int32, w13_output_buffer_1_iter,
346 + input_scales, w13_scale_ptr_1_iter,
347 + actual_token_num, gemm_n_tile_size,
348 + actual_n_tile_size);
349 + if (w13_bias != nullptr) {
350 + cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>(
351 + w13_output_buffer_1_iter, w13_output_buffer_1_iter,
352 + w13_bias_ptr_1_iter, actual_token_num, actual_n_tile_size,
353 + actual_n_tile_size);
354 + w13_bias_ptr_1_iter += gemm_n_tile_size;
355 + }
356 +
357 + w13_weight_ptr_0_iter += w13_n_tile_stride;
358 + w13_weight_ptr_1_iter += w13_n_tile_stride;
359 + w13_scale_ptr_0_iter += gemm_n_tile_size;
360 + w13_scale_ptr_1_iter += gemm_n_tile_size;
361 + w13_output_buffer_0_iter += gemm_n_tile_size;
362 + w13_output_buffer_1_iter += gemm_n_tile_size;
363 + }
364 +
365 + apply_gated_act(act_type, gemm_output_buffer,
366 + curr_w13_gemm_output_buffer, actual_token_num,
367 + actual_n_tile_size, actual_n_tile_size, input_size_2);
368 + curr_w13_gemm_output_buffer += gemm_m_tile_size * input_size_2;
369 + }
370 + }
371 + }
372 + }
373 +
374 + auto* __restrict__ w13_gemm_output_buffer = reinterpret_cast<scalar_t*>(
375 + common_buffer_start + w13_gemm_output_buffer_offset);
376 + float* __restrict__ w13_output_scale_buffer = reinterpret_cast<float*>(
377 + common_buffer_start + w13_output_scale_buffer_offset);
378 +
379 +// quantize w2 inputs - in place
380 +#pragma omp parallel for
381 + for (int32_t token_idx = 0; token_idx < expanded_token_num; ++token_idx) {
382 + scalar_t* input_row = w13_gemm_output_buffer + token_idx * input_size_2;
383 + int8_t* output_row = reinterpret_cast<int8_t*>(input_row);
384 + gemm_t::quantize_row(input_row, output_row,
385 + w13_output_scale_buffer[token_idx], input_size_2);
386 + }
387 +
388 + {
389 + alignas(64) cpu_utils::Counter counter;
390 + cpu_utils::Counter* counter_ptr = &counter;
391 +
392 +// w2 gemm
393 +#pragma omp parallel for schedule(static, 1)
394 + for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) {
395 + const int32_t task_num_per_expert =
396 + (output_size_2 + w2_n_tile_size - 1) / w2_n_tile_size;
397 + const int32_t task_num = task_num_per_expert * expert_num;
398 + uint8_t* __restrict__ thread_buffer =
399 + thread_buffer_start + thread_id * thread_buffer_size;
400 + int8_t* __restrict__ gemm_input_buffer =
401 + reinterpret_cast<int8_t*>(thread_buffer + gemm_input_buffer_offset);
402 + float* __restrict__ gemm_output_buffer =
403 + reinterpret_cast<float*>(thread_buffer + gemm_output_buffer_offset);
404 + float* __restrict__ w2_gemm_output_buffer = reinterpret_cast<float*>(
405 + common_buffer_start + w2_gemm_output_buffer_offset);
406 + gemm_t gemm;
407 +
408 + const int32_t w2_n_group_stride =
409 + gemm_t::WeightOCGroupSize * input_size_2;
410 + const int32_t w2_n_tile_stride = gemm_n_tile_size * input_size_2;
411 +
412 + for (;;) {
413 + const int32_t task_id = counter_ptr->acquire_counter();
414 + if (task_id >= task_num) {
415 + break;
416 + }
417 + const int32_t curr_expert_id = task_id / task_num_per_expert;
418 + const int32_t curr_output_group_id = task_id % task_num_per_expert;
419 + const int32_t curr_token_num =
420 + token_num_per_group_buffer[curr_expert_id];
421 + if (curr_token_num == 0) {
422 + continue;
423 + }
424 +
425 + const int32_t actual_n_tile_size =
426 + std::min(w2_n_tile_size,
427 + output_size_2 - curr_output_group_id * w2_n_tile_size);
428 + scalar_t* __restrict__ curr_w13_gemm_output_buffer =
429 + w13_gemm_output_buffer +
430 + cu_token_num_per_group_buffer[curr_expert_id] * input_size_2;
431 + float* __restrict__ curr_w13_output_scale_buffer =
432 + w13_output_scale_buffer +
433 + cu_token_num_per_group_buffer[curr_expert_id];
434 + float* __restrict__ curr_w2_gemm_output_buffer =
435 + w2_gemm_output_buffer +
436 + cu_token_num_per_group_buffer[curr_expert_id] * output_size_2 +
437 + curr_output_group_id * w2_n_tile_size;
438 + const int8_t* __restrict__ w2_weight_ptr =
439 + w2 + curr_expert_id * output_size_2 * input_size_2 +
440 + curr_output_group_id * w2_n_tile_size * input_size_2;
441 + const float* __restrict__ w2_scale_ptr =
442 + w2_scales + curr_expert_id * output_size_2 +
443 + curr_output_group_id * w2_n_tile_size;
444 + scalar_t* w2_bias_ptr = nullptr;
445 + if (w2_bias != nullptr) {
446 + w2_bias_ptr = w2_bias + curr_expert_id * output_size_2 +
447 + curr_output_group_id * w2_n_tile_size;
448 + }
449 +
450 + for (int32_t token_idx = 0; token_idx < curr_token_num;
451 + token_idx += gemm_m_tile_size) {
452 + const int32_t actual_token_num =
453 + std::min(gemm_m_tile_size, curr_token_num - token_idx);
454 + const int8_t* input_rows[gemm_m_tile_size];
455 + alignas(64) float input_scales[gemm_m_tile_size];
456 + for (int32_t i = 0; i < actual_token_num; ++i) {
457 + input_rows[i] = reinterpret_cast<const int8_t*>(
458 + curr_w13_gemm_output_buffer + i * input_size_2);
459 + input_scales[i] = curr_w13_output_scale_buffer[i];
460 + }
461 + gemm_t::pack_input_from_rows(input_rows, gemm_input_buffer,
462 + actual_token_num, input_size_2);
463 +
464 + const int8_t* w2_weight_ptr_iter = w2_weight_ptr;
465 + const float* w2_scale_ptr_iter = w2_scale_ptr;
466 + scalar_t* w2_bias_ptr_iter = w2_bias_ptr;
467 + float* curr_w2_gemm_output_buffer_iter = curr_w2_gemm_output_buffer;
468 + for (int32_t i = 0; i < actual_n_tile_size; i += gemm_n_tile_size) {
469 + auto* output_int32 = reinterpret_cast<int32_t*>(gemm_output_buffer);
470 + gemm.gemm(gemm_input_buffer, w2_weight_ptr_iter, output_int32,
471 + actual_token_num, input_size_2, w2_n_group_stride,
472 + gemm_n_tile_size);
473 + gemm_t::dequantize_tile(output_int32, gemm_output_buffer,
474 + input_scales, w2_scale_ptr_iter,
475 + actual_token_num, gemm_n_tile_size,
476 + gemm_n_tile_size);
477 + if (w2_bias != nullptr) {
478 + cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>(
479 + gemm_output_buffer, gemm_output_buffer, w2_bias_ptr_iter,
480 + actual_token_num, gemm_n_tile_size, gemm_n_tile_size);
481 + w2_bias_ptr_iter += gemm_n_tile_size;
482 + }
483 + for (int32_t m_idx = 0; m_idx < actual_token_num; ++m_idx) {
484 + std::memcpy(
485 + curr_w2_gemm_output_buffer_iter + m_idx * output_size_2,
486 + gemm_output_buffer + m_idx * gemm_n_tile_size,
487 + gemm_n_tile_size * sizeof(float));
488 + }
489 +
490 + w2_weight_ptr_iter += w2_n_tile_stride;
491 + w2_scale_ptr_iter += gemm_n_tile_size;
492 + curr_w2_gemm_output_buffer_iter += gemm_n_tile_size;
493 + }
494 +
495 + curr_w13_gemm_output_buffer += gemm_m_tile_size * input_size_2;
496 + curr_w13_output_scale_buffer += gemm_m_tile_size;
497 + curr_w2_gemm_output_buffer += gemm_m_tile_size * output_size_2;
498 + }
499 + }
500 + }
501 + }
502 +
503 + {
504 + alignas(64) cpu_utils::Counter counter;
505 + cpu_utils::Counter* counter_ptr = &counter;
506 +
507 +#pragma omp parallel for schedule(static, 1)
508 + for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) {
509 + uint8_t* __restrict__ thread_buffer =
510 + thread_buffer_start + thread_id * thread_buffer_size;
511 + float* __restrict__ ws_output_buffer =
512 + reinterpret_cast<float*>(thread_buffer + ws_output_buffer_offset);
513 + float* __restrict__ w2_gemm_output_buffer = reinterpret_cast<float*>(
514 + common_buffer_start + w2_gemm_output_buffer_offset);
515 +
516 + for (;;) {
517 + const int32_t token_id = counter_ptr->acquire_counter();
518 + if (token_id >= token_num) {
519 + break;
520 + }
521 + int32_t* __restrict__ curr_expand_token_id_index_buffer =
522 + expand_token_id_index_buffer + token_id * topk_num;
523 + const float* __restrict__ curr_weight =
524 + topk_weights + token_id * topk_num;
525 + const float first_weight = skip_weighted ? 1.0f : curr_weight[0];
526 + scalar_t* __restrict__ curr_output_buffer =
527 + output + token_id * output_size_2;
528 +
529 + if (topk_num > 1) {
530 + int32_t w2_output_idx = curr_expand_token_id_index_buffer[0];
531 + float* w2_output_iter =
532 + w2_gemm_output_buffer + w2_output_idx * output_size_2;
533 + float* ws_output_buffer_iter = ws_output_buffer;
534 + vec_op::FP32Vec16 weight_vec(first_weight);
535 + for (int32_t i = 0; i < output_size_2; i += 16) {
536 + vec_op::FP32Vec16 vec(w2_output_iter);
537 + (vec * weight_vec).save(ws_output_buffer_iter);
538 + w2_output_iter += 16;
539 + ws_output_buffer_iter += 16;
540 + }
541 +
542 + for (int32_t idx = 1; idx < topk_num - 1; ++idx) {
543 + w2_output_idx = curr_expand_token_id_index_buffer[idx];
544 + w2_output_iter =
545 + w2_gemm_output_buffer + w2_output_idx * output_size_2;
546 + ws_output_buffer_iter = ws_output_buffer;
547 + weight_vec = vec_op::FP32Vec16(curr_weight[idx]);
548 + for (int32_t i = 0; i < output_size_2; i += 16) {
549 + vec_op::FP32Vec16 vec(w2_output_iter);
550 + vec_op::FP32Vec16 sum(ws_output_buffer_iter);
551 + (sum + vec * weight_vec).save(ws_output_buffer_iter);
552 + w2_output_iter += 16;
553 + ws_output_buffer_iter += 16;
554 + }
555 + }
556 +
557 + const int32_t last_idx = topk_num - 1;
558 + w2_output_idx = curr_expand_token_id_index_buffer[last_idx];
559 + w2_output_iter =
560 + w2_gemm_output_buffer + w2_output_idx * output_size_2;
561 + ws_output_buffer_iter = ws_output_buffer;
562 + scalar_t* curr_output_buffer_iter = curr_output_buffer;
563 + weight_vec = vec_op::FP32Vec16(curr_weight[last_idx]);
564 + for (int32_t i = 0; i < output_size_2; i += 16) {
565 + vec_op::FP32Vec16 vec(w2_output_iter);
566 + vec_op::FP32Vec16 sum(ws_output_buffer_iter);
567 + scalar_vec_t(sum + vec * weight_vec).save(curr_output_buffer_iter);
568 + w2_output_iter += 16;
569 + ws_output_buffer_iter += 16;
570 + curr_output_buffer_iter += 16;
571 + }
572 + } else {
573 + const int32_t w2_output_idx = curr_expand_token_id_index_buffer[0];
574 + float* w2_output_iter =
575 + w2_gemm_output_buffer + w2_output_idx * output_size_2;
576 + scalar_t* curr_output_buffer_iter = curr_output_buffer;
577 + vec_op::FP32Vec16 weight_vec(first_weight);
578 + for (int32_t i = 0; i < output_size_2; i += 16) {
579 + vec_op::FP32Vec16 vec(w2_output_iter);
580 + scalar_vec_t(vec * weight_vec).save(curr_output_buffer_iter);
581 + w2_output_iter += 16;
582 + curr_output_buffer_iter += 16;
583 + }
584 + }
585 + }
586 + }
587 + }
588 +}
589 +} // namespace
590 +
591 +void prepack_moe_weight_int8(
592 + const torch::Tensor& weight, // [expert_num, output_size, input_size]
593 + torch::Tensor& packed_weight, const std::string& isa) {
594 + TORCH_CHECK(weight.is_contiguous());
595 + const int32_t expert_num = weight.size(0);
596 + const int32_t output_size = weight.size(1);
597 + const int32_t input_size = weight.size(2);
598 + const int64_t expert_stride = weight.stride(0);
599 + const cpu_utils::ISA isa_type = cpu_utils::get_isa(isa);
600 + TORCH_CHECK_EQ(output_size % 32, 0);
601 +
602 + CPU_INT8_ISA_DISPATCH_IMPL(isa_type, c10::BFloat16, [&]() {
603 + TORCH_CHECK_EQ(input_size % gemm_t::K, 0);
604 + prepack_moe_weight_int8_impl<gemm_t>(
605 + weight.data_ptr<int8_t>(), packed_weight.data_ptr<int8_t>(), expert_num,
606 + output_size, input_size, expert_stride);
607 + });
608 +}
609 +
610 +void cpu_fused_moe_int8(torch::Tensor& output, const torch::Tensor& input,
611 + const torch::Tensor& w13, const torch::Tensor& w2,
612 + const torch::Tensor& w13_scale,
613 + const torch::Tensor& w2_scale,
614 + const std::optional<torch::Tensor>& w13_bias,
615 + const std::optional<torch::Tensor>& w2_bias,
616 + const torch::Tensor& topk_weights,
617 + const torch::Tensor& topk_id, const bool skip_weighted,
618 + const std::string& act, const std::string& isa) {
619 + const int32_t token_num = input.size(0);
620 + const int32_t input_size_13 = input.size(1);
621 + const int64_t input_stride = input.stride(0);
622 + TORCH_CHECK_EQ(input_stride, input_size_13);
623 + const int32_t expert_num = w13.size(0);
624 + const int32_t output_size_13 = w13.size(1);
625 + const int32_t input_size_2 = w2.size(2);
626 + const int32_t output_size_2 = w2.size(1);
627 + const int32_t topk_num = topk_id.size(1);
628 + const FusedMOEAct act_type = cpu_fused_moe_utils::get_act_type(act);
629 + const cpu_utils::ISA isa_type = cpu_utils::get_isa(isa);
630 + TORCH_CHECK(!skip_weighted || topk_num == 1,
631 + "skip_weighted is only supported for topk=1 on CPU");
632 +
633 + VLLM_DISPATCH_FLOATING_TYPES(
634 + input.scalar_type(), "cpu_fused_moe_int8", [&]() {
635 + CPU_INT8_ISA_DISPATCH_IMPL(isa_type, scalar_t, [&]() {
636 + fused_moe_int8_impl<scalar_t, gemm_t>(
637 + output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
638 + w13.data_ptr<int8_t>(), w2.data_ptr<int8_t>(),
639 + w13_scale.data_ptr<float>(), w2_scale.data_ptr<float>(),
640 + w13_bias.has_value() ? w13_bias->data_ptr<scalar_t>() : nullptr,
641 + w2_bias.has_value() ? w2_bias->data_ptr<scalar_t>() : nullptr,
642 + topk_weights.data_ptr<float>(), topk_id.data_ptr<int32_t>(),
643 + act_type, token_num, expert_num, topk_num, input_size_13,
644 + output_size_13, input_size_2, output_size_2, skip_weighted);
645 + });
646 + });
647 +}
csrc/cpu/micro_gemm/cpu_micro_gemm_impl.hpp
+3
@@ -31,6 +31,9 @@ class MicroGemm {
31 }
32 };
33
34 +template <cpu_utils::ISA isa, typename scalar_t>
35 +class MicroGemmINT8;
36 +
37 template <int32_t n_size, typename scalar_t>
38 FORCE_INLINE void default_epilogue(float* __restrict__ c_ptr,
39 scalar_t* __restrict__ d_ptr,
csrc/cpu/micro_gemm/cpu_micro_gemm_int8_neon.hpp new
+424
@@ -0,0 +1,424 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +#ifndef CPU_MICRO_GEMM_INT8_NEON_HPP
5 +#define CPU_MICRO_GEMM_INT8_NEON_HPP
6 +
7 +#include <algorithm>
8 +#include <cstdint>
9 +
10 +#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
11 +
12 +#include <arm_bf16.h>
13 +#include <arm_neon.h>
14 +#include <c10/util/BFloat16.h>
15 +#include <c10/util/Exception.h>
16 +#include <c10/util/Half.h>
17 +
18 +namespace cpu_micro_gemm {
19 +
20 +namespace neon_smmla {
21 +
22 +constexpr int32_t K = 8;
23 +constexpr int32_t Cols = 2;
24 +constexpr int32_t TileSize = K * Cols;
25 +
26 +FORCE_INLINE float32x4x2_t load_as_f32(const float* input) {
27 + float32x4x2_t result;
28 + result.val[0] = vld1q_f32(input);
29 + result.val[1] = vld1q_f32(input + 4);
30 + return result;
31 +}
32 +
33 +FORCE_INLINE float32x4x2_t load_as_f32(const c10::Half* input) {
34 + const auto input_vec = vld1q_f16(reinterpret_cast<const float16_t*>(input));
35 + float32x4x2_t result;
36 + result.val[0] = vcvt_f32_f16(vget_low_f16(input_vec));
37 + result.val[1] = vcvt_f32_f16(vget_high_f16(input_vec));
38 + return result;
39 +}
40 +
41 +FORCE_INLINE float32x4x2_t load_as_f32(const c10::BFloat16* input) {
42 + const auto input_vec = vld1q_bf16(reinterpret_cast<const bfloat16_t*>(input));
43 + float32x4x2_t result;
44 + result.val[0] = vcvt_f32_bf16(vget_low_bf16(input_vec));
45 + result.val[1] = vcvt_f32_bf16(vget_high_bf16(input_vec));
46 + return result;
47 +}
48 +
49 +FORCE_INLINE void store_acc_rowpair(const int32x4_t acc01,
50 + const int32x4_t acc23,
51 + const int32x4_t acc45,
52 + const int32x4_t acc67,
53 + int32_t* __restrict__ c_ptr,
54 + const int64_t ldc, const int32_t m_rows) {
55 + if (m_rows == 0) {
56 + return;
57 + }
58 +
59 + vst1q_s32(c_ptr, vcombine_s32(vget_low_s32(acc01), vget_low_s32(acc23)));
60 + vst1q_s32(c_ptr + 4, vcombine_s32(vget_low_s32(acc45), vget_low_s32(acc67)));
61 +
62 + if (m_rows == 2) {
63 + vst1q_s32(c_ptr + ldc,
64 + vcombine_s32(vget_high_s32(acc01), vget_high_s32(acc23)));
65 + vst1q_s32(c_ptr + ldc + 4,
66 + vcombine_s32(vget_high_s32(acc45), vget_high_s32(acc67)));
67 + }
68 +}
69 +
70 +FORCE_INLINE void gemm_micro_smmla_8x8_packed_a(
71 + const int8_t* __restrict__ a_packed, const int8_t* __restrict__ b_packed,
72 + int32_t* __restrict__ c_ptr, const int32_t m, const int32_t k_size,
73 + const int64_t ldc) {
74 + const int32x4_t zero = vdupq_n_s32(0);
75 + int32x4_t acc0101 = zero, acc0123 = zero, acc0145 = zero, acc0167 = zero;
76 + int32x4_t acc2301 = zero, acc2323 = zero, acc2345 = zero, acc2367 = zero;
77 + int32x4_t acc4501 = zero, acc4523 = zero, acc4545 = zero, acc4567 = zero;
78 + int32x4_t acc6701 = zero, acc6723 = zero, acc6745 = zero, acc6767 = zero;
79 +
80 + const int8_t* __restrict__ a_tile = a_packed;
81 + const int8_t* __restrict__ b_tile = b_packed;
82 +
83 +#pragma GCC unroll 8
84 + for (int32_t k_idx = 0; k_idx < k_size; k_idx += K) {
85 + const int8x16_t a_tile01 = vld1q_s8(a_tile);
86 + const int8x16_t a_tile23 = vld1q_s8(a_tile + TileSize);
87 + const int8x16_t a_tile45 = vld1q_s8(a_tile + 2 * TileSize);
88 + const int8x16_t a_tile67 = vld1q_s8(a_tile + 3 * TileSize);
89 +
90 + const int8x16_t b_tile01 = vld1q_s8(b_tile);
91 + const int8x16_t b_tile23 = vld1q_s8(b_tile + TileSize);
92 + const int8x16_t b_tile45 = vld1q_s8(b_tile + 2 * TileSize);
93 + const int8x16_t b_tile67 = vld1q_s8(b_tile + 3 * TileSize);
94 +
95 + acc0101 = vmmlaq_s32(acc0101, a_tile01, b_tile01);
96 + acc2301 = vmmlaq_s32(acc2301, a_tile23, b_tile01);
97 + acc4501 = vmmlaq_s32(acc4501, a_tile45, b_tile01);
98 + acc6701 = vmmlaq_s32(acc6701, a_tile67, b_tile01);
99 +
100 + acc0123 = vmmlaq_s32(acc0123, a_tile01, b_tile23);
101 + acc2323 = vmmlaq_s32(acc2323, a_tile23, b_tile23);
102 + acc4523 = vmmlaq_s32(acc4523, a_tile45, b_tile23);
103 + acc6723 = vmmlaq_s32(acc6723, a_tile67, b_tile23);
104 +
105 + acc0145 = vmmlaq_s32(acc0145, a_tile01, b_tile45);
106 + acc2345 = vmmlaq_s32(acc2345, a_tile23, b_tile45);
107 + acc4545 = vmmlaq_s32(acc4545, a_tile45, b_tile45);
108 + acc6745 = vmmlaq_s32(acc6745, a_tile67, b_tile45);
109 +
110 + acc0167 = vmmlaq_s32(acc0167, a_tile01, b_tile67);
111 + acc2367 = vmmlaq_s32(acc2367, a_tile23, b_tile67);
112 + acc4567 = vmmlaq_s32(acc4567, a_tile45, b_tile67);
113 + acc6767 = vmmlaq_s32(acc6767, a_tile67, b_tile67);
114 +
115 + a_tile += 4 * TileSize;
116 + b_tile += 4 * TileSize;
117 + }
118 +
119 + store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc,
120 + std::min(2, m));
121 + store_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc,
122 + std::min(2, std::max(0, m - 2)));
123 + store_acc_rowpair(acc4501, acc4523, acc4545, acc4567, c_ptr + 4 * ldc, ldc,
124 + std::min(2, std::max(0, m - 4)));
125 + store_acc_rowpair(acc6701, acc6723, acc6745, acc6767, c_ptr + 6 * ldc, ldc,
126 + std::min(2, std::max(0, m - 6)));
127 +}
128 +
129 +FORCE_INLINE void gemm_micro_smmla_4x16_packed_a(
130 + const int8_t* __restrict__ a_packed, const int8_t* __restrict__ b_packed,
131 + int32_t* __restrict__ c_ptr, const int32_t m, const int32_t k_size,
132 + const int64_t b_n_group_stride, const int64_t ldc) {
133 + const int32_t m_rows_01 = std::min(2, m);
134 + const int32_t m_rows_23 = std::min(2, std::max(0, m - 2));
135 + const int32x4_t zero = vdupq_n_s32(0);
136 +
137 + int32x4_t acc0101 = zero, acc0123 = zero, acc0145 = zero, acc0167 = zero;
138 + int32x4_t acc2301 = zero, acc2323 = zero, acc2345 = zero, acc2367 = zero;
139 + int32x4_t acc0189 = zero, acc011011 = zero, acc011213 = zero,
140 + acc011415 = zero;
141 + int32x4_t acc2389 = zero, acc231011 = zero, acc231213 = zero,
142 + acc231415 = zero;
143 +
144 + const int8_t* __restrict__ a_tile = a_packed;
145 + // note: b packs 8 panels contiguously, so we need 2 b_tile ptrs
146 + // for the 4x16 microkernel
147 + const int8_t* __restrict__ b_tile0 = b_packed;
148 + const int8_t* __restrict__ b_tile1 = b_packed + b_n_group_stride;
149 +
150 +#pragma GCC unroll 8
151 + for (int32_t k_idx = 0; k_idx < k_size; k_idx += K) {
152 + const int8x16_t a_tile01 = vld1q_s8(a_tile);
153 + const int8x16_t a_tile23 = vld1q_s8(a_tile + TileSize);
154 + const int8x16_t b_tile01 = vld1q_s8(b_tile0);
155 + const int8x16_t b_tile23 = vld1q_s8(b_tile0 + TileSize);
156 + const int8x16_t b_tile45 = vld1q_s8(b_tile0 + 2 * TileSize);
157 + const int8x16_t b_tile67 = vld1q_s8(b_tile0 + 3 * TileSize);
158 + const int8x16_t b_tile89 = vld1q_s8(b_tile1);
159 + const int8x16_t b_tile1011 = vld1q_s8(b_tile1 + TileSize);
160 + const int8x16_t b_tile1213 = vld1q_s8(b_tile1 + 2 * TileSize);
161 + const int8x16_t b_tile1415 = vld1q_s8(b_tile1 + 3 * TileSize);
162 +
163 + acc0101 = vmmlaq_s32(acc0101, a_tile01, b_tile01);
164 + acc2301 = vmmlaq_s32(acc2301, a_tile23, b_tile01);
165 + acc0123 = vmmlaq_s32(acc0123, a_tile01, b_tile23);
166 + acc2323 = vmmlaq_s32(acc2323, a_tile23, b_tile23);
167 +
168 + acc0145 = vmmlaq_s32(acc0145, a_tile01, b_tile45);
169 + acc2345 = vmmlaq_s32(acc2345, a_tile23, b_tile45);
170 + acc0167 = vmmlaq_s32(acc0167, a_tile01, b_tile67);
171 + acc2367 = vmmlaq_s32(acc2367, a_tile23, b_tile67);
172 +
173 + acc0189 = vmmlaq_s32(acc0189, a_tile01, b_tile89);
174 + acc2389 = vmmlaq_s32(acc2389, a_tile23, b_tile89);
175 + acc011011 = vmmlaq_s32(acc011011, a_tile01, b_tile1011);
176 + acc231011 = vmmlaq_s32(acc231011, a_tile23, b_tile1011);
177 +
178 + acc011213 = vmmlaq_s32(acc011213, a_tile01, b_tile1213);
179 + acc231213 = vmmlaq_s32(acc231213, a_tile23, b_tile1213);
180 + acc011415 = vmmlaq_s32(acc011415, a_tile01, b_tile1415);
181 + acc231415 = vmmlaq_s32(acc231415, a_tile23, b_tile1415);
182 +
183 + a_tile += 2 * TileSize;
184 + b_tile0 += 4 * TileSize;
185 + b_tile1 += 4 * TileSize;
186 + }
187 +
188 + // rows 0-1, columns 0-7
189 + store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, m_rows_01);
190 + // rows 0-1, columns 8-15
191 + store_acc_rowpair(acc0189, acc011011, acc011213, acc011415, c_ptr + 8, ldc,
192 + m_rows_01);
193 + // rows 2-3, columns 0-7
194 + store_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc,
195 + m_rows_23);
196 + // rows 2-3, columns 8-15
197 + store_acc_rowpair(acc2389, acc231011, acc231213, acc231415,
198 + c_ptr + 2 * ldc + 8, ldc, m_rows_23);
199 +}
200 +
201 +} // namespace neon_smmla
202 +
203 +template <typename scalar_t>
204 +class MicroGemmINT8<cpu_utils::ISA::NEON, scalar_t> {
205 + public:
206 + static constexpr int32_t K = neon_smmla::K;
207 + static constexpr int32_t Mr = 8;
208 + static constexpr int32_t Nr = 8;
209 + static constexpr int32_t NrGemv = 16;
210 + static constexpr int32_t MaxMSize = 8;
211 + static constexpr int32_t NSize = 32;
212 + static constexpr int32_t WeightOCGroupSize = Nr;
213 + static_assert(MaxMSize % Mr == 0);
214 +
215 + static FORCE_INLINE void quantize_row(const scalar_t* input, int8_t* output,
216 + float& scale, const int32_t size) {
217 + TORCH_CHECK_EQ(size % K, 0);
218 + float32x4_t max_vec = vdupq_n_f32(0.0f);
219 +
220 + for (int32_t i = 0; i < size; i += K) {
221 + const float32x4x2_t input_vec = neon_smmla::load_as_f32(input + i);
222 + max_vec = vmaxq_f32(max_vec, vabsq_f32(input_vec.val[0]));
223 + max_vec = vmaxq_f32(max_vec, vabsq_f32(input_vec.val[1]));
224 + }
225 +
226 + const float abs_max = std::max(vmaxvq_f32(max_vec), 1.0e-7f);
227 + scale = abs_max / 127.0f;
228 + const float32x4_t inv_scale_vec = vdupq_n_f32(127.0f / abs_max);
229 +
230 + for (int32_t i = 0; i < size; i += K) {
231 + const float32x4x2_t input_vec = neon_smmla::load_as_f32(input + i);
232 + const int32x4_t output_low =
233 + vcvtnq_s32_f32(vmulq_f32(input_vec.val[0], inv_scale_vec));
234 + const int32x4_t output_high =
235 + vcvtnq_s32_f32(vmulq_f32(input_vec.val[1], inv_scale_vec));
236 + const int16x8_t output_s16 =
237 + vcombine_s16(vqmovn_s32(output_low), vqmovn_s32(output_high));
238 + vst1_s8(output + i, vqmovn_s16(output_s16));
239 + }
240 + }
241 +
242 + // with current code, fusing this into the gemm micro kernel didn't move the
243 + // needle
244 + static FORCE_INLINE void dequantize_tile(
245 + int32_t* input, float* output, const float* __restrict__ input_scales,
246 + const float* __restrict__ weight_scales, const int32_t m, const int32_t n,
247 + const int32_t stride) {
248 + TORCH_CHECK_EQ(n % 4, 0);
249 + for (int32_t m_idx = 0; m_idx < m; ++m_idx) {
250 + const float32x4_t input_scale_vec = vdupq_n_f32(input_scales[m_idx]);
251 + for (int32_t n_idx = 0; n_idx < n; n_idx += 4) {
252 + const int32x4_t input_vec = vld1q_s32(input + m_idx * stride + n_idx);
253 + const float32x4_t weight_scale_vec = vld1q_f32(weight_scales + n_idx);
254 + const float32x4_t output_vec =
255 + vmulq_f32(vcvtq_f32_s32(input_vec),
256 + vmulq_f32(input_scale_vec, weight_scale_vec));
257 + vst1q_f32(output + m_idx * stride + n_idx, output_vec);
258 + }
259 + }
260 + }
261 +
262 + // physical layout [
263 + // M / (8 or 4); Mr is 8 or 4
264 + // K / 8; K for smmla is 8
265 + // 4, ; 4 row-pairs for each 8 rows
266 + // 2, ; row-pair is 2 rows
267 + // 4 ; 4 elements per row
268 + // ]
269 + static void pack_input_from_rows(const int8_t* const* __restrict__ rows,
270 + int8_t* __restrict__ a_packed,
271 + const int32_t m, const int32_t k) {
272 + TORCH_CHECK(m > 0 && m <= MaxMSize);
273 + TORCH_CHECK(k % K == 0);
274 + const int8x8_t zero = vdup_n_s8(0);
275 +
276 + for (int32_t row_base = 0; row_base < m; row_base += Mr) {
277 + const int32_t panel_m = std::min(Mr, m - row_base);
278 + const int8_t* const* panel_rows = rows + row_base;
279 + int8_t* __restrict__ out = a_packed + row_base * k;
280 +
281 + // fast path for full 8-row panels (fast path for 4-row panels didn't move
282 + // the needle)
283 + if (panel_m == Mr) {
284 + const int8_t* __restrict__ row0 = panel_rows[0];
285 + const int8_t* __restrict__ row1 = panel_rows[1];
286 + const int8_t* __restrict__ row2 = panel_rows[2];
287 + const int8_t* __restrict__ row3 = panel_rows[3];
288 + const int8_t* __restrict__ row4 = panel_rows[4];
289 + const int8_t* __restrict__ row5 = panel_rows[5];
290 + const int8_t* __restrict__ row6 = panel_rows[6];
291 + const int8_t* __restrict__ row7 = panel_rows[7];
292 + int32_t k_idx = 0;
293 + for (; k_idx + 2 * K <= k; k_idx += 2 * K) {
294 + int8_t* __restrict__ block0 = out;
295 + int8_t* __restrict__ block1 = out + 4 * neon_smmla::TileSize;
296 +
297 + int8x16_t a0 = vld1q_s8(row0 + k_idx);
298 + int8x16_t a1 = vld1q_s8(row1 + k_idx);
299 + vst1q_s8(block0, vcombine_s8(vget_low_s8(a0), vget_low_s8(a1)));
300 + vst1q_s8(block1, vcombine_s8(vget_high_s8(a0), vget_high_s8(a1)));
301 +
302 + a0 = vld1q_s8(row2 + k_idx);
303 + a1 = vld1q_s8(row3 + k_idx);
304 + vst1q_s8(block0 + neon_smmla::TileSize,
305 + vcombine_s8(vget_low_s8(a0), vget_low_s8(a1)));
306 + vst1q_s8(block1 + neon_smmla::TileSize,
307 + vcombine_s8(vget_high_s8(a0), vget_high_s8(a1)));
308 +
309 + a0 = vld1q_s8(row4 + k_idx);
310 + a1 = vld1q_s8(row5 + k_idx);
311 + vst1q_s8(block0 + 2 * neon_smmla::TileSize,
312 + vcombine_s8(vget_low_s8(a0), vget_low_s8(a1)));
313 + vst1q_s8(block1 + 2 * neon_smmla::TileSize,
314 + vcombine_s8(vget_high_s8(a0), vget_high_s8(a1)));
315 +
316 + a0 = vld1q_s8(row6 + k_idx);
317 + a1 = vld1q_s8(row7 + k_idx);
318 + vst1q_s8(block0 + 3 * neon_smmla::TileSize,
319 + vcombine_s8(vget_low_s8(a0), vget_low_s8(a1)));
320 + vst1q_s8(block1 + 3 * neon_smmla::TileSize,
321 + vcombine_s8(vget_high_s8(a0), vget_high_s8(a1)));
322 +
323 + out += 8 * neon_smmla::TileSize;
324 + }
325 +
326 + for (; k_idx < k; k_idx += K) {
327 + int8x8_t a0 = vld1_s8(row0 + k_idx);
328 + int8x8_t a1 = vld1_s8(row1 + k_idx);
329 + vst1q_s8(out, vcombine_s8(a0, a1));
330 +
331 + a0 = vld1_s8(row2 + k_idx);
332 + a1 = vld1_s8(row3 + k_idx);
333 + vst1q_s8(out + neon_smmla::TileSize, vcombine_s8(a0, a1));
334 +
335 + a0 = vld1_s8(row4 + k_idx);
336 + a1 = vld1_s8(row5 + k_idx);
337 + vst1q_s8(out + 2 * neon_smmla::TileSize, vcombine_s8(a0, a1));
338 +
339 + a0 = vld1_s8(row6 + k_idx);
340 + a1 = vld1_s8(row7 + k_idx);
341 + vst1q_s8(out + 3 * neon_smmla::TileSize, vcombine_s8(a0, a1));
342 +
343 + out += 4 * neon_smmla::TileSize;
344 + }
345 + continue;
346 + }
347 +
348 + const int32_t row_pairs = (panel_m <= 4) ? 2 : Mr / 2;
349 + for (int32_t k_idx = 0; k_idx < k; k_idx += K) {
350 + for (int32_t pair_idx = 0; pair_idx < row_pairs; ++pair_idx) {
351 + const int32_t row_idx = pair_idx * 2;
352 + const int8x8_t row0 =
353 + (row_idx < panel_m) ? vld1_s8(panel_rows[row_idx] + k_idx) : zero;
354 + const int8x8_t row1 = (row_idx + 1 < panel_m)
355 + ? vld1_s8(panel_rows[row_idx + 1] + k_idx)
356 + : zero;
357 + vst1q_s8(out, vcombine_s8(row0, row1));
358 + out += neon_smmla::TileSize;
359 + }
360 + }
361 + }
362 + }
363 +
364 + // physical layout [
365 + // N / 8; Nr is 8
366 + // K / 8; K for smmla is 8
367 + // 4, ; 4 col-pairs for each 8 cols
368 + // 2, ; col-pair is 2 cols
369 + // 4 ; 4 elements per col
370 + // ]
371 + static void pack_weight(const int8_t* __restrict__ weight,
372 + int8_t* __restrict__ packed_weight,
373 + const int32_t output_size, const int32_t input_size) {
374 + TORCH_CHECK(output_size % NSize == 0);
375 + TORCH_CHECK(input_size % K == 0);
376 +
377 + for (int32_t o_idx = 0; o_idx < output_size; o_idx += Nr) {
378 + int8_t* __restrict__ dst = packed_weight + o_idx * input_size;
379 + for (int32_t k_idx = 0; k_idx < input_size; k_idx += K) {
380 + for (int32_t pair_idx = 0; pair_idx < Nr;
381 + pair_idx += neon_smmla::Cols) {
382 + const int8_t* __restrict__ row0 =
383 + weight + (o_idx + pair_idx) * input_size + k_idx;
384 + const int8_t* __restrict__ row1 = row0 + input_size;
385 + vst1q_s8(dst, vcombine_s8(vld1_s8(row0), vld1_s8(row1)));
386 + dst += neon_smmla::TileSize;
387 + }
388 + }
389 + }
390 + }
391 +
392 + void gemm(const int8_t* __restrict__ a_packed,
393 + const int8_t* __restrict__ b_packed, int32_t* __restrict__ c,
394 + const int32_t m, const int32_t k, const int64_t b_n_group_stride,
395 + const int64_t ldc) const {
396 + TORCH_CHECK(m > 0 && m <= MaxMSize);
397 + TORCH_CHECK(k % K == 0);
398 +
399 + for (int32_t n_idx = 0; n_idx < NSize; n_idx += NrGemv) {
400 + const int8_t* __restrict__ b_panel = b_packed + n_idx * k;
401 +
402 + for (int32_t row_base = 0; row_base < m; row_base += Mr) {
403 + const int32_t panel_m = std::min(Mr, m - row_base);
404 + const int8_t* __restrict__ a_panel = a_packed + row_base * k;
405 + int32_t* __restrict__ c_panel = c + row_base * ldc + n_idx;
406 +
407 + if (panel_m <= 4) {
408 + neon_smmla::gemm_micro_smmla_4x16_packed_a(
409 + a_panel, b_panel, c_panel, panel_m, k, b_n_group_stride, ldc);
410 + } else {
411 + neon_smmla::gemm_micro_smmla_8x8_packed_a(a_panel, b_panel, c_panel,
412 + panel_m, k, ldc);
413 + neon_smmla::gemm_micro_smmla_8x8_packed_a(
414 + a_panel, b_panel + b_n_group_stride, c_panel + Nr, panel_m, k,
415 + ldc);
416 + }
417 + }
418 + }
419 + }
420 +};
421 +
422 +} // namespace cpu_micro_gemm
423 +
424 +#endif
csrc/cpu/micro_gemm/cpu_micro_gemm_neon.hpp
+14 -8
@@ -1,3 +1,6 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 #ifndef CPU_MICRO_GEMM_NEON_HPP
5 #define CPU_MICRO_GEMM_NEON_HPP
6
@@ -16,9 +19,6 @@ namespace {
19 constexpr int32_t K = 4;
20 constexpr int32_t Cols = 2;
21 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) {
@@ -132,7 +132,7 @@ FORCE_INLINE void gemm_micro_bfmmla_8x8_packed_a(
132 acc6767 = vbfmmlaq_f32(acc6767, a_tile67, b_tile67);
133
134 a_tile += 4 * TileSize;
135 - b_tile += Nr * K;
135 + b_tile += 4 * TileSize;
136 }
137
138 store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc,
@@ -205,8 +205,8 @@ FORCE_INLINE void gemm_micro_bfmmla_4x16_packed_a(
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;
208 + b_tile0 += 4 * TileSize;
209 + b_tile1 += 4 * TileSize;
210 }
211
212 store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, m_rows_01);
@@ -223,6 +223,9 @@ FORCE_INLINE void gemm_micro_bfmmla_4x16_packed_a(
223 template <typename scalar_t>
224 class MicroGemm<cpu_utils::ISA::NEON, scalar_t> {
225 public:
226 + static constexpr int32_t Mr = 8;
227 + static constexpr int32_t Nr = 8;
228 + static constexpr int32_t NrGemv = 16;
229 static constexpr int32_t MaxMSize = 8;
230 static constexpr int32_t NSize = 32;
231 static constexpr int32_t WeightOCGroupSize = Nr;
@@ -246,6 +249,9 @@ class MicroGemm<cpu_utils::ISA::NEON, c10::BFloat16> {
249 public:
250 using scalar_t = c10::BFloat16;
251
252 + static constexpr int32_t Mr = 8;
253 + static constexpr int32_t Nr = 8;
254 + static constexpr int32_t NrGemv = 16;
255 static constexpr int32_t MaxMSize = 8;
256 static constexpr int32_t NSize = 32;
257 static constexpr int32_t WeightOCGroupSize = Nr;
@@ -253,7 +259,7 @@ class MicroGemm<cpu_utils::ISA::NEON, c10::BFloat16> {
259
260 public:
261 // physical layout [
256 - // M / 8; Mr is 8
262 + // M / (8 or 4); Mr is 8 or 4
263 // K / 4; K for bfmmla is 4
264 // 4, ; 4 row-pairs for each 8 rows
265 // 2, ; row-pair is 2 rows
@@ -439,7 +445,7 @@ class MicroGemm<cpu_utils::ISA::NEON, c10::BFloat16> {
445 (void)lda; // A is packed, so lda is not needed
446 TORCH_CHECK_EQ(k % K, 0);
447
442 - for (int32_t n_idx = 0; n_idx < NSize; n_idx += Nr_gemv) {
448 + for (int32_t n_idx = 0; n_idx < NSize; n_idx += NrGemv) {
449 const bfloat16_t* __restrict__ b_panel =
450 reinterpret_cast<const bfloat16_t*>(b_ptr) + n_idx * k;
451
csrc/cpu/torch_bindings.cpp
+31 -3
@@ -207,6 +207,20 @@ void cpu_fused_moe(torch::Tensor& output, const torch::Tensor& input,
207 const torch::Tensor& topk_id, const bool skip_weighted,
208 const std::string& act, const std::string& isa);
209
210 +void prepack_moe_weight_int8(const torch::Tensor& weight,
211 + torch::Tensor& packed_weight,
212 + const std::string& isa);
213 +
214 +void cpu_fused_moe_int8(torch::Tensor& output, const torch::Tensor& input,
215 + const torch::Tensor& w13, const torch::Tensor& w2,
216 + const torch::Tensor& w13_scale,
217 + const torch::Tensor& w2_scale,
218 + const std::optional<torch::Tensor>& w13_bias,
219 + const std::optional<torch::Tensor>& w2_bias,
220 + const torch::Tensor& topk_weights,
221 + const torch::Tensor& topk_id, const bool skip_weighted,
222 + const std::string& act, const std::string& isa);
223 +
224 void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
225 const torch::Tensor positions,
226 const torch::Tensor block_table,
@@ -596,8 +610,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
610 #endif
611
612 // fused moe
599 -#if defined(__AVX512F__) || \
600 - (defined(__aarch64__) && !defined(__APPLE__) && defined(ARM_BF16_SUPPORT))
613 +#if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT) && !defined(__APPLE__))
614 ops.def(
615 "prepack_moe_weight(Tensor weight, Tensor(a1!) packed_weight, str isa) "
616 "-> ()");
@@ -608,7 +621,22 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
621 "bool skip_weighted, "
622 "str act, str isa) -> ()");
623 ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe);
611 -#endif
624 +#endif // #if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT) &&
625 + // !defined(__APPLE__))
626 +#if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) && \
627 + !defined(__APPLE__)
628 + ops.def(
629 + "prepack_moe_weight_int8(Tensor weight, Tensor(a1!) packed_weight, "
630 + "str isa) -> ()");
631 + ops.impl("prepack_moe_weight_int8", torch::kCPU, &prepack_moe_weight_int8);
632 + ops.def(
633 + "cpu_fused_moe_int8(Tensor(a0!) output, Tensor input, Tensor w13, "
634 + "Tensor w2, Tensor w13_scale, Tensor w2_scale, Tensor? w13_bias, "
635 + "Tensor? w2_bias, Tensor topk_weights, Tensor topk_id, bool "
636 + "skip_weighted, str act, str isa) -> ()");
637 + ops.impl("cpu_fused_moe_int8", torch::kCPU, &cpu_fused_moe_int8);
638 +#endif // #if defined(ARM_I8MM_SUPPORT) && defined(ARM_BF16_SUPPORT) &&
639 + // !defined(__APPLE__)
640 ops.def(
641 "mla_decode_kvcache("
642 " Tensor! out, Tensor query, Tensor kv_cache,"
tests/kernels/moe/test_cpu_fused_moe.py
+171 -1
@@ -7,10 +7,14 @@ import torch
7 from tests.kernels.allclose_default import get_default_atol, get_default_rtol
8 from vllm._custom_ops import (
9 cpu_fused_moe,
10 + cpu_fused_moe_int8,
11 cpu_prepack_moe_weight,
12 + cpu_prepack_moe_weight_int8,
13 )
14 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
15 +from vllm.model_executor.layers.fused_moe.cpu_fused_moe import (
16 + _CPU_MOE_ACT_FN,
17 +)
18 from vllm.platforms import CpuArchEnum, current_platform
19 from vllm.utils.torch_utils import set_random_seed
20
@@ -101,6 +105,90 @@ def ref_fused_moe(
105 return final_out
106
107
108 +def quantize_per_channel(
109 + weight: torch.Tensor,
110 +) -> tuple[torch.Tensor, torch.Tensor]:
111 + """Symmetrically quantize each weight output channel."""
112 + weight_fp32 = weight.float()
113 + scale = weight_fp32.abs().amax(dim=-1).clamp_min(1e-12) / 127.0
114 + quantized = (
115 + (weight_fp32 / scale.unsqueeze(-1)).round().clamp(-127, 127).to(torch.int8)
116 + )
117 + return quantized, scale
118 +
119 +
120 +def quantize_per_token(
121 + input: torch.Tensor,
122 +) -> tuple[torch.Tensor, torch.Tensor]:
123 + """Symmetrically quantize each input token."""
124 + input_fp32 = input.float()
125 + scale = input_fp32.abs().amax(dim=-1, keepdim=True).clamp_min(1e-7) / 127.0
126 + quantized = (input_fp32 / scale).round().clamp(-127, 127).to(torch.int8)
127 + return quantized, scale
128 +
129 +
130 +def ref_fused_moe_int8(
131 + input: torch.Tensor,
132 + w13: torch.Tensor,
133 + w2: torch.Tensor,
134 + w13_scale: torch.Tensor,
135 + w2_scale: torch.Tensor,
136 + w13_bias: torch.Tensor | None,
137 + w2_bias: torch.Tensor | None,
138 + topk_weights: torch.Tensor,
139 + topk_ids: torch.Tensor,
140 + activation: MoEActivation,
141 +) -> torch.Tensor:
142 + """Reference the two dynamically quantized INT8 GEMMs."""
143 + input_int8, input_scale = quantize_per_token(input)
144 + expert_num = w13.size(0)
145 +
146 + counts = topk_ids.new_zeros((topk_ids.size(0), expert_num))
147 + counts.scatter_(1, topk_ids.to(torch.int64), 1)
148 + tokens_per_expert = counts.sum(dim=0).tolist()
149 + sorted_route_ids = topk_ids.view(-1).argsort()
150 + sorted_token_ids = sorted_route_ids // topk_ids.size(1)
151 +
152 + outputs = []
153 + start_idx = 0
154 + for expert_idx, token_count in enumerate(tokens_per_expert):
155 + end_idx = start_idx + token_count
156 + if token_count == 0:
157 + continue
158 +
159 + token_ids = sorted_token_ids[start_idx:end_idx]
160 + gate_up = torch.matmul(
161 + input_int8[token_ids].float(),
162 + w13[expert_idx].float().T,
163 + )
164 + gate_up *= input_scale[token_ids] * w13_scale[expert_idx]
165 + if w13_bias is not None:
166 + gate_up += w13_bias[expert_idx].float()
167 +
168 + intermediate = _CPU_MOE_ACT_FN[activation](gate_up).to(input.dtype)
169 + intermediate_int8, intermediate_scale = quantize_per_token(intermediate)
170 + output = torch.matmul(
171 + intermediate_int8.float(),
172 + w2[expert_idx].float().T,
173 + )
174 + output *= intermediate_scale * w2_scale[expert_idx]
175 + if w2_bias is not None:
176 + output += w2_bias[expert_idx].float()
177 +
178 + outputs.append(output)
179 + start_idx = end_idx
180 +
181 + sorted_output = torch.cat(outputs, dim=0)
182 + routed_output = torch.empty_like(sorted_output)
183 + routed_output[sorted_route_ids] = sorted_output
184 + return (
185 + routed_output.view(*topk_ids.shape, -1)
186 + .mul_(topk_weights.unsqueeze(-1))
187 + .sum(dim=1)
188 + .to(input.dtype)
189 + )
190 +
191 +
192 @pytest.mark.parametrize("batch_size", BATCH_SIZE)
193 @pytest.mark.parametrize("expert_num", EXPERT_NUM)
194 @pytest.mark.parametrize("hidden_size", HIDDEN_DIM)
@@ -176,3 +264,85 @@ def test_cpu_fused_moe(
264 torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol),
265 f"{torch.max(torch.abs(output - ref_output))}",
266 )
267 +
268 +
269 +@pytest.mark.skipif(
270 + current_platform.get_cpu_architecture() != CpuArchEnum.ARM,
271 + reason="Requires Arm CPU",
272 +)
273 +@pytest.mark.parametrize("batch_size", BATCH_SIZE)
274 +@pytest.mark.parametrize("expert_num", EXPERT_NUM)
275 +@pytest.mark.parametrize("hidden_size", HIDDEN_DIM)
276 +@pytest.mark.parametrize("intermediate_size", INTERMEDIATE_DIM)
277 +@pytest.mark.parametrize("use_bias", USE_BIAS)
278 +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
279 +@pytest.mark.parametrize("act", ACT)
280 +@pytest.mark.parametrize("isa", ["neon"])
281 +def test_cpu_fused_moe_int8(
282 + batch_size: int,
283 + expert_num: int,
284 + hidden_size: int,
285 + intermediate_size: int,
286 + use_bias: bool,
287 + dtype: torch.dtype,
288 + act: MoEActivation,
289 + isa: str,
290 +):
291 + set_random_seed(0)
292 + topk_num = max(expert_num // 2, 1)
293 + up_dim = 2 * intermediate_size
294 +
295 + input = torch.randn((batch_size, hidden_size), dtype=dtype) / (
296 + 0.5 * hidden_size**0.5
297 + )
298 + w13_fp = torch.randn((expert_num, up_dim, hidden_size), dtype=dtype) / (
299 + 0.5 * hidden_size**0.5
300 + )
301 + w2_fp = torch.randn((expert_num, hidden_size, intermediate_size), dtype=dtype) / (
302 + 0.5 * intermediate_size**0.5
303 + )
304 + w13, w13_scale = quantize_per_channel(w13_fp)
305 + w2, w2_scale = quantize_per_channel(w2_fp)
306 +
307 + w13_bias = None
308 + w2_bias = None
309 + if use_bias:
310 + w13_bias = torch.randn((expert_num, up_dim), dtype=dtype) / (0.5 * up_dim**0.5)
311 + w2_bias = torch.randn((expert_num, hidden_size), dtype=dtype) / (
312 + 0.5 * hidden_size**0.5
313 + )
314 +
315 + router_logits = torch.randn((batch_size, expert_num), dtype=dtype)
316 + score = torch.softmax(router_logits, dim=-1, dtype=torch.float32)
317 + topk_weights, topk_ids = torch.topk(score, topk_num)
318 + topk_ids = topk_ids.to(torch.int32)
319 +
320 + ref_output = ref_fused_moe_int8(
321 + input,
322 + w13,
323 + w2,
324 + w13_scale,
325 + w2_scale,
326 + w13_bias,
327 + w2_bias,
328 + topk_weights,
329 + topk_ids,
330 + act,
331 + )
332 + packed_w13 = cpu_prepack_moe_weight_int8(w13, isa)
333 + packed_w2 = cpu_prepack_moe_weight_int8(w2, isa)
334 + output = cpu_fused_moe_int8(
335 + input,
336 + packed_w13,
337 + packed_w2,
338 + w13_scale,
339 + w2_scale,
340 + w13_bias,
341 + w2_bias,
342 + topk_weights,
343 + topk_ids,
344 + act.value,
345 + isa,
346 + )
347 +
348 + torch.testing.assert_close(output, ref_output, atol=2e-2, rtol=2e-2)
vllm/_custom_ops.py
+42
@@ -3787,6 +3787,15 @@ def cpu_prepack_moe_weight(
3787 return output
3788
3789
3790 +def cpu_prepack_moe_weight_int8(
3791 + weight: torch.Tensor,
3792 + isa: str,
3793 +) -> torch.Tensor:
3794 + output = torch.empty_like(weight)
3795 + torch.ops._C.prepack_moe_weight_int8(weight, output, isa)
3796 + return output
3797 +
3798 +
3799 def cpu_fused_moe(
3800 input: torch.Tensor,
3801 w13: torch.Tensor,
@@ -3816,6 +3825,39 @@ def cpu_fused_moe(
3825 return output
3826
3827
3828 +def cpu_fused_moe_int8(
3829 + input: torch.Tensor,
3830 + w13: torch.Tensor,
3831 + w2: torch.Tensor,
3832 + w13_scale: torch.Tensor,
3833 + w2_scale: torch.Tensor,
3834 + w13_bias: torch.Tensor | None,
3835 + w2_bias: torch.Tensor | None,
3836 + topk_weights: torch.Tensor,
3837 + topk_ids: torch.Tensor,
3838 + act: str,
3839 + isa: str,
3840 + skip_weighted: bool = False,
3841 +) -> torch.Tensor:
3842 + output = torch.empty_like(input)
3843 + torch.ops._C.cpu_fused_moe_int8(
3844 + output,
3845 + input,
3846 + w13,
3847 + w2,
3848 + w13_scale,
3849 + w2_scale,
3850 + w13_bias,
3851 + w2_bias,
3852 + topk_weights,
3853 + topk_ids,
3854 + skip_weighted,
3855 + act,
3856 + isa,
3857 + )
3858 + return output
3859 +
3860 +
3861 if hasattr(torch.ops._qutlass_C, "matmul_mxf4_bf16_tn"):
3862
3863 @register_fake("_qutlass_C::matmul_mxf4_bf16_tn")
vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
+179 -3
@@ -9,6 +9,8 @@ from vllm._custom_ops import (
9 CPUQuantAlgo,
10 CPUQuantMethod,
11 convert_weight_packed_scale_zp,
12 + cpu_fused_moe_int8,
13 + cpu_prepack_moe_weight_int8,
14 fused_experts_cpu,
15 )
16 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
@@ -27,7 +29,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
29 kInt8StaticChannelSym,
30 kMxfp4Static,
31 )
30 -from vllm.platforms import current_platform
32 +from vllm.platforms import CpuArchEnum, current_platform
33
34 # ===========================================================================
35 # FP8 W8A16 MoE
@@ -556,7 +558,14 @@ def prepare_int8_moe_layer_for_cpu(
558 w13: torch.Tensor,
559 w2: torch.Tensor,
560 ) -> tuple[torch.Tensor, torch.Tensor]:
559 - """VNNI-prepack INT8 MoE weights for CPU kernel."""
561 + """Prepack INT8 MoE weights for the current CPU architecture."""
562 + # SMMLA packing for AArch64
563 + if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
564 + return (
565 + cpu_prepack_moe_weight_int8(w13, "neon"),
566 + cpu_prepack_moe_weight_int8(w2, "neon"),
567 + )
568 + # VNNI packing for x86
569 packed_w13 = torch.ops._C.convert_weight_packed(w13)
570 packed_w2 = torch.ops._C.convert_weight_packed(w2)
571 return packed_w13, packed_w2
@@ -586,7 +595,10 @@ class CPUExpertsInt8(mk.FusedMoEExpertsMonolithic):
595
596 @staticmethod
597 def _supports_current_device() -> bool:
589 - return current_platform.is_cpu()
598 + return (
599 + current_platform.is_cpu()
600 + and current_platform.get_cpu_architecture() == CpuArchEnum.X86
601 + )
602
603 @staticmethod
604 def _supports_no_act_and_mul() -> bool:
@@ -701,3 +713,167 @@ class CPUExpertsInt8(mk.FusedMoEExpertsMonolithic):
713 None, # limit
714 True, # is_vnni
715 )
716 +
717 +
718 +class ArmCPUExpertsInt8(mk.FusedMoEExpertsMonolithic):
719 + """Arm INT8 MoE with per-token activation and channelwise weight quantization."""
720 +
721 + @property
722 + def expects_unquantized_inputs(self) -> bool:
723 + return True
724 +
725 + @staticmethod
726 + def activation_format() -> mk.FusedMoEActivationFormat:
727 + return mk.FusedMoEActivationFormat.Standard
728 +
729 + @staticmethod
730 + def is_supported_config(
731 + cls: type[mk.FusedMoEExperts],
732 + moe_config: FusedMoEConfig,
733 + weight_key: QuantKey | None,
734 + activation_key: QuantKey | None,
735 + activation_format: mk.FusedMoEActivationFormat,
736 + ) -> tuple[bool, str | None]:
737 + supported, reason = mk.FusedMoEExperts.is_supported_config(
738 + cls,
739 + moe_config,
740 + weight_key,
741 + activation_key,
742 + activation_format,
743 + )
744 + if not supported:
745 + return supported, reason
746 + if moe_config.in_dtype not in (
747 + torch.float32,
748 + torch.float16,
749 + torch.bfloat16,
750 + ):
751 + return False, "kernel requires float32, float16, or bfloat16 activations"
752 + if moe_config.hidden_dim % 32 != 0:
753 + return False, "kernel requires hidden dim divisible by 32"
754 + if moe_config.intermediate_size_per_partition % 32 != 0:
755 + return False, "kernel requires intermediate dim divisible by 32"
756 + return True, None
757 +
758 + @staticmethod
759 + def _supports_current_device() -> bool:
760 + return (
761 + current_platform.is_cpu()
762 + and current_platform.get_cpu_architecture() == CpuArchEnum.ARM
763 + and hasattr(torch.ops._C, "cpu_fused_moe_int8")
764 + )
765 +
766 + @staticmethod
767 + def _supports_no_act_and_mul() -> bool:
768 + return False
769 +
770 + @staticmethod
771 + def _supports_activation(activation: MoEActivation) -> bool:
772 + return activation in (
773 + MoEActivation.SILU,
774 + MoEActivation.SWIGLUOAI,
775 + MoEActivation.GELU,
776 + MoEActivation.GELU_TANH,
777 + )
778 +
779 + @staticmethod
780 + def _supports_parallel_config(
781 + moe_parallel_config: FusedMoEParallelConfig,
782 + ) -> bool:
783 + return not moe_parallel_config.use_ep
784 +
785 + @staticmethod
786 + def _supports_quant_scheme(
787 + weight_key: QuantKey | None,
788 + activation_key: QuantKey | None,
789 + ) -> bool:
790 + return (weight_key, activation_key) == (
791 + kInt8StaticChannelSym,
792 + kInt8DynamicTokenSym,
793 + )
794 +
795 + @staticmethod
796 + def _supports_routing_method(
797 + routing_method: RoutingMethodType,
798 + weight_key: QuantKey | None,
799 + activation_key: QuantKey | None,
800 + ) -> bool:
801 + return routing_method in [
802 + RoutingMethodType.Default,
803 + RoutingMethodType.Renormalize,
804 + RoutingMethodType.RenormalizeNaive,
805 + ]
806 +
807 + @staticmethod
808 + def _supports_router_logits_dtype(
809 + router_logits_dtype: torch.dtype | None,
810 + routing_method: RoutingMethodType,
811 + ) -> bool:
812 + return True
813 +
814 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
815 + from vllm.model_executor.utils import replace_parameter
816 +
817 + w13, w2 = prepare_int8_moe_layer_for_cpu(layer.w13_weight, layer.w2_weight)
818 + replace_parameter(layer, "w13_weight", w13)
819 + replace_parameter(layer, "w2_weight", w2)
820 +
821 + def apply(
822 + self,
823 + hidden_states: torch.Tensor,
824 + w1: torch.Tensor,
825 + w2: torch.Tensor,
826 + router_logits: torch.Tensor,
827 + activation: MoEActivation,
828 + global_num_experts: int,
829 + expert_map: torch.Tensor | None,
830 + a1q_scale: torch.Tensor | None,
831 + apply_router_weight_on_input: bool,
832 + num_expert_group: int | None = None,
833 + e_score_correction_bias: torch.Tensor | None = None,
834 + routed_scaling_factor: float | None = None,
835 + topk_group: int | None = None,
836 + ) -> torch.Tensor:
837 + from vllm.model_executor.layers.fused_moe.cpu_fused_moe import (
838 + select_experts,
839 + )
840 +
841 + topk_weights, topk_ids = select_experts(
842 + hidden_states=hidden_states,
843 + router_logits=router_logits,
844 + use_grouped_topk=num_expert_group is not None,
845 + top_k=self.moe_config.experts_per_token,
846 + renormalize=self.moe_config.routing_method
847 + in (
848 + RoutingMethodType.Renormalize,
849 + RoutingMethodType.RenormalizeNaive,
850 + ),
851 + topk_group=topk_group,
852 + num_expert_group=num_expert_group,
853 + scoring_func="softmax",
854 + routed_scaling_factor=(
855 + routed_scaling_factor if routed_scaling_factor is not None else 1.0
856 + ),
857 + e_score_correction_bias=e_score_correction_bias,
858 + )
859 +
860 + if apply_router_weight_on_input:
861 + assert topk_ids.size(1) == 1
862 + hidden_states.mul_(topk_weights.to(hidden_states.dtype))
863 +
864 + assert self.w1_scale is not None
865 + assert self.w2_scale is not None
866 + return cpu_fused_moe_int8(
867 + hidden_states,
868 + w1,
869 + w2,
870 + self.w1_scale,
871 + self.w2_scale,
872 + self.w1_bias,
873 + self.w2_bias,
874 + topk_weights,
875 + topk_ids,
876 + activation.value,
877 + "neon",
878 + skip_weighted=apply_router_weight_on_input,
879 + )
vllm/model_executor/layers/fused_moe/oracle/int8.py
+3 -13
@@ -24,7 +24,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
24 kInt8StaticChannelSym,
25 )
26 from vllm.model_executor.utils import replace_parameter
27 -from vllm.platforms import current_platform
27
28 logger = init_logger(__name__)
29
@@ -41,20 +40,12 @@ def _get_priority_backends(
40 """
41 Get available backends in priority order based on platform and config.
42 """
44 - _AVAILABLE_BACKENDS = [
43 + return [
44 Int8MoeBackend.TRITON,
45 Int8MoeBackend.HUMMING,
46 Int8MoeBackend.CPU,
47 ]
48
50 - def _move_to_front(backends: list[Int8MoeBackend], backend: Int8MoeBackend) -> None:
51 - backends.insert(0, backends.pop(backends.index(backend)))
52 -
53 - if current_platform.is_cpu():
54 - _move_to_front(_AVAILABLE_BACKENDS, Int8MoeBackend.CPU)
55 -
56 - return _AVAILABLE_BACKENDS
57 -
49
50 def backend_to_kernel_cls(
51 backend: Int8MoeBackend,
@@ -78,14 +69,13 @@ def backend_to_kernel_cls(
69 HummingGroupedExperts,
70 HummingIndexedExperts,
71 ]
81 -
72 elif backend == Int8MoeBackend.CPU:
73 from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
74 + ArmCPUExpertsInt8,
75 CPUExpertsInt8,
76 )
77
87 - return [CPUExpertsInt8]
88 -
78 + return [ArmCPUExpertsInt8, CPUExpertsInt8]
79 else:
80 raise ValueError(f"Unknown Int8 MoE backend: {backend.value}")
81