| 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 | #if defined(__powerpc64__) |
| 30 | #include "cpu/micro_gemm/cpu_micro_gemm_int8_vsx.hpp" |
| 31 | #define VSX_DISPATCH(SCALAR_TYPE, ...) \ |
| 32 | case cpu_utils::ISA::VSX: { \ |
| 33 | using gemm_t = \ |
| 34 | cpu_micro_gemm::MicroGemmINT8<cpu_utils::ISA::VSX, SCALAR_TYPE>; \ |
| 35 | return __VA_ARGS__(); \ |
| 36 | } |
| 37 | #else |
| 38 | #define VSX_DISPATCH(SCALAR_TYPE, ...) case cpu_utils::ISA::VSX: |
| 39 | #endif |
| 40 | |
| 41 | #define CPU_INT8_ISA_DISPATCH_IMPL(ISA_TYPE, SCALAR_TYPE, ...) \ |
| 42 | [&] { \ |
| 43 | switch (ISA_TYPE) { \ |
| 44 | NEON_DISPATCH(SCALAR_TYPE, __VA_ARGS__) \ |
| 45 | VSX_DISPATCH(SCALAR_TYPE, __VA_ARGS__) \ |
| 46 | default: { \ |
| 47 | TORCH_CHECK(false, "Invalid CPU ISA type."); \ |
| 48 | } \ |
| 49 | } \ |
| 50 | }() |
| 51 | |
| 52 | namespace { |
| 53 | using cpu_fused_moe_utils::apply_gated_act; |
| 54 | using cpu_fused_moe_utils::FusedMOEAct; |
| 55 | |
| 56 | template <typename gemm_t> |
| 57 | void prepack_moe_weight_int8_impl(const int8_t* __restrict__ weight_ptr, |
| 58 | int8_t* __restrict__ packed_weight_ptr, |
| 59 | const int32_t expert_num, |
| 60 | const int32_t output_size, |
| 61 | const int32_t input_size, |
| 62 | const int64_t expert_stride) { |
| 63 | #pragma omp parallel for |
| 64 | for (int32_t e_idx = 0; e_idx < expert_num; ++e_idx) { |
| 65 | gemm_t::pack_weight(weight_ptr + expert_stride * e_idx, |
| 66 | packed_weight_ptr + expert_stride * e_idx, output_size, |
| 67 | input_size); |
| 68 | } |
| 69 | } |
| 70 | |
| 71 | // INT8 MoE kernel, based on the original BF16 kernel in cpu_fused_moe.cpp |
| 72 | template <typename scalar_t, typename gemm_t> |
| 73 | void fused_moe_int8_impl( |
| 74 | scalar_t* __restrict__ output, const scalar_t* __restrict__ input, |
| 75 | const int8_t* __restrict__ w13, const int8_t* __restrict__ w2, |
| 76 | const float* __restrict__ w13_scales, const float* __restrict__ w2_scales, |
| 77 | scalar_t* __restrict__ w13_bias, scalar_t* __restrict__ w2_bias, |
| 78 | const float* __restrict__ topk_weights, const int32_t* __restrict__ topk_id, |
| 79 | const FusedMOEAct act_type, const int32_t token_num, |
| 80 | const int32_t expert_num, const int32_t topk_num, |
| 81 | const int32_t input_size_13, const int32_t output_size_13, |
| 82 | const int32_t input_size_2, const int32_t output_size_2, |
| 83 | const bool skip_weighted) { |
| 84 | using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t; |
| 85 | constexpr int32_t gemm_n_tile_size = gemm_t::NSize; |
| 86 | constexpr int32_t gemm_m_tile_size = gemm_t::MaxMSize; |
| 87 | constexpr int32_t min_w13_n_tile_size = 2 * gemm_n_tile_size; |
| 88 | |
| 89 | TORCH_CHECK_EQ(input_size_13 % gemm_t::K, 0); |
| 90 | TORCH_CHECK_EQ(input_size_2 % gemm_t::K, 0); |
| 91 | TORCH_CHECK_EQ(output_size_13 % min_w13_n_tile_size, 0); |
| 92 | TORCH_CHECK_EQ(output_size_2 % gemm_n_tile_size, 0); |
| 93 | TORCH_CHECK_EQ(output_size_13 / 2, input_size_2); |
| 94 | |
| 95 | const int32_t thread_num = cpu_utils::get_max_threads(); |
| 96 | const int32_t w13_input_buffer_size = cpu_utils::round_up<64>( |
| 97 | (gemm_m_tile_size * input_size_13 + 64) * sizeof(int8_t)); |
| 98 | const int32_t w2_input_buffer_size = cpu_utils::round_up<64>( |
| 99 | (gemm_m_tile_size * input_size_2 + 64) * sizeof(int8_t)); |
| 100 | |
| 101 | const int32_t w13_n_tile_size = [&]() { |
| 102 | const int64_t cache_size = cpu_utils::get_available_l2_size(); |
| 103 | const int32_t n_size_cache_limit = |
| 104 | (cache_size - w13_input_buffer_size) / |
| 105 | (gemm_m_tile_size * sizeof(float) + input_size_13 * sizeof(int8_t)); |
| 106 | const int32_t n_size_thread_limit = |
| 107 | output_size_13 / std::max(1, thread_num / topk_num); |
| 108 | const int32_t n_size = cpu_utils::round_down<min_w13_n_tile_size>( |
| 109 | std::min(n_size_cache_limit, n_size_thread_limit)); |
| 110 | return std::max(n_size, min_w13_n_tile_size); |
| 111 | }(); |
| 112 | |
| 113 | const int32_t w2_n_tile_size = [&]() { |
| 114 | const int64_t cache_size = cpu_utils::get_available_l2_size(); |
| 115 | const int32_t n_size_cache_limit = |
| 116 | (cache_size - w2_input_buffer_size) / (input_size_2 * sizeof(int8_t)); |
| 117 | const int32_t n_size_thread_limit = |
| 118 | output_size_2 / std::max(1, thread_num / topk_num); |
| 119 | const int32_t n_size = cpu_utils::round_down<gemm_n_tile_size>( |
| 120 | std::min(n_size_cache_limit, n_size_thread_limit)); |
| 121 | return std::max(n_size, gemm_n_tile_size); |
| 122 | }(); |
| 123 | |
| 124 | int32_t common_buffer_offset = 0; |
| 125 | const int32_t token_num_per_group_buffer_offset = common_buffer_offset; |
| 126 | common_buffer_offset += cpu_utils::round_up<64>(expert_num * sizeof(int32_t)); |
| 127 | const int32_t cu_token_num_per_group_buffer_offset = common_buffer_offset; |
| 128 | common_buffer_offset += |
| 129 | cpu_utils::round_up<64>((expert_num + 1) * sizeof(int32_t)); |
| 130 | const int32_t expanded_token_num = token_num * topk_num; |
| 131 | const int32_t expand_token_id_buffer_offset = common_buffer_offset; |
| 132 | common_buffer_offset += |
| 133 | cpu_utils::round_up<64>(expanded_token_num * sizeof(int32_t)); |
| 134 | const int32_t expand_token_id_index_buffer_offset = common_buffer_offset; |
| 135 | common_buffer_offset += |
| 136 | cpu_utils::round_up<64>(expanded_token_num * sizeof(int32_t)); |
| 137 | const int32_t input_quant_buffer_offset = common_buffer_offset; |
| 138 | common_buffer_offset += |
| 139 | cpu_utils::round_up<64>(token_num * input_size_13 * sizeof(int8_t)); |
| 140 | const int32_t input_scale_buffer_offset = common_buffer_offset; |
| 141 | common_buffer_offset += cpu_utils::round_up<64>(token_num * sizeof(float)); |
| 142 | const int32_t w13_gemm_output_buffer_offset = common_buffer_offset; |
| 143 | common_buffer_offset += cpu_utils::round_up<64>( |
| 144 | expanded_token_num * input_size_2 * sizeof(scalar_t)); |
| 145 | const int32_t w13_output_scale_buffer_offset = common_buffer_offset; |
| 146 | common_buffer_offset += |
| 147 | cpu_utils::round_up<64>(expanded_token_num * sizeof(float)); |
| 148 | const int32_t w2_gemm_output_buffer_offset = common_buffer_offset; |
| 149 | common_buffer_offset += cpu_utils::round_up<64>( |
| 150 | expanded_token_num * output_size_2 * sizeof(float)); |
| 151 | |
| 152 | int32_t gemm_thread_buffer_offset = 0; |
| 153 | const int32_t gemm_input_buffer_offset = gemm_thread_buffer_offset; |
| 154 | gemm_thread_buffer_offset += |
| 155 | std::max(w13_input_buffer_size, w2_input_buffer_size); |
| 156 | const int32_t gemm_output_buffer_offset = gemm_thread_buffer_offset; |
| 157 | gemm_thread_buffer_offset += cpu_utils::round_up<64>( |
| 158 | gemm_m_tile_size * std::max(w13_n_tile_size, w2_n_tile_size) * |
| 159 | sizeof(int32_t)); |
| 160 | |
| 161 | const int32_t ws_output_buffer_offset = 0; |
| 162 | const int32_t ws_thread_buffer_size = |
| 163 | cpu_utils::round_up<64>(output_size_2 * sizeof(float)); |
| 164 | const int32_t thread_buffer_size = |
| 165 | std::max(gemm_thread_buffer_offset, ws_thread_buffer_size); |
| 166 | const int32_t buffer_size = |
| 167 | common_buffer_offset + thread_buffer_size * thread_num; |
| 168 | cpu_utils::ScratchPadManager::get_scratchpad_manager()->realloc(buffer_size); |
| 169 | uint8_t* common_buffer_start = |
| 170 | cpu_utils::ScratchPadManager::get_scratchpad_manager() |
| 171 | ->get_data<uint8_t>(); |
| 172 | uint8_t* thread_buffer_start = common_buffer_start + common_buffer_offset; |
| 173 | |
| 174 | int32_t* __restrict__ token_num_per_group_buffer = reinterpret_cast<int32_t*>( |
| 175 | common_buffer_start + token_num_per_group_buffer_offset); |
| 176 | int32_t* __restrict__ cu_token_num_per_group_buffer = |
| 177 | reinterpret_cast<int32_t*>(common_buffer_start + |
| 178 | cu_token_num_per_group_buffer_offset); |
| 179 | int32_t* __restrict__ expand_token_id_buffer = reinterpret_cast<int32_t*>( |
| 180 | common_buffer_start + expand_token_id_buffer_offset); |
| 181 | int32_t* __restrict__ expand_token_id_index_buffer = |
| 182 | reinterpret_cast<int32_t*>(common_buffer_start + |
| 183 | expand_token_id_index_buffer_offset); |
| 184 | int8_t* __restrict__ input_quant_buffer = reinterpret_cast<int8_t*>( |
| 185 | common_buffer_start + input_quant_buffer_offset); |
| 186 | float* __restrict__ input_scale_buffer = |
| 187 | reinterpret_cast<float*>(common_buffer_start + input_scale_buffer_offset); |
| 188 | |
| 189 | std::memset(token_num_per_group_buffer, 0, expert_num * sizeof(int32_t)); |
| 190 | for (int32_t i = 0; i < expanded_token_num; ++i) { |
| 191 | ++token_num_per_group_buffer[topk_id[i]]; |
| 192 | } |
| 193 | |
| 194 | int32_t token_num_sum = 0; |
| 195 | cu_token_num_per_group_buffer[0] = 0; |
| 196 | int32_t* token_index_buffer = cu_token_num_per_group_buffer + 1; |
| 197 | for (int32_t i = 0; i < expert_num; ++i) { |
| 198 | token_index_buffer[i] = token_num_sum; |
| 199 | token_num_sum += token_num_per_group_buffer[i]; |
| 200 | } |
| 201 | |
| 202 | for (int32_t i = 0; i < token_num; ++i) { |
| 203 | const int32_t* curr_topk_id = topk_id + i * topk_num; |
| 204 | int32_t* curr_index_buffer = expand_token_id_index_buffer + i * topk_num; |
| 205 | for (int32_t j = 0; j < topk_num; ++j) { |
| 206 | const int32_t curr_expert_id = curr_topk_id[j]; |
| 207 | const int32_t curr_index = token_index_buffer[curr_expert_id]++; |
| 208 | expand_token_id_buffer[curr_index] = i; |
| 209 | curr_index_buffer[j] = curr_index; |
| 210 | } |
| 211 | } |
| 212 | |
| 213 | // quantize inputs |
| 214 | #pragma omp parallel for |
| 215 | for (int32_t token_idx = 0; token_idx < token_num; ++token_idx) { |
| 216 | gemm_t::quantize_row(input + token_idx * input_size_13, |
| 217 | input_quant_buffer + token_idx * input_size_13, |
| 218 | input_scale_buffer[token_idx], input_size_13); |
| 219 | } |
| 220 | |
| 221 | { |
| 222 | alignas(64) cpu_utils::Counter counter; |
| 223 | cpu_utils::Counter* counter_ptr = &counter; |
| 224 | |
| 225 | // w13 GEMM + act |
| 226 | #pragma omp parallel for schedule(static, 1) |
| 227 | for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) { |
| 228 | const int32_t task_num_per_expert = |
| 229 | (output_size_13 + w13_n_tile_size - 1) / w13_n_tile_size; |
| 230 | const int32_t task_num = task_num_per_expert * expert_num; |
| 231 | uint8_t* __restrict__ thread_buffer = |
| 232 | thread_buffer_start + thread_id * thread_buffer_size; |
| 233 | int8_t* __restrict__ gemm_input_buffer = |
| 234 | reinterpret_cast<int8_t*>(thread_buffer + gemm_input_buffer_offset); |
| 235 | float* __restrict__ gemm_output_buffer = |
| 236 | reinterpret_cast<float*>(thread_buffer + gemm_output_buffer_offset); |
| 237 | auto* __restrict__ w13_gemm_output_buffer = reinterpret_cast<scalar_t*>( |
| 238 | common_buffer_start + w13_gemm_output_buffer_offset); |
| 239 | gemm_t gemm; |
| 240 | |
| 241 | const int32_t w13_n_group_stride = |
| 242 | gemm_t::WeightOCGroupSize * input_size_13; |
| 243 | const int32_t w13_n_tile_stride = gemm_n_tile_size * input_size_13; |
| 244 | |
| 245 | for (;;) { |
| 246 | const int32_t task_id = counter_ptr->acquire_counter(); |
| 247 | if (task_id >= task_num) { |
| 248 | break; |
| 249 | } |
| 250 | const int32_t curr_expert_id = task_id / task_num_per_expert; |
| 251 | const int32_t curr_output_group_id = task_id % task_num_per_expert; |
| 252 | const int32_t curr_token_num = |
| 253 | token_num_per_group_buffer[curr_expert_id]; |
| 254 | if (curr_token_num == 0) { |
| 255 | continue; |
| 256 | } |
| 257 | |
| 258 | const int32_t actual_n_tile_size = |
| 259 | std::min(w13_n_tile_size, |
| 260 | output_size_13 - curr_output_group_id * w13_n_tile_size); |
| 261 | const int32_t* __restrict__ curr_expand_token_id_buffer = |
| 262 | expand_token_id_buffer + |
| 263 | cu_token_num_per_group_buffer[curr_expert_id]; |
| 264 | scalar_t* __restrict__ curr_w13_gemm_output_buffer = |
| 265 | w13_gemm_output_buffer + |
| 266 | cu_token_num_per_group_buffer[curr_expert_id] * input_size_2 + |
| 267 | curr_output_group_id * w13_n_tile_size / 2; |
| 268 | |
| 269 | const int8_t* w13_weight_ptr_0 = nullptr; |
| 270 | const int8_t* w13_weight_ptr_1 = nullptr; |
| 271 | const float* w13_scale_ptr_0 = nullptr; |
| 272 | const float* w13_scale_ptr_1 = nullptr; |
| 273 | scalar_t* w13_bias_ptr_0 = nullptr; |
| 274 | scalar_t* w13_bias_ptr_1 = nullptr; |
| 275 | if (act_type == FusedMOEAct::SwigluOAIAndMul) { |
| 276 | const int32_t output_offset = curr_output_group_id * w13_n_tile_size; |
| 277 | w13_weight_ptr_0 = w13 + |
| 278 | curr_expert_id * input_size_13 * output_size_13 + |
| 279 | output_offset * input_size_13; |
| 280 | w13_weight_ptr_1 = |
| 281 | w13_weight_ptr_0 + actual_n_tile_size / 2 * input_size_13; |
| 282 | w13_scale_ptr_0 = |
| 283 | w13_scales + curr_expert_id * output_size_13 + output_offset; |
| 284 | w13_scale_ptr_1 = w13_scale_ptr_0 + actual_n_tile_size / 2; |
| 285 | if (w13_bias != nullptr) { |
| 286 | w13_bias_ptr_0 = |
| 287 | w13_bias + curr_expert_id * output_size_13 + output_offset; |
| 288 | w13_bias_ptr_1 = w13_bias_ptr_0 + actual_n_tile_size / 2; |
| 289 | } |
| 290 | } else { |
| 291 | const int32_t output_offset = |
| 292 | curr_output_group_id * (w13_n_tile_size / 2); |
| 293 | w13_weight_ptr_0 = w13 + |
| 294 | curr_expert_id * input_size_13 * output_size_13 + |
| 295 | output_offset * input_size_13; |
| 296 | w13_weight_ptr_1 = |
| 297 | w13_weight_ptr_0 + output_size_13 / 2 * input_size_13; |
| 298 | w13_scale_ptr_0 = |
| 299 | w13_scales + curr_expert_id * output_size_13 + output_offset; |
| 300 | w13_scale_ptr_1 = w13_scale_ptr_0 + output_size_13 / 2; |
| 301 | if (w13_bias != nullptr) { |
| 302 | w13_bias_ptr_0 = |
| 303 | w13_bias + curr_expert_id * output_size_13 + output_offset; |
| 304 | w13_bias_ptr_1 = w13_bias_ptr_0 + output_size_13 / 2; |
| 305 | } |
| 306 | } |
| 307 | |
| 308 | for (int32_t token_idx = 0; token_idx < curr_token_num; |
| 309 | token_idx += gemm_m_tile_size) { |
| 310 | const int32_t actual_token_num = |
| 311 | std::min(gemm_m_tile_size, curr_token_num - token_idx); |
| 312 | const int8_t* input_rows[gemm_m_tile_size]; |
| 313 | alignas(64) float input_scales[gemm_m_tile_size]; |
| 314 | // gather and pack |
| 315 | for (int32_t i = 0; i < actual_token_num; ++i) { |
| 316 | const int32_t curr_token_id = curr_expand_token_id_buffer[i]; |
| 317 | input_rows[i] = input_quant_buffer + curr_token_id * input_size_13; |
| 318 | input_scales[i] = input_scale_buffer[curr_token_id]; |
| 319 | } |
| 320 | gemm_t::pack_input_from_rows(input_rows, gemm_input_buffer, |
| 321 | actual_token_num, input_size_13); |
| 322 | curr_expand_token_id_buffer += actual_token_num; |
| 323 | |
| 324 | const int8_t* w13_weight_ptr_0_iter = w13_weight_ptr_0; |
| 325 | const int8_t* w13_weight_ptr_1_iter = w13_weight_ptr_1; |
| 326 | const float* w13_scale_ptr_0_iter = w13_scale_ptr_0; |
| 327 | const float* w13_scale_ptr_1_iter = w13_scale_ptr_1; |
| 328 | scalar_t* w13_bias_ptr_0_iter = w13_bias_ptr_0; |
| 329 | scalar_t* w13_bias_ptr_1_iter = w13_bias_ptr_1; |
| 330 | float* w13_output_buffer_0_iter = gemm_output_buffer; |
| 331 | float* w13_output_buffer_1_iter = |
| 332 | gemm_output_buffer + actual_n_tile_size / 2; |
| 333 | |
| 334 | for (int32_t i = 0; i < actual_n_tile_size; |
| 335 | i += min_w13_n_tile_size) { |
| 336 | auto* output_0_int32 = |
| 337 | reinterpret_cast<int32_t*>(w13_output_buffer_0_iter); |
| 338 | gemm.gemm(gemm_input_buffer, w13_weight_ptr_0_iter, output_0_int32, |
| 339 | actual_token_num, input_size_13, w13_n_group_stride, |
| 340 | actual_n_tile_size); |
| 341 | gemm_t::dequantize_tile(output_0_int32, w13_output_buffer_0_iter, |
| 342 | input_scales, w13_scale_ptr_0_iter, |
| 343 | actual_token_num, gemm_n_tile_size, |
| 344 | actual_n_tile_size); |
| 345 | if (w13_bias != nullptr) { |
| 346 | cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>( |
| 347 | w13_output_buffer_0_iter, w13_output_buffer_0_iter, |
| 348 | w13_bias_ptr_0_iter, actual_token_num, actual_n_tile_size, |
| 349 | actual_n_tile_size); |
| 350 | w13_bias_ptr_0_iter += gemm_n_tile_size; |
| 351 | } |
| 352 | |
| 353 | auto* output_1_int32 = |
| 354 | reinterpret_cast<int32_t*>(w13_output_buffer_1_iter); |
| 355 | gemm.gemm(gemm_input_buffer, w13_weight_ptr_1_iter, output_1_int32, |
| 356 | actual_token_num, input_size_13, w13_n_group_stride, |
| 357 | actual_n_tile_size); |
| 358 | gemm_t::dequantize_tile(output_1_int32, w13_output_buffer_1_iter, |
| 359 | input_scales, w13_scale_ptr_1_iter, |
| 360 | actual_token_num, gemm_n_tile_size, |
| 361 | actual_n_tile_size); |
| 362 | if (w13_bias != nullptr) { |
| 363 | cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>( |
| 364 | w13_output_buffer_1_iter, w13_output_buffer_1_iter, |
| 365 | w13_bias_ptr_1_iter, actual_token_num, actual_n_tile_size, |
| 366 | actual_n_tile_size); |
| 367 | w13_bias_ptr_1_iter += gemm_n_tile_size; |
| 368 | } |
| 369 | |
| 370 | w13_weight_ptr_0_iter += w13_n_tile_stride; |
| 371 | w13_weight_ptr_1_iter += w13_n_tile_stride; |
| 372 | w13_scale_ptr_0_iter += gemm_n_tile_size; |
| 373 | w13_scale_ptr_1_iter += gemm_n_tile_size; |
| 374 | w13_output_buffer_0_iter += gemm_n_tile_size; |
| 375 | w13_output_buffer_1_iter += gemm_n_tile_size; |
| 376 | } |
| 377 | |
| 378 | apply_gated_act(act_type, gemm_output_buffer, |
| 379 | curr_w13_gemm_output_buffer, actual_token_num, |
| 380 | actual_n_tile_size, actual_n_tile_size, input_size_2); |
| 381 | curr_w13_gemm_output_buffer += gemm_m_tile_size * input_size_2; |
| 382 | } |
| 383 | } |
| 384 | } |
| 385 | } |
| 386 | |
| 387 | auto* __restrict__ w13_gemm_output_buffer = reinterpret_cast<scalar_t*>( |
| 388 | common_buffer_start + w13_gemm_output_buffer_offset); |
| 389 | float* __restrict__ w13_output_scale_buffer = reinterpret_cast<float*>( |
| 390 | common_buffer_start + w13_output_scale_buffer_offset); |
| 391 | |
| 392 | // quantize w2 inputs - in place |
| 393 | #pragma omp parallel for |
| 394 | for (int32_t token_idx = 0; token_idx < expanded_token_num; ++token_idx) { |
| 395 | scalar_t* input_row = w13_gemm_output_buffer + token_idx * input_size_2; |
| 396 | int8_t* output_row = reinterpret_cast<int8_t*>(input_row); |
| 397 | gemm_t::quantize_row(input_row, output_row, |
| 398 | w13_output_scale_buffer[token_idx], input_size_2); |
| 399 | } |
| 400 | |
| 401 | { |
| 402 | alignas(64) cpu_utils::Counter counter; |
| 403 | cpu_utils::Counter* counter_ptr = &counter; |
| 404 | |
| 405 | // w2 gemm |
| 406 | #pragma omp parallel for schedule(static, 1) |
| 407 | for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) { |
| 408 | const int32_t task_num_per_expert = |
| 409 | (output_size_2 + w2_n_tile_size - 1) / w2_n_tile_size; |
| 410 | const int32_t task_num = task_num_per_expert * expert_num; |
| 411 | uint8_t* __restrict__ thread_buffer = |
| 412 | thread_buffer_start + thread_id * thread_buffer_size; |
| 413 | int8_t* __restrict__ gemm_input_buffer = |
| 414 | reinterpret_cast<int8_t*>(thread_buffer + gemm_input_buffer_offset); |
| 415 | float* __restrict__ gemm_output_buffer = |
| 416 | reinterpret_cast<float*>(thread_buffer + gemm_output_buffer_offset); |
| 417 | float* __restrict__ w2_gemm_output_buffer = reinterpret_cast<float*>( |
| 418 | common_buffer_start + w2_gemm_output_buffer_offset); |
| 419 | gemm_t gemm; |
| 420 | |
| 421 | const int32_t w2_n_group_stride = |
| 422 | gemm_t::WeightOCGroupSize * input_size_2; |
| 423 | const int32_t w2_n_tile_stride = gemm_n_tile_size * input_size_2; |
| 424 | |
| 425 | for (;;) { |
| 426 | const int32_t task_id = counter_ptr->acquire_counter(); |
| 427 | if (task_id >= task_num) { |
| 428 | break; |
| 429 | } |
| 430 | const int32_t curr_expert_id = task_id / task_num_per_expert; |
| 431 | const int32_t curr_output_group_id = task_id % task_num_per_expert; |
| 432 | const int32_t curr_token_num = |
| 433 | token_num_per_group_buffer[curr_expert_id]; |
| 434 | if (curr_token_num == 0) { |
| 435 | continue; |
| 436 | } |
| 437 | |
| 438 | const int32_t actual_n_tile_size = |
| 439 | std::min(w2_n_tile_size, |
| 440 | output_size_2 - curr_output_group_id * w2_n_tile_size); |
| 441 | scalar_t* __restrict__ curr_w13_gemm_output_buffer = |
| 442 | w13_gemm_output_buffer + |
| 443 | cu_token_num_per_group_buffer[curr_expert_id] * input_size_2; |
| 444 | float* __restrict__ curr_w13_output_scale_buffer = |
| 445 | w13_output_scale_buffer + |
| 446 | cu_token_num_per_group_buffer[curr_expert_id]; |
| 447 | float* __restrict__ curr_w2_gemm_output_buffer = |
| 448 | w2_gemm_output_buffer + |
| 449 | cu_token_num_per_group_buffer[curr_expert_id] * output_size_2 + |
| 450 | curr_output_group_id * w2_n_tile_size; |
| 451 | const int8_t* __restrict__ w2_weight_ptr = |
| 452 | w2 + curr_expert_id * output_size_2 * input_size_2 + |
| 453 | curr_output_group_id * w2_n_tile_size * input_size_2; |
| 454 | const float* __restrict__ w2_scale_ptr = |
| 455 | w2_scales + curr_expert_id * output_size_2 + |
| 456 | curr_output_group_id * w2_n_tile_size; |
| 457 | scalar_t* w2_bias_ptr = nullptr; |
| 458 | if (w2_bias != nullptr) { |
| 459 | w2_bias_ptr = w2_bias + curr_expert_id * output_size_2 + |
| 460 | curr_output_group_id * w2_n_tile_size; |
| 461 | } |
| 462 | |
| 463 | for (int32_t token_idx = 0; token_idx < curr_token_num; |
| 464 | token_idx += gemm_m_tile_size) { |
| 465 | const int32_t actual_token_num = |
| 466 | std::min(gemm_m_tile_size, curr_token_num - token_idx); |
| 467 | const int8_t* input_rows[gemm_m_tile_size]; |
| 468 | alignas(64) float input_scales[gemm_m_tile_size]; |
| 469 | for (int32_t i = 0; i < actual_token_num; ++i) { |
| 470 | input_rows[i] = reinterpret_cast<const int8_t*>( |
| 471 | curr_w13_gemm_output_buffer + i * input_size_2); |
| 472 | input_scales[i] = curr_w13_output_scale_buffer[i]; |
| 473 | } |
| 474 | gemm_t::pack_input_from_rows(input_rows, gemm_input_buffer, |
| 475 | actual_token_num, input_size_2); |
| 476 | |
| 477 | const int8_t* w2_weight_ptr_iter = w2_weight_ptr; |
| 478 | const float* w2_scale_ptr_iter = w2_scale_ptr; |
| 479 | scalar_t* w2_bias_ptr_iter = w2_bias_ptr; |
| 480 | float* curr_w2_gemm_output_buffer_iter = curr_w2_gemm_output_buffer; |
| 481 | for (int32_t i = 0; i < actual_n_tile_size; i += gemm_n_tile_size) { |
| 482 | auto* output_int32 = reinterpret_cast<int32_t*>(gemm_output_buffer); |
| 483 | gemm.gemm(gemm_input_buffer, w2_weight_ptr_iter, output_int32, |
| 484 | actual_token_num, input_size_2, w2_n_group_stride, |
| 485 | gemm_n_tile_size); |
| 486 | gemm_t::dequantize_tile(output_int32, gemm_output_buffer, |
| 487 | input_scales, w2_scale_ptr_iter, |
| 488 | actual_token_num, gemm_n_tile_size, |
| 489 | gemm_n_tile_size); |
| 490 | if (w2_bias != nullptr) { |
| 491 | cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>( |
| 492 | gemm_output_buffer, gemm_output_buffer, w2_bias_ptr_iter, |
| 493 | actual_token_num, gemm_n_tile_size, gemm_n_tile_size); |
| 494 | w2_bias_ptr_iter += gemm_n_tile_size; |
| 495 | } |
| 496 | for (int32_t m_idx = 0; m_idx < actual_token_num; ++m_idx) { |
| 497 | std::memcpy( |
| 498 | curr_w2_gemm_output_buffer_iter + m_idx * output_size_2, |
| 499 | gemm_output_buffer + m_idx * gemm_n_tile_size, |
| 500 | gemm_n_tile_size * sizeof(float)); |
| 501 | } |
| 502 | |
| 503 | w2_weight_ptr_iter += w2_n_tile_stride; |
| 504 | w2_scale_ptr_iter += gemm_n_tile_size; |
| 505 | curr_w2_gemm_output_buffer_iter += gemm_n_tile_size; |
| 506 | } |
| 507 | |
| 508 | curr_w13_gemm_output_buffer += gemm_m_tile_size * input_size_2; |
| 509 | curr_w13_output_scale_buffer += gemm_m_tile_size; |
| 510 | curr_w2_gemm_output_buffer += gemm_m_tile_size * output_size_2; |
| 511 | } |
| 512 | } |
| 513 | } |
| 514 | } |
| 515 | |
| 516 | { |
| 517 | alignas(64) cpu_utils::Counter counter; |
| 518 | cpu_utils::Counter* counter_ptr = &counter; |
| 519 | |
| 520 | #pragma omp parallel for schedule(static, 1) |
| 521 | for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) { |
| 522 | uint8_t* __restrict__ thread_buffer = |
| 523 | thread_buffer_start + thread_id * thread_buffer_size; |
| 524 | float* __restrict__ ws_output_buffer = |
| 525 | reinterpret_cast<float*>(thread_buffer + ws_output_buffer_offset); |
| 526 | float* __restrict__ w2_gemm_output_buffer = reinterpret_cast<float*>( |
| 527 | common_buffer_start + w2_gemm_output_buffer_offset); |
| 528 | |
| 529 | for (;;) { |
| 530 | const int32_t token_id = counter_ptr->acquire_counter(); |
| 531 | if (token_id >= token_num) { |
| 532 | break; |
| 533 | } |
| 534 | int32_t* __restrict__ curr_expand_token_id_index_buffer = |
| 535 | expand_token_id_index_buffer + token_id * topk_num; |
| 536 | const float* __restrict__ curr_weight = |
| 537 | topk_weights + token_id * topk_num; |
| 538 | const float first_weight = skip_weighted ? 1.0f : curr_weight[0]; |
| 539 | scalar_t* __restrict__ curr_output_buffer = |
| 540 | output + token_id * output_size_2; |
| 541 | |
| 542 | if (topk_num > 1) { |
| 543 | int32_t w2_output_idx = curr_expand_token_id_index_buffer[0]; |
| 544 | float* w2_output_iter = |
| 545 | w2_gemm_output_buffer + w2_output_idx * output_size_2; |
| 546 | float* ws_output_buffer_iter = ws_output_buffer; |
| 547 | vec_op::FP32Vec16 weight_vec(first_weight); |
| 548 | for (int32_t i = 0; i < output_size_2; i += 16) { |
| 549 | vec_op::FP32Vec16 vec(w2_output_iter); |
| 550 | (vec * weight_vec).save(ws_output_buffer_iter); |
| 551 | w2_output_iter += 16; |
| 552 | ws_output_buffer_iter += 16; |
| 553 | } |
| 554 | |
| 555 | for (int32_t idx = 1; idx < topk_num - 1; ++idx) { |
| 556 | w2_output_idx = curr_expand_token_id_index_buffer[idx]; |
| 557 | w2_output_iter = |
| 558 | w2_gemm_output_buffer + w2_output_idx * output_size_2; |
| 559 | ws_output_buffer_iter = ws_output_buffer; |
| 560 | weight_vec = vec_op::FP32Vec16(curr_weight[idx]); |
| 561 | for (int32_t i = 0; i < output_size_2; i += 16) { |
| 562 | vec_op::FP32Vec16 vec(w2_output_iter); |
| 563 | vec_op::FP32Vec16 sum(ws_output_buffer_iter); |
| 564 | (sum + vec * weight_vec).save(ws_output_buffer_iter); |
| 565 | w2_output_iter += 16; |
| 566 | ws_output_buffer_iter += 16; |
| 567 | } |
| 568 | } |
| 569 | |
| 570 | const int32_t last_idx = topk_num - 1; |
| 571 | w2_output_idx = curr_expand_token_id_index_buffer[last_idx]; |
| 572 | w2_output_iter = |
| 573 | w2_gemm_output_buffer + w2_output_idx * output_size_2; |
| 574 | ws_output_buffer_iter = ws_output_buffer; |
| 575 | scalar_t* curr_output_buffer_iter = curr_output_buffer; |
| 576 | weight_vec = vec_op::FP32Vec16(curr_weight[last_idx]); |
| 577 | for (int32_t i = 0; i < output_size_2; i += 16) { |
| 578 | vec_op::FP32Vec16 vec(w2_output_iter); |
| 579 | vec_op::FP32Vec16 sum(ws_output_buffer_iter); |
| 580 | scalar_vec_t(sum + vec * weight_vec).save(curr_output_buffer_iter); |
| 581 | w2_output_iter += 16; |
| 582 | ws_output_buffer_iter += 16; |
| 583 | curr_output_buffer_iter += 16; |
| 584 | } |
| 585 | } else { |
| 586 | const int32_t w2_output_idx = curr_expand_token_id_index_buffer[0]; |
| 587 | float* w2_output_iter = |
| 588 | w2_gemm_output_buffer + w2_output_idx * output_size_2; |
| 589 | scalar_t* curr_output_buffer_iter = curr_output_buffer; |
| 590 | vec_op::FP32Vec16 weight_vec(first_weight); |
| 591 | for (int32_t i = 0; i < output_size_2; i += 16) { |
| 592 | vec_op::FP32Vec16 vec(w2_output_iter); |
| 593 | scalar_vec_t(vec * weight_vec).save(curr_output_buffer_iter); |
| 594 | w2_output_iter += 16; |
| 595 | curr_output_buffer_iter += 16; |
| 596 | } |
| 597 | } |
| 598 | } |
| 599 | } |
| 600 | } |
| 601 | } |
| 602 | } // namespace |
| 603 | |
| 604 | void prepack_moe_weight_int8( |
| 605 | const torch::Tensor& weight, // [expert_num, output_size, input_size] |
| 606 | torch::Tensor& packed_weight, const std::string& isa) { |
| 607 | TORCH_CHECK(weight.is_contiguous()); |
| 608 | const int32_t expert_num = weight.size(0); |
| 609 | const int32_t output_size = weight.size(1); |
| 610 | const int32_t input_size = weight.size(2); |
| 611 | const int64_t expert_stride = weight.stride(0); |
| 612 | const cpu_utils::ISA isa_type = cpu_utils::get_isa(isa); |
| 613 | TORCH_CHECK_EQ(output_size % 32, 0); |
| 614 | |
| 615 | CPU_INT8_ISA_DISPATCH_IMPL(isa_type, c10::BFloat16, [&]() { |
| 616 | TORCH_CHECK_EQ(input_size % gemm_t::K, 0); |
| 617 | prepack_moe_weight_int8_impl<gemm_t>( |
| 618 | weight.data_ptr<int8_t>(), packed_weight.data_ptr<int8_t>(), expert_num, |
| 619 | output_size, input_size, expert_stride); |
| 620 | }); |
| 621 | } |
| 622 | |
| 623 | void cpu_fused_moe_int8(torch::Tensor& output, const torch::Tensor& input, |
| 624 | const torch::Tensor& w13, const torch::Tensor& w2, |
| 625 | const torch::Tensor& w13_scale, |
| 626 | const torch::Tensor& w2_scale, |
| 627 | const std::optional<torch::Tensor>& w13_bias, |
| 628 | const std::optional<torch::Tensor>& w2_bias, |
| 629 | const torch::Tensor& topk_weights, |
| 630 | const torch::Tensor& topk_id, const bool skip_weighted, |
| 631 | const std::string& act, const std::string& isa) { |
| 632 | const int32_t token_num = input.size(0); |
| 633 | const int32_t input_size_13 = input.size(1); |
| 634 | const int64_t input_stride = input.stride(0); |
| 635 | TORCH_CHECK_EQ(input_stride, input_size_13); |
| 636 | const int32_t expert_num = w13.size(0); |
| 637 | const int32_t output_size_13 = w13.size(1); |
| 638 | const int32_t input_size_2 = w2.size(2); |
| 639 | const int32_t output_size_2 = w2.size(1); |
| 640 | const int32_t topk_num = topk_id.size(1); |
| 641 | const FusedMOEAct act_type = cpu_fused_moe_utils::get_act_type(act); |
| 642 | const cpu_utils::ISA isa_type = cpu_utils::get_isa(isa); |
| 643 | TORCH_CHECK(!skip_weighted || topk_num == 1, |
| 644 | "skip_weighted is only supported for topk=1 on CPU"); |
| 645 | |
| 646 | VLLM_DISPATCH_FLOATING_TYPES( |
| 647 | input.scalar_type(), "cpu_fused_moe_int8", [&]() { |
| 648 | CPU_INT8_ISA_DISPATCH_IMPL(isa_type, scalar_t, [&]() { |
| 649 | fused_moe_int8_impl<scalar_t, gemm_t>( |
| 650 | output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), |
| 651 | w13.data_ptr<int8_t>(), w2.data_ptr<int8_t>(), |
| 652 | w13_scale.data_ptr<float>(), w2_scale.data_ptr<float>(), |
| 653 | w13_bias.has_value() ? w13_bias->data_ptr<scalar_t>() : nullptr, |
| 654 | w2_bias.has_value() ? w2_bias->data_ptr<scalar_t>() : nullptr, |
| 655 | topk_weights.data_ptr<float>(), topk_id.data_ptr<int32_t>(), |
| 656 | act_type, token_num, expert_num, topk_num, input_size_13, |
| 657 | output_size_13, input_size_2, output_size_2, skip_weighted); |
| 658 | }); |
| 659 | }); |
| 660 | } |