| 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 | |
| 7 | #ifdef CPU_CAPABILITY_AMXBF16 |
| 8 | #include "cpu/micro_gemm/cpu_micro_gemm_amx.hpp" |
| 9 | #define AMX_DISPATCH(...) \ |
| 10 | case cpu_utils::ISA::AMX: { \ |
| 11 | using gemm_t = cpu_micro_gemm::MicroGemm<cpu_utils::ISA::AMX, scalar_t>; \ |
| 12 | return __VA_ARGS__(); \ |
| 13 | } |
| 14 | #else |
| 15 | #define AMX_DISPATCH(...) case cpu_utils::ISA::AMX: |
| 16 | #endif |
| 17 | |
| 18 | #if defined(ARM_BF16_SUPPORT) |
| 19 | #include "cpu/micro_gemm/cpu_micro_gemm_neon.hpp" |
| 20 | #define NEON_DISPATCH(...) \ |
| 21 | case cpu_utils::ISA::NEON: { \ |
| 22 | using gemm_t = \ |
| 23 | cpu_micro_gemm::MicroGemm<cpu_utils::ISA::NEON, scalar_t>; \ |
| 24 | return __VA_ARGS__(); \ |
| 25 | } |
| 26 | #else |
| 27 | #define NEON_DISPATCH(...) case cpu_utils::ISA::NEON: |
| 28 | #endif |
| 29 | |
| 30 | #ifdef __powerpc64__ |
| 31 | #include "cpu/micro_gemm/cpu_micro_gemm_vsx.hpp" |
| 32 | #define VSX_DISPATCH(...) \ |
| 33 | case cpu_utils::ISA::VSX: { \ |
| 34 | using gemm_t = cpu_micro_gemm::MicroGemm<cpu_utils::ISA::VSX, scalar_t>; \ |
| 35 | return __VA_ARGS__(); \ |
| 36 | } |
| 37 | #else |
| 38 | #define VSX_DISPATCH(...) case cpu_utils::ISA::VSX: |
| 39 | #endif |
| 40 | |
| 41 | #define CPU_ISA_DISPATCH_IMPL(ISA_TYPE, ...) \ |
| 42 | [&] { \ |
| 43 | switch (ISA_TYPE) { \ |
| 44 | AMX_DISPATCH(__VA_ARGS__) \ |
| 45 | VSX_DISPATCH(__VA_ARGS__) \ |
| 46 | case cpu_utils::ISA::VEC: { \ |
| 47 | using gemm_t = \ |
| 48 | cpu_micro_gemm::MicroGemm<cpu_utils::ISA::VEC, scalar_t>; \ |
| 49 | return __VA_ARGS__(); \ |
| 50 | } \ |
| 51 | NEON_DISPATCH(__VA_ARGS__) \ |
| 52 | default: { \ |
| 53 | TORCH_CHECK(false, "Invalid CPU ISA type."); \ |
| 54 | } \ |
| 55 | } \ |
| 56 | }() |
| 57 | |
| 58 | namespace { |
| 59 | |
| 60 | using cpu_fused_moe_utils::apply_gated_act; |
| 61 | using cpu_fused_moe_utils::FusedMOEAct; |
| 62 | |
| 63 | template <typename scalar_t, typename gemm_t> |
| 64 | void prepack_moe_weight_impl(scalar_t* __restrict__ weight_ptr, |
| 65 | scalar_t* __restrict__ packed_weight_ptr, |
| 66 | const int32_t expert_num, |
| 67 | const int32_t output_size, |
| 68 | const int32_t input_size, |
| 69 | const int64_t expert_stride) { |
| 70 | #pragma omp parallel for |
| 71 | for (int32_t e_idx = 0; e_idx < expert_num; ++e_idx) { |
| 72 | gemm_t::pack_weight(weight_ptr + expert_stride * e_idx, |
| 73 | packed_weight_ptr + expert_stride * e_idx, output_size, |
| 74 | input_size); |
| 75 | } |
| 76 | } |
| 77 | |
| 78 | template <typename scalar_t, typename w_t, typename gemm_t> |
| 79 | void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input, |
| 80 | w_t* __restrict__ w13, w_t* __restrict__ w2, |
| 81 | w_t* __restrict__ w13_bias, w_t* __restrict__ w2_bias, |
| 82 | float* __restrict__ topk_weights, |
| 83 | int32_t* __restrict__ topk_id, FusedMOEAct act_type, |
| 84 | const int32_t token_num, const int32_t expert_num, |
| 85 | const int32_t topk_num, const int32_t input_size_13, |
| 86 | const int32_t output_size_13, const int32_t input_size_2, |
| 87 | const int32_t output_size_2, const bool skip_weighted) { |
| 88 | using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t; |
| 89 | constexpr int32_t gemm_n_tile_size = gemm_t::NSize; |
| 90 | constexpr int32_t gemm_m_tile_size = gemm_t::MaxMSize; |
| 91 | constexpr int32_t min_w13_n_tile_size = 2 * gemm_n_tile_size; |
| 92 | constexpr bool pack_a = gemm_t::PackA; |
| 93 | static_assert(gemm_n_tile_size % 16 == 0); |
| 94 | |
| 95 | TORCH_CHECK_EQ(output_size_13 % min_w13_n_tile_size, 0); |
| 96 | TORCH_CHECK_EQ(output_size_2 % gemm_n_tile_size, 0); |
| 97 | TORCH_CHECK_EQ(output_size_13 / 2, input_size_2); |
| 98 | |
| 99 | const int32_t thread_num = cpu_utils::get_max_threads(); |
| 100 | |
| 101 | const int32_t w13_input_buffer_size = cpu_utils::round_up<64>( |
| 102 | gemm_m_tile_size * input_size_13 * sizeof(scalar_t)); |
| 103 | |
| 104 | const int32_t w13_n_tile_size = [&]() { |
| 105 | const int64_t cache_size = cpu_utils::get_available_l2_size(); |
| 106 | // input buffer + output buffer + weight |
| 107 | const int32_t n_size_cache_limit = |
| 108 | (cache_size - w13_input_buffer_size) / |
| 109 | (gemm_m_tile_size * sizeof(float) + input_size_13 * sizeof(scalar_t)); |
| 110 | const int32_t n_size_thread_limit = |
| 111 | output_size_13 / std::max(1, thread_num / topk_num); |
| 112 | const int32_t n_size = cpu_utils::round_down<min_w13_n_tile_size>( |
| 113 | std::min(n_size_cache_limit, n_size_thread_limit)); |
| 114 | return std::max(n_size, min_w13_n_tile_size); |
| 115 | }(); |
| 116 | |
| 117 | const int32_t w2_input_tile_size = cpu_utils::round_up<64>( |
| 118 | gemm_m_tile_size * input_size_2 * sizeof(scalar_t)); |
| 119 | // use w2 input buffer only when we need to pack input |
| 120 | const int32_t w2_input_buffer_size = |
| 121 | pack_a ? cpu_utils::round_up<64>(gemm_m_tile_size * input_size_2 * |
| 122 | sizeof(scalar_t)) |
| 123 | : 0; |
| 124 | |
| 125 | const int32_t w2_n_tile_size = [&]() { |
| 126 | const int64_t cache_size = cpu_utils::get_available_l2_size(); |
| 127 | // input tile + optional packed input + weight |
| 128 | const int32_t n_size_cache_limit = |
| 129 | (cache_size - (pack_a ? w2_input_buffer_size : w2_input_tile_size)) / |
| 130 | (input_size_2 * sizeof(scalar_t)); |
| 131 | const int32_t n_size_thread_limit = |
| 132 | output_size_2 / std::max(1, thread_num / topk_num); |
| 133 | const int32_t n_size = cpu_utils::round_down<gemm_n_tile_size>( |
| 134 | std::min(n_size_cache_limit, n_size_thread_limit)); |
| 135 | return std::max(n_size, gemm_n_tile_size); |
| 136 | }(); |
| 137 | |
| 138 | // allocate buffers |
| 139 | // Sizes/offsets below are byte counts across the whole (token_num * |
| 140 | // topk_num) batch, which for large hidden sizes and batch sizes can |
| 141 | // exceed INT32_MAX. Use int64_t throughout to avoid silently wrapping |
| 142 | // into a negative size that gets sign-extended into a huge value when |
| 143 | // passed to ScratchPadManager::realloc(size_t). |
| 144 | int64_t common_buffer_offset = 0; |
| 145 | int64_t w13_thread_buffer_offset = 0; |
| 146 | int64_t ws_thread_buffer_offset = 0; |
| 147 | |
| 148 | // common buffers |
| 149 | const int64_t token_num_per_group_buffer_size = |
| 150 | cpu_utils::round_up<64>(expert_num * sizeof(int32_t)); |
| 151 | const int64_t token_num_per_group_buffer_offset = common_buffer_offset; |
| 152 | common_buffer_offset += token_num_per_group_buffer_size; |
| 153 | |
| 154 | const int64_t cu_token_num_per_group_buffer_size = |
| 155 | cpu_utils::round_up<64>((expert_num + 1) * sizeof(int32_t)); |
| 156 | const int64_t cu_token_num_per_group_buffer_offset = common_buffer_offset; |
| 157 | common_buffer_offset += cu_token_num_per_group_buffer_size; |
| 158 | |
| 159 | const int64_t expand_token_id_buffer_size = cpu_utils::round_up<64>( |
| 160 | static_cast<int64_t>(token_num) * topk_num * sizeof(int32_t)); |
| 161 | const int64_t expand_token_id_buffer_offset = common_buffer_offset; |
| 162 | common_buffer_offset += expand_token_id_buffer_size; |
| 163 | |
| 164 | const int64_t expand_token_id_index_buffer_size = cpu_utils::round_up<64>( |
| 165 | static_cast<int64_t>(token_num) * topk_num * sizeof(int32_t)); |
| 166 | const int64_t expand_token_id_index_buffer_offset = common_buffer_offset; |
| 167 | common_buffer_offset += expand_token_id_index_buffer_size; |
| 168 | |
| 169 | const int64_t w13_gemm_output_buffer_size = |
| 170 | cpu_utils::round_up<64>(static_cast<int64_t>(token_num) * topk_num * |
| 171 | (output_size_13 / 2) * sizeof(scalar_t)); |
| 172 | const int64_t w13_gemm_output_buffer_offset = common_buffer_offset; |
| 173 | common_buffer_offset += w13_gemm_output_buffer_size; |
| 174 | |
| 175 | const int64_t w2_gemm_output_buffer_size = |
| 176 | cpu_utils::round_up<64>(static_cast<int64_t>(token_num) * topk_num * |
| 177 | output_size_2 * sizeof(float)); |
| 178 | const int64_t w2_gemm_output_buffer_offset = common_buffer_offset; |
| 179 | common_buffer_offset += w2_gemm_output_buffer_size; |
| 180 | |
| 181 | // w13 GEMM thread buffers |
| 182 | const int64_t w13_input_buffer_offset = w13_thread_buffer_offset; |
| 183 | w13_thread_buffer_offset += w13_input_buffer_size; |
| 184 | |
| 185 | const int64_t w13_output_buffer_size = cpu_utils::round_up<64>( |
| 186 | gemm_m_tile_size * w13_n_tile_size * sizeof(float)); |
| 187 | const int64_t w13_output_buffer_offset = w13_thread_buffer_offset; |
| 188 | w13_thread_buffer_offset += w13_output_buffer_size; |
| 189 | |
| 190 | const int64_t w2_input_buffer_offset = w13_thread_buffer_offset; |
| 191 | w13_thread_buffer_offset += w2_input_buffer_size; |
| 192 | |
| 193 | // Weighted sum thread buffer |
| 194 | const int64_t ws_output_buffer_size = |
| 195 | cpu_utils::round_up<64>(output_size_2 * sizeof(float)); |
| 196 | const int64_t ws_output_buffer_offset = ws_thread_buffer_offset; |
| 197 | ws_thread_buffer_offset += ws_output_buffer_size; |
| 198 | |
| 199 | const int64_t buffer_size = |
| 200 | common_buffer_offset + |
| 201 | std::max(w13_thread_buffer_offset, ws_thread_buffer_offset) * thread_num; |
| 202 | cpu_utils::ScratchPadManager::get_scratchpad_manager()->realloc(buffer_size); |
| 203 | uint8_t* common_buffer_start = |
| 204 | cpu_utils::ScratchPadManager::get_scratchpad_manager() |
| 205 | ->get_data<uint8_t>(); |
| 206 | uint8_t* thread_buffer_start = common_buffer_start + common_buffer_offset; |
| 207 | |
| 208 | int32_t* __restrict__ token_num_per_group_buffer = reinterpret_cast<int32_t*>( |
| 209 | common_buffer_start + token_num_per_group_buffer_offset); |
| 210 | int32_t* __restrict__ cu_token_num_per_group_buffer = |
| 211 | reinterpret_cast<int32_t*>(common_buffer_start + |
| 212 | cu_token_num_per_group_buffer_offset); |
| 213 | int32_t* __restrict__ expand_token_id_buffer = reinterpret_cast<int32_t*>( |
| 214 | common_buffer_start + expand_token_id_buffer_offset); |
| 215 | int32_t* __restrict__ expand_token_id_index_buffer = |
| 216 | reinterpret_cast<int32_t*>(common_buffer_start + |
| 217 | expand_token_id_index_buffer_offset); |
| 218 | |
| 219 | // prepare token-expert mappings |
| 220 | { |
| 221 | std::memset(token_num_per_group_buffer, 0, expert_num * sizeof(int32_t)); |
| 222 | for (int32_t i = 0; i < token_num * topk_num; ++i) { |
| 223 | int32_t curr_expert_id = topk_id[i]; |
| 224 | ++token_num_per_group_buffer[curr_expert_id]; |
| 225 | } |
| 226 | |
| 227 | int32_t token_num_sum = 0; |
| 228 | cu_token_num_per_group_buffer[0] = 0; |
| 229 | int32_t* token_index_buffer = cu_token_num_per_group_buffer + 1; |
| 230 | for (int32_t i = 0; i < expert_num; ++i) { |
| 231 | token_index_buffer[i] = token_num_sum; |
| 232 | token_num_sum += token_num_per_group_buffer[i]; |
| 233 | } |
| 234 | |
| 235 | for (int32_t i = 0; i < token_num; ++i) { |
| 236 | int32_t* curr_topk_id = topk_id + i * topk_num; |
| 237 | int32_t* curr_index_buffer = expand_token_id_index_buffer + i * topk_num; |
| 238 | for (int32_t j = 0; j < topk_num; ++j) { |
| 239 | int32_t curr_expert_id = curr_topk_id[j]; |
| 240 | int32_t curr_index = token_index_buffer[curr_expert_id]; |
| 241 | ++token_index_buffer[curr_expert_id]; |
| 242 | expand_token_id_buffer[curr_index] = i; |
| 243 | curr_index_buffer[j] = curr_index; |
| 244 | } |
| 245 | } |
| 246 | } |
| 247 | |
| 248 | // w13 GEMM + act |
| 249 | { |
| 250 | alignas(64) cpu_utils::Counter counter; |
| 251 | cpu_utils::Counter* counter_ptr = &counter; |
| 252 | |
| 253 | #pragma omp parallel for schedule(static, 1) |
| 254 | for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) { |
| 255 | const int32_t task_num_per_expert = |
| 256 | (output_size_13 + w13_n_tile_size - 1) / w13_n_tile_size; |
| 257 | const int32_t task_num = task_num_per_expert * expert_num; |
| 258 | |
| 259 | uint8_t* __restrict__ thread_buffer = |
| 260 | thread_buffer_start + thread_id * w13_thread_buffer_offset; |
| 261 | scalar_t* __restrict__ w13_input_buffer = |
| 262 | reinterpret_cast<scalar_t*>(thread_buffer + w13_input_buffer_offset); |
| 263 | float* __restrict__ w13_output_buffer = |
| 264 | reinterpret_cast<float*>(thread_buffer + w13_output_buffer_offset); |
| 265 | scalar_t* __restrict__ w13_gemm_output_buffer = |
| 266 | reinterpret_cast<scalar_t*>(common_buffer_start + |
| 267 | w13_gemm_output_buffer_offset); |
| 268 | |
| 269 | gemm_t gemm; |
| 270 | |
| 271 | const int32_t input_size_13_bytes = input_size_13 * sizeof(scalar_t); |
| 272 | const int32_t w13_n_group_stride = |
| 273 | gemm_t::WeightOCGroupSize * input_size_13; |
| 274 | const int32_t w13_n_tile_stride = gemm_n_tile_size * input_size_13; |
| 275 | |
| 276 | for (;;) { |
| 277 | int32_t task_id = counter_ptr->acquire_counter(); |
| 278 | if (task_id >= task_num) { |
| 279 | break; |
| 280 | } |
| 281 | |
| 282 | const int32_t curr_expert_id = task_id / task_num_per_expert; |
| 283 | const int32_t curr_output_group_id = task_id % task_num_per_expert; |
| 284 | const int32_t curr_token_num = |
| 285 | token_num_per_group_buffer[curr_expert_id]; |
| 286 | if (curr_token_num == 0) { |
| 287 | continue; |
| 288 | } |
| 289 | |
| 290 | const int32_t actual_n_tile_size = |
| 291 | std::min(w13_n_tile_size, |
| 292 | output_size_13 - curr_output_group_id * w13_n_tile_size); |
| 293 | const int32_t* __restrict__ curr_expand_token_id_buffer = |
| 294 | expand_token_id_buffer + |
| 295 | cu_token_num_per_group_buffer[curr_expert_id]; |
| 296 | scalar_t* __restrict__ curr_w13_gemm_output_buffer = |
| 297 | w13_gemm_output_buffer + |
| 298 | cu_token_num_per_group_buffer[curr_expert_id] * |
| 299 | (output_size_13 / 2) + |
| 300 | curr_output_group_id * w13_n_tile_size / 2; |
| 301 | |
| 302 | w_t* __restrict__ w13_weight_ptr_0 = nullptr; |
| 303 | w_t* __restrict__ w13_weight_ptr_1 = nullptr; |
| 304 | w_t* __restrict__ w13_bias_ptr_0 = nullptr; |
| 305 | w_t* __restrict__ w13_bias_ptr_1 = nullptr; |
| 306 | if (act_type == FusedMOEAct::SwigluOAIAndMul) { |
| 307 | // For SwigluOAIAndMul, up and down weights are interleaved |
| 308 | w13_weight_ptr_0 = |
| 309 | w13 + curr_expert_id * input_size_13 * output_size_13 + |
| 310 | curr_output_group_id * w13_n_tile_size * input_size_13; |
| 311 | w13_weight_ptr_1 = |
| 312 | w13_weight_ptr_0 + actual_n_tile_size / 2 * input_size_13; |
| 313 | if (w13_bias != nullptr) { |
| 314 | w13_bias_ptr_0 = w13_bias + curr_expert_id * output_size_13 + |
| 315 | curr_output_group_id * w13_n_tile_size; |
| 316 | w13_bias_ptr_1 = w13_bias_ptr_0 + actual_n_tile_size / 2; |
| 317 | } |
| 318 | } else { |
| 319 | w13_weight_ptr_0 = |
| 320 | w13 + curr_expert_id * input_size_13 * output_size_13 + |
| 321 | curr_output_group_id * (w13_n_tile_size / 2) * input_size_13; |
| 322 | w13_weight_ptr_1 = |
| 323 | w13_weight_ptr_0 + output_size_13 / 2 * input_size_13; |
| 324 | if (w13_bias != nullptr) { |
| 325 | w13_bias_ptr_0 = w13_bias + curr_expert_id * output_size_13 + |
| 326 | curr_output_group_id * (w13_n_tile_size / 2); |
| 327 | w13_bias_ptr_1 = w13_bias_ptr_0 + output_size_13 / 2; |
| 328 | } |
| 329 | } |
| 330 | |
| 331 | scalar_t* __restrict__ curr_w13_input_buffer = w13_input_buffer; |
| 332 | for (int32_t token_idx = 0; token_idx < curr_token_num; |
| 333 | token_idx += gemm_m_tile_size) { |
| 334 | const int32_t actual_token_num = |
| 335 | std::min(gemm_m_tile_size, curr_token_num - token_idx); |
| 336 | |
| 337 | scalar_t* __restrict__ curr_w13_gemm_input_buffer = nullptr; |
| 338 | if constexpr (pack_a) { |
| 339 | // copy and pack inputs |
| 340 | curr_w13_gemm_input_buffer = w13_input_buffer; |
| 341 | const scalar_t* w13_input_rows[gemm_m_tile_size]; |
| 342 | for (int32_t i = 0; i < actual_token_num; ++i) { |
| 343 | w13_input_rows[i] = |
| 344 | input + curr_expand_token_id_buffer[i] * input_size_13; |
| 345 | } |
| 346 | gemm_t::pack_input_from_rows(w13_input_rows, |
| 347 | curr_w13_gemm_input_buffer, |
| 348 | actual_token_num, input_size_13); |
| 349 | curr_expand_token_id_buffer += actual_token_num; |
| 350 | } else { |
| 351 | // copy inputs |
| 352 | curr_w13_gemm_input_buffer = curr_w13_input_buffer; |
| 353 | scalar_t* __restrict__ curr_w13_input_buffer_iter = |
| 354 | curr_w13_input_buffer; |
| 355 | for (int32_t i = 0; i < actual_token_num; ++i) { |
| 356 | const int32_t curr_token_id = curr_expand_token_id_buffer[i]; |
| 357 | int8_t* __restrict__ curr_input_iter = reinterpret_cast<int8_t*>( |
| 358 | input + curr_token_id * input_size_13); |
| 359 | int8_t* __restrict__ curr_output_iter = |
| 360 | reinterpret_cast<int8_t*>(curr_w13_input_buffer_iter); |
| 361 | int32_t j = 0; |
| 362 | for (; j < input_size_13_bytes - 64; j += 64) { |
| 363 | vec_op::INT8Vec64 vec(curr_input_iter); |
| 364 | vec.save(curr_output_iter); |
| 365 | curr_input_iter += 64; |
| 366 | curr_output_iter += 64; |
| 367 | } |
| 368 | vec_op::INT8Vec64 vec(curr_input_iter); |
| 369 | vec.save(curr_output_iter, input_size_13_bytes - j); |
| 370 | |
| 371 | // update |
| 372 | curr_w13_input_buffer_iter += input_size_13; |
| 373 | } |
| 374 | // update |
| 375 | curr_expand_token_id_buffer += actual_token_num; |
| 376 | } |
| 377 | |
| 378 | // gemm + act |
| 379 | { |
| 380 | scalar_t* __restrict__ w13_weight_ptr_0_iter = w13_weight_ptr_0; |
| 381 | scalar_t* __restrict__ w13_weight_ptr_1_iter = w13_weight_ptr_1; |
| 382 | scalar_t* __restrict__ w13_bias_ptr_0_iter = w13_bias_ptr_0; |
| 383 | scalar_t* __restrict__ w13_bias_ptr_1_iter = w13_bias_ptr_1; |
| 384 | float* __restrict__ w13_output_buffer_0_iter = w13_output_buffer; |
| 385 | float* __restrict__ w13_output_buffer_1_iter = |
| 386 | w13_output_buffer + actual_n_tile_size / 2; |
| 387 | for (int32_t i = 0; i < actual_n_tile_size; |
| 388 | i += min_w13_n_tile_size) { |
| 389 | gemm.gemm(curr_w13_gemm_input_buffer, w13_weight_ptr_0_iter, |
| 390 | w13_output_buffer_0_iter, actual_token_num, |
| 391 | input_size_13, input_size_13, w13_n_group_stride, |
| 392 | actual_n_tile_size, false); |
| 393 | |
| 394 | if (w13_bias != nullptr) { |
| 395 | cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>( |
| 396 | w13_output_buffer_0_iter, w13_output_buffer_0_iter, |
| 397 | w13_bias_ptr_0_iter, actual_token_num, actual_n_tile_size, |
| 398 | actual_n_tile_size); |
| 399 | w13_bias_ptr_0_iter += gemm_n_tile_size; |
| 400 | } |
| 401 | |
| 402 | gemm.gemm(curr_w13_gemm_input_buffer, w13_weight_ptr_1_iter, |
| 403 | w13_output_buffer_1_iter, actual_token_num, |
| 404 | input_size_13, input_size_13, w13_n_group_stride, |
| 405 | actual_n_tile_size, false); |
| 406 | |
| 407 | if (w13_bias != nullptr) { |
| 408 | cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>( |
| 409 | w13_output_buffer_1_iter, w13_output_buffer_1_iter, |
| 410 | w13_bias_ptr_1_iter, actual_token_num, actual_n_tile_size, |
| 411 | actual_n_tile_size); |
| 412 | w13_bias_ptr_1_iter += gemm_n_tile_size; |
| 413 | } |
| 414 | |
| 415 | // update |
| 416 | w13_weight_ptr_0_iter += w13_n_tile_stride; |
| 417 | w13_weight_ptr_1_iter += w13_n_tile_stride; |
| 418 | w13_output_buffer_0_iter += gemm_n_tile_size; |
| 419 | w13_output_buffer_1_iter += gemm_n_tile_size; |
| 420 | } |
| 421 | |
| 422 | apply_gated_act(act_type, w13_output_buffer, |
| 423 | curr_w13_gemm_output_buffer, actual_token_num, |
| 424 | actual_n_tile_size, actual_n_tile_size, |
| 425 | output_size_13 / 2); |
| 426 | |
| 427 | // update |
| 428 | curr_w13_gemm_output_buffer += |
| 429 | gemm_m_tile_size * (output_size_13 / 2); |
| 430 | } |
| 431 | } |
| 432 | } |
| 433 | } |
| 434 | } |
| 435 | |
| 436 | // w2 GEMM |
| 437 | { |
| 438 | alignas(64) cpu_utils::Counter counter; |
| 439 | cpu_utils::Counter* counter_ptr = &counter; |
| 440 | |
| 441 | #pragma omp parallel for schedule(static, 1) |
| 442 | for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) { |
| 443 | const int32_t task_num_per_expert = |
| 444 | (output_size_2 + w2_n_tile_size - 1) / w2_n_tile_size; |
| 445 | const int32_t task_num = task_num_per_expert * expert_num; |
| 446 | scalar_t* __restrict__ w13_gemm_output_buffer = |
| 447 | reinterpret_cast<scalar_t*>(common_buffer_start + |
| 448 | w13_gemm_output_buffer_offset); |
| 449 | float* __restrict__ w2_gemm_output_buffer = reinterpret_cast<float*>( |
| 450 | common_buffer_start + w2_gemm_output_buffer_offset); |
| 451 | |
| 452 | gemm_t gemm; |
| 453 | |
| 454 | const int32_t w2_n_tile_stride = gemm_n_tile_size * input_size_2; |
| 455 | const int32_t w2_n_group_stride = |
| 456 | gemm_t::WeightOCGroupSize * input_size_2; |
| 457 | |
| 458 | for (;;) { |
| 459 | int32_t task_id = counter_ptr->acquire_counter(); |
| 460 | if (task_id >= task_num) { |
| 461 | break; |
| 462 | } |
| 463 | |
| 464 | const int32_t curr_expert_id = task_id / task_num_per_expert; |
| 465 | const int32_t curr_output_group_id = task_id % task_num_per_expert; |
| 466 | const int32_t curr_token_num = |
| 467 | token_num_per_group_buffer[curr_expert_id]; |
| 468 | if (curr_token_num == 0) { |
| 469 | continue; |
| 470 | } |
| 471 | |
| 472 | const int32_t actual_n_tile_size = |
| 473 | std::min(w2_n_tile_size, |
| 474 | output_size_2 - curr_output_group_id * w2_n_tile_size); |
| 475 | scalar_t* __restrict__ curr_w13_gemm_output_buffer = |
| 476 | w13_gemm_output_buffer + |
| 477 | cu_token_num_per_group_buffer[curr_expert_id] * input_size_2; |
| 478 | float* __restrict__ curr_w2_gemm_output_buffer = |
| 479 | w2_gemm_output_buffer + |
| 480 | cu_token_num_per_group_buffer[curr_expert_id] * output_size_2 + |
| 481 | curr_output_group_id * w2_n_tile_size; |
| 482 | scalar_t* __restrict__ w2_weight_ptr = |
| 483 | w2 + curr_expert_id * output_size_2 * input_size_2 + |
| 484 | curr_output_group_id * w2_n_tile_size * input_size_2; |
| 485 | scalar_t* __restrict__ w2_bias_ptr = nullptr; |
| 486 | if (w2_bias != nullptr) { |
| 487 | w2_bias_ptr = w2_bias + curr_expert_id * output_size_2 + |
| 488 | curr_output_group_id * w2_n_tile_size; |
| 489 | } |
| 490 | |
| 491 | for (int32_t token_idx = 0; token_idx < curr_token_num; |
| 492 | token_idx += gemm_m_tile_size) { |
| 493 | const int32_t actual_token_num = |
| 494 | std::min(gemm_m_tile_size, curr_token_num - token_idx); |
| 495 | scalar_t* __restrict__ curr_w2_gemm_input_buffer = |
| 496 | curr_w13_gemm_output_buffer; |
| 497 | if constexpr (pack_a) { |
| 498 | uint8_t* __restrict__ thread_buffer = |
| 499 | thread_buffer_start + thread_id * w13_thread_buffer_offset; |
| 500 | scalar_t* __restrict__ w2_input_buffer = |
| 501 | reinterpret_cast<scalar_t*>(thread_buffer + |
| 502 | w2_input_buffer_offset); |
| 503 | curr_w2_gemm_input_buffer = w2_input_buffer; |
| 504 | const scalar_t* w2_input_rows[gemm_m_tile_size]; |
| 505 | for (int32_t i = 0; i < actual_token_num; ++i) { |
| 506 | w2_input_rows[i] = curr_w13_gemm_output_buffer + i * input_size_2; |
| 507 | } |
| 508 | gemm_t::pack_input_from_rows(w2_input_rows, |
| 509 | curr_w2_gemm_input_buffer, |
| 510 | actual_token_num, input_size_2); |
| 511 | } |
| 512 | |
| 513 | scalar_t* __restrict__ w2_weight_ptr_iter = w2_weight_ptr; |
| 514 | scalar_t* __restrict__ w2_bias_ptr_iter = w2_bias_ptr; |
| 515 | float* __restrict__ curr_w2_gemm_output_buffer_iter = |
| 516 | curr_w2_gemm_output_buffer; |
| 517 | for (int32_t i = 0; i < actual_n_tile_size; i += gemm_n_tile_size) { |
| 518 | gemm.gemm(curr_w2_gemm_input_buffer, w2_weight_ptr_iter, |
| 519 | curr_w2_gemm_output_buffer_iter, actual_token_num, |
| 520 | input_size_2, input_size_2, w2_n_group_stride, |
| 521 | output_size_2, false); |
| 522 | |
| 523 | if (w2_bias != nullptr) { |
| 524 | cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>( |
| 525 | curr_w2_gemm_output_buffer_iter, |
| 526 | curr_w2_gemm_output_buffer_iter, w2_bias_ptr_iter, |
| 527 | actual_token_num, output_size_2, output_size_2); |
| 528 | w2_bias_ptr_iter += gemm_n_tile_size; |
| 529 | } |
| 530 | |
| 531 | w2_weight_ptr_iter += w2_n_tile_stride; |
| 532 | curr_w2_gemm_output_buffer_iter += gemm_n_tile_size; |
| 533 | } |
| 534 | |
| 535 | // update |
| 536 | curr_w13_gemm_output_buffer += gemm_m_tile_size * input_size_2; |
| 537 | curr_w2_gemm_output_buffer += gemm_m_tile_size * output_size_2; |
| 538 | } |
| 539 | } |
| 540 | } |
| 541 | } |
| 542 | |
| 543 | // weighted sum |
| 544 | { |
| 545 | alignas(64) cpu_utils::Counter counter; |
| 546 | cpu_utils::Counter* counter_ptr = &counter; |
| 547 | |
| 548 | #pragma omp parallel for schedule(static, 1) |
| 549 | for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) { |
| 550 | const int32_t task_num = token_num; |
| 551 | uint8_t* __restrict__ thread_buffer = |
| 552 | thread_buffer_start + thread_id * ws_thread_buffer_offset; |
| 553 | float* __restrict__ ws_output_buffer = |
| 554 | reinterpret_cast<float*>(thread_buffer + ws_output_buffer_offset); |
| 555 | float* __restrict__ w2_gemm_output_buffer = reinterpret_cast<float*>( |
| 556 | common_buffer_start + w2_gemm_output_buffer_offset); |
| 557 | |
| 558 | for (;;) { |
| 559 | int32_t task_id = counter_ptr->acquire_counter(); |
| 560 | if (task_id >= task_num) { |
| 561 | break; |
| 562 | } |
| 563 | |
| 564 | int32_t token_id = task_id; |
| 565 | int32_t* __restrict__ curr_expand_token_id_index_buffer = |
| 566 | expand_token_id_index_buffer + token_id * topk_num; |
| 567 | float* __restrict__ curr_weight = topk_weights + token_id * topk_num; |
| 568 | scalar_t* __restrict__ curr_output_buffer = |
| 569 | output + token_id * output_size_2; |
| 570 | |
| 571 | if (skip_weighted) { |
| 572 | // Only for topk_num == 1 |
| 573 | *curr_weight = 1.0f; |
| 574 | } |
| 575 | |
| 576 | if (topk_num > 1) { |
| 577 | { |
| 578 | int32_t w2_output_idx = curr_expand_token_id_index_buffer[0]; |
| 579 | float* __restrict__ w2_output_iter = |
| 580 | w2_gemm_output_buffer + w2_output_idx * output_size_2; |
| 581 | float* __restrict__ ws_output_buffer_iter = ws_output_buffer; |
| 582 | vec_op::FP32Vec16 weight_vec(curr_weight[0]); |
| 583 | for (int32_t i = 0; i < output_size_2; i += 16) { |
| 584 | vec_op::FP32Vec16 vec(w2_output_iter); |
| 585 | vec = vec * weight_vec; |
| 586 | vec.save(ws_output_buffer_iter); |
| 587 | |
| 588 | // update |
| 589 | w2_output_iter += 16; |
| 590 | ws_output_buffer_iter += 16; |
| 591 | } |
| 592 | } |
| 593 | |
| 594 | { |
| 595 | for (int32_t idx = 1; idx < topk_num - 1; ++idx) { |
| 596 | int32_t w2_output_idx = curr_expand_token_id_index_buffer[idx]; |
| 597 | float* __restrict__ w2_output_iter = |
| 598 | w2_gemm_output_buffer + w2_output_idx * output_size_2; |
| 599 | float* __restrict__ ws_output_buffer_iter = ws_output_buffer; |
| 600 | vec_op::FP32Vec16 weight_vec(curr_weight[idx]); |
| 601 | for (int32_t i = 0; i < output_size_2; i += 16) { |
| 602 | vec_op::FP32Vec16 vec(w2_output_iter); |
| 603 | vec_op::FP32Vec16 sum(ws_output_buffer_iter); |
| 604 | sum = sum + vec * weight_vec; |
| 605 | sum.save(ws_output_buffer_iter); |
| 606 | |
| 607 | // update |
| 608 | w2_output_iter += 16; |
| 609 | ws_output_buffer_iter += 16; |
| 610 | } |
| 611 | } |
| 612 | } |
| 613 | |
| 614 | { |
| 615 | int32_t idx = topk_num - 1; |
| 616 | int32_t w2_output_idx = curr_expand_token_id_index_buffer[idx]; |
| 617 | float* __restrict__ w2_output_iter = |
| 618 | w2_gemm_output_buffer + w2_output_idx * output_size_2; |
| 619 | float* __restrict__ ws_output_buffer_iter = ws_output_buffer; |
| 620 | scalar_t* __restrict__ curr_output_buffer_iter = curr_output_buffer; |
| 621 | vec_op::FP32Vec16 weight_vec(curr_weight[idx]); |
| 622 | for (int32_t i = 0; i < output_size_2; i += 16) { |
| 623 | vec_op::FP32Vec16 vec(w2_output_iter); |
| 624 | vec_op::FP32Vec16 sum(ws_output_buffer_iter); |
| 625 | sum = sum + vec * weight_vec; |
| 626 | scalar_vec_t out_vec(sum); |
| 627 | out_vec.save(curr_output_buffer_iter); |
| 628 | |
| 629 | // update |
| 630 | w2_output_iter += 16; |
| 631 | ws_output_buffer_iter += 16; |
| 632 | curr_output_buffer_iter += 16; |
| 633 | } |
| 634 | } |
| 635 | } else { |
| 636 | int32_t w2_output_idx = curr_expand_token_id_index_buffer[0]; |
| 637 | float* __restrict__ w2_output_iter = |
| 638 | w2_gemm_output_buffer + w2_output_idx * output_size_2; |
| 639 | scalar_t* __restrict__ curr_output_buffer_iter = curr_output_buffer; |
| 640 | vec_op::FP32Vec16 weight_vec(curr_weight[0]); |
| 641 | for (int32_t i = 0; i < output_size_2; i += 16) { |
| 642 | vec_op::FP32Vec16 vec(w2_output_iter); |
| 643 | vec = vec * weight_vec; |
| 644 | scalar_vec_t out_vec(vec); |
| 645 | out_vec.save(curr_output_buffer_iter); |
| 646 | |
| 647 | // update |
| 648 | w2_output_iter += 16; |
| 649 | curr_output_buffer_iter += 16; |
| 650 | } |
| 651 | } |
| 652 | } |
| 653 | } |
| 654 | } |
| 655 | } |
| 656 | |
| 657 | } // namespace |
| 658 | |
| 659 | void prepack_moe_weight( |
| 660 | const torch::Tensor& weight, // [expert_num, output_size, input_size] |
| 661 | torch::Tensor& packed_weight, const std::string& isa) { |
| 662 | TORCH_CHECK(weight.is_contiguous()); |
| 663 | const int32_t expert_num = weight.size(0); |
| 664 | const int32_t output_size = weight.size(1); |
| 665 | const int32_t input_size = weight.size(2); |
| 666 | TORCH_CHECK_EQ(output_size % 32, 0); |
| 667 | const int64_t expert_stride = weight.stride(0); |
| 668 | cpu_utils::ISA isa_type = cpu_utils::get_isa(isa); |
| 669 | |
| 670 | VLLM_DISPATCH_FLOATING_TYPES( |
| 671 | weight.scalar_type(), "prepack_moe_weight", [&]() { |
| 672 | CPU_ISA_DISPATCH_IMPL(isa_type, [&]() { |
| 673 | scalar_t* weight_ptr = weight.data_ptr<scalar_t>(); |
| 674 | scalar_t* packed_weight_ptr = packed_weight.data_ptr<scalar_t>(); |
| 675 | prepack_moe_weight_impl<scalar_t, gemm_t>( |
| 676 | weight_ptr, packed_weight_ptr, expert_num, output_size, |
| 677 | input_size, expert_stride); |
| 678 | }); |
| 679 | }); |
| 680 | } |
| 681 | |
| 682 | void cpu_fused_moe( |
| 683 | torch::Tensor& output, // [token_num, output_size_2] |
| 684 | const torch::Tensor& input, // [token_num, input_size_13] |
| 685 | const torch::Tensor& |
| 686 | w13, // [expert_num, output_size_13, input_size_13], packed |
| 687 | const torch::Tensor& |
| 688 | w2, // [expert_num, output_size_2, input_size_2], packed |
| 689 | const std::optional<torch::Tensor>& |
| 690 | w13_bias, // [expert_num, output_size_13] |
| 691 | const std::optional<torch::Tensor>& w2_bias, // [expert_num, output_size_2] |
| 692 | const torch::Tensor& topk_weights, // [token_num, k], float32 |
| 693 | const torch::Tensor& topk_id, // [token_num, k], int32 |
| 694 | const bool skip_weighted, const std::string& act, const std::string& isa) { |
| 695 | const int32_t token_num = input.size(0); |
| 696 | const int32_t input_size_13 = input.size(1); |
| 697 | const int64_t input_stride = input.stride(0); |
| 698 | TORCH_CHECK_EQ(input_stride, input_size_13); |
| 699 | const int32_t expert_num = w13.size(0); |
| 700 | const int32_t output_size_13 = w13.size(1); |
| 701 | const int32_t input_size_2 = w2.size(2); |
| 702 | const int32_t output_size_2 = w2.size(1); |
| 703 | const int32_t topk_num = topk_id.size(1); |
| 704 | const FusedMOEAct act_type = cpu_fused_moe_utils::get_act_type(act); |
| 705 | cpu_utils::ISA isa_type = cpu_utils::get_isa(isa); |
| 706 | TORCH_CHECK(!skip_weighted || topk_num == 1, |
| 707 | "skip_weighted is only supported for topk=1 on CPU"); |
| 708 | |
| 709 | VLLM_DISPATCH_FLOATING_TYPES(w13.scalar_type(), "cpu_fused_moe", [&]() { |
| 710 | CPU_ISA_DISPATCH_IMPL(isa_type, [&]() { |
| 711 | fused_moe_impl<scalar_t, scalar_t, gemm_t>( |
| 712 | output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), |
| 713 | w13.data_ptr<scalar_t>(), w2.data_ptr<scalar_t>(), |
| 714 | w13_bias.has_value() ? w13_bias->data_ptr<scalar_t>() : nullptr, |
| 715 | w2_bias.has_value() ? w2_bias->data_ptr<scalar_t>() : nullptr, |
| 716 | topk_weights.data_ptr<float>(), topk_id.data_ptr<int32_t>(), act_type, |
| 717 | token_num, expert_num, topk_num, input_size_13, output_size_13, |
| 718 | input_size_2, output_size_2, skip_weighted); |
| 719 | }); |
| 720 | }); |
| 721 | } |