Skip to content

[CPU] Refactor CPU fused MOE #30531

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jikunshang merged 5 commits into from
Dec 18, 2025
Merged

[CPU] Refactor CPU fused MOE #30531

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
8c955cc
update
bigPYJ1151 Dec 11, 2025
711e647
fix
bigPYJ1151 Dec 12, 2025
12451f7
fix
bigPYJ1151 Dec 12, 2025
3218b41
fix
bigPYJ1151 Dec 15, 2025
fde02ce
Merge branch 'main' into grouped_gemm
bigPYJ1151 Dec 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .buildkite/scripts/hardware_ci/run-cpu-test.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ function cpu_tests() {
docker exec cpu-test-"$NUMA_NODE" bash -c "
set -e
pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py

@fadara01 fadara01 Dec 12, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Given you enabled this through ISA::Vec which is CPU agnostic, would it be a good idea to enable this for all run-cpu-tests, i.e:

  • .buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
  • .buildkite/scripts/hardware_ci/run-cpu-test-ppc64le.sh
  • .buildkite/scripts/hardware_ci/run-cpu-test-s390x.sh

@bigPYJ1151 bigPYJ1151 Dec 15, 2025

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not exactly. This kerenl is only avaliable on AVX512 because of some new vec_op. Other platforms need a bit further change and verification.

pytest -x -v -s tests/kernels/test_onednn.py"

# Run basic model test
Expand Down
4 changes: 2 additions & 2 deletions cmake/cpu_extension.cmake
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -330,7 +330,7 @@ if ((AVX512_FOUND AND NOT AVX512_DISABLED) OR (ASIMD_FOUND AND NOT APPLE_SILICON
PUBLIC ${oneDNN_BINARY_DIR}/include
PRIVATE ${oneDNN_SOURCE_DIR}/src
)
target_link_libraries(dnnl_ext dnnl)
target_link_libraries(dnnl_ext dnnl torch)
target_compile_options(dnnl_ext PRIVATE ${CXX_COMPILE_FLAGS} -fPIC)
list(APPEND LIBS dnnl_ext)
set(USE_ONEDNN ON)
Expand Down Expand Up @@ -358,13 +358,13 @@ set(VLLM_EXT_SRC
"csrc/cpu/pos_encoding.cpp"
"csrc/moe/dynamic_4bit_int_moe_cpu.cpp"
"csrc/cpu/cpu_attn.cpp"
"csrc/cpu/scratchpad_manager.cpp"
"csrc/cpu/torch_bindings.cpp")

if (AVX512_FOUND AND NOT AVX512_DISABLED)
set(VLLM_EXT_SRC
"csrc/cpu/shm.cpp"
"csrc/cpu/cpu_wna16.cpp"
"csrc/cpu/cpu_fused_moe.cpp"

@fadara01 fadara01 Dec 14, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

See my comments above about isa::vec being platform agnostic.

${VLLM_EXT_SRC})
if (ENABLE_AVX512BF16 AND ENABLE_AVX512VNNI)
set(VLLM_EXT_SRC
Expand Down
10 changes: 5 additions & 5 deletions csrc/cpu/cpu_attn_macros.h → csrc/cpu/cpu_arch_macros.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
#ifndef CPU_ATTN_MACROS_H
#define CPU_ATTN_MACROS_H
#ifndef CPU_ARCH_MACROS_H
#define CPU_ARCH_MACROS_H

// x86_64
#ifdef __x86_64__
Expand All @@ -26,7 +26,7 @@
_mm512_castsi512_ps(_mm512_set1_epi32(0x42b17218)); \
const __m512i vec_127 = _mm512_set1_epi32(0x0000007f); \
const int n_mantissa_bits = 23; \
auto fast_exp = [&](vec_op::FP32Vec16& vec) __attribute__(( \
auto fast_exp = [&](const vec_op::FP32Vec16& vec) __attribute__(( \
Comment thread
bigPYJ1151 marked this conversation as resolved.
always_inline)) { \
__m512 values = vec.reg; \
auto less_ln_flt_min_mask = \
Expand Down Expand Up @@ -98,7 +98,7 @@
poly = vbslq_f32(hi_mask, inf, poly); \
return vbslq_f32(lo_mask, zero, poly); \
}; \
auto fast_exp = [&](vec_op::FP32Vec16& vec) \
auto fast_exp = [&](const vec_op::FP32Vec16& vec) \
__attribute__((always_inline)) { \
float32x4x4_t result; \
result.val[0] = neon_expf(vec.reg.val[0]); \
Expand All @@ -110,4 +110,4 @@

#endif // __aarch64__

#endif
#endif
41 changes: 9 additions & 32 deletions csrc/cpu/cpu_attn_impl.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,10 +8,8 @@
#include <sys/sysctl.h>
#endif

#include "cpu_types.hpp"
#include "scratchpad_manager.h"
#include "cpu_attn_macros.h"
#include "utils.hpp"
#include "cpu/cpu_arch_macros.h"
#include "cpu/utils.hpp"

namespace cpu_attention {
enum class ISA { AMX, VEC, VEC16, NEON };
Expand Down Expand Up @@ -378,12 +376,13 @@ class AttentionScheduler {

static constexpr int32_t MaxQTileIterNum = 128;

AttentionScheduler() : available_cache_size_(get_available_l2_size()) {}
AttentionScheduler()
: available_cache_size_(cpu_utils::get_available_l2_size()) {}

torch::Tensor schedule(const ScheduleInput& input) const {
const bool casual = input.casual;
const int32_t thread_num = omp_get_max_threads();
const int64_t cache_size = get_available_l2_size();
const int64_t cache_size = cpu_utils::get_available_l2_size();
const int32_t max_num_q_per_iter = input.max_num_q_per_iter;
const int32_t kv_len_alignment = input.kv_block_alignment;
int32_t q_head_per_kv = input.num_heads_q / input.num_heads_kv;
Expand Down Expand Up @@ -659,15 +658,15 @@ class AttentionScheduler {
metadata_ptr->thread_num +
metadata_ptr->reduction_scratchpad_size_per_kv_head *
(use_gqa ? input.num_heads_kv : input.num_heads_q);
DNNLScratchPadManager::get_dnnl_scratchpad_manager()->realloc(
cpu_utils::ScratchPadManager::get_scratchpad_manager()->realloc(
scratchpad_size);

// metadata_ptr->print();

// test out of boundary access
// {
// float* cache_ptr =
// DNNLScratchPadManager::get_dnnl_scratchpad_manager()->get_data<float>();
// cpu_utils::ScratchPadManager::getl_scratchpad_manager()->get_data<float>();
// for (int64_t i = 0; i < scratchpad_size / sizeof(float); ++i) {
// cache_ptr[i] = std::numeric_limits<float>::quiet_NaN();
// }
Expand Down Expand Up @@ -749,27 +748,6 @@ class AttentionScheduler {
return std::max(rounded_tile_size, round_size);
}

static int64_t get_available_l2_size() {
static int64_t size = []() {
#if defined(__APPLE__)
// macOS doesn't have _SC_LEVEL2_CACHE_SIZE. Use sysctlbyname.
int64_t l2_cache_size = 0;
size_t len = sizeof(l2_cache_size);
if (sysctlbyname("hw.l2cachesize", &l2_cache_size, &len, NULL, 0) == 0 &&
l2_cache_size > 0) {
return l2_cache_size >> 1; // use 50% of L2 cache
}
// Fallback if sysctlbyname fails
return 128LL * 1024 >> 1; // use 50% of 128KB
#else
long l2_cache_size = sysconf(_SC_LEVEL2_CACHE_SIZE);
TORCH_CHECK_NE(l2_cache_size, -1);
return l2_cache_size >> 1; // use 50% of L2 cache
#endif
}();
return size;
}

private:
int64_t available_cache_size_;
};
Expand Down Expand Up @@ -1402,7 +1380,7 @@ class AttentionMainLoop {

// init buffers
void* scratchpad_ptr =
DNNLScratchPadManager::get_dnnl_scratchpad_manager()
cpu_utils::ScratchPadManager::get_scratchpad_manager()
->get_data<void>();
AttentionScratchPad buffer_manager(thread_id, metadata, scratchpad_ptr);

Expand All @@ -1422,8 +1400,7 @@ class AttentionMainLoop {
}
}

const int64_t available_cache_size =
AttentionScheduler::get_available_l2_size();
const int64_t available_cache_size = cpu_utils::get_available_l2_size();
const int32_t default_tile_size =
AttentionScheduler::calcu_default_tile_size(
available_cache_size, head_dim, sizeof(kv_cache_t),
Expand Down
Loading