From bf33e8e68f1cac4d97340c04fc6276d543779299 Mon Sep 17 00:00:00 2001 From: Shaojie Li Date: Sat, 30 May 2026 20:35:37 +0000 Subject: [PATCH] [Perf][CPU] Restore vectorized silu_and_mul on CPU; handle uneven hidden dim PR #22145 (Aug 2025) routed all CPU SiluAndMul calls through PyTorch native to dodge a TORCH_CHECK that fired on Qwen2.5 with non-aligned hidden sizes. On AMD EPYC 9R14 the eager-mode "very slight" drop is 4-8x on prefill. Three fixes: 1. csrc/cpu/activation.cpp -- replace the TORCH_CHECK(d % VEC_ELEM_NUM == 0) with a scalar tail loop that handles the remainder elements. 2. csrc/cpu/cpu_types_x86.hpp -- BF16Vec8 and FP16Vec8 ::save(void*) were doing aligned MOVDQA (*reinterpret_cast<__m128i*>(ptr) = reg). Switch to _mm_storeu_si128. The aligned variant segfaults whenever the per-token write offset (output + i*d) is not 16-byte aligned -- any odd d once i >= 1. 3. vllm/model_executor/layers/activation.py -- with the kernel fixed, restore the C++ dispatch on CPU for SiluAndMul. The kernel template also serves gelu_and_mul, gelu_tanh_and_mul, gelu_new, gelu_fast, gelu_quick. Those were never explicitly Python-fallbacked, so they relied on the CustomOp default (forward_cpu -> forward_native) and any caller that used self.op directly on CPU would have crashed on non-aligned d. This fix makes the kernel correct for them too. Microbench, AMD EPYC 9R14 (Zen 4, AVX-512+BF16, no AMX), 64 threads, bf16: shape native (us) cpp (us) speedup Llama-3-8B prefill M=4096 28387 6068 4.68x Llama-3-70B prefill M=4096 17219 2147 8.02x D=18944 M=4096 38363 4893 7.84x D=18943 M=4096 (uneven) 40378 4974 8.12x D=14329 M=4096 (uneven) 29467 3711 7.94x Llama-3-8B decode M=16 62 49 1.28x Llama-3-70B decode M=16 62 36 1.74x C++ runs the activation in fp32 then casts back; on D=14329 the C++ output is closer to the fp32 reference than forward_native (max abs diff 0.0156 vs 0.0330). Test plan: pytest tests/kernels/core/test_cpu_activation.py # 96 passed, 16 skipped (Arm-only paths) D parametrization is extended with 513 and 2049 to exercise the new scalar-tail path. Signed-off-by: Shaojie Li --- csrc/cpu/activation.cpp | 46 ++++++++++++++++++++--- csrc/cpu/cpu_types_x86.hpp | 4 +- tests/kernels/core/test_cpu_activation.py | 5 ++- vllm/model_executor/layers/activation.py | 11 ++++-- 4 files changed, 54 insertions(+), 12 deletions(-) diff --git a/csrc/cpu/activation.cpp b/csrc/cpu/activation.cpp index 039b8d5c30d4..d6ea4f948b10 100644 --- a/csrc/cpu/activation.cpp +++ b/csrc/cpu/activation.cpp @@ -8,16 +8,20 @@ void activation_kernel(int num_tokens, int d, scalar_t* __restrict__ input, using scalar_vec_t = vec_op::vec_t; constexpr int VEC_ELEM_NUM = scalar_vec_t::get_elem_num(); - TORCH_CHECK(d % VEC_ELEM_NUM == 0); + // Round d down to a vector boundary; the remainder runs through a scalar + // tail below. Pre-#22145 this kernel required d % VEC_ELEM_NUM == 0, which + // crashed on models with non-aligned hidden sizes (e.g. Qwen2.5). + const int d_vec = (d / VEC_ELEM_NUM) * VEC_ELEM_NUM; + const int tail = d - d_vec; #pragma omp parallel for for (int i = 0; i < num_tokens; ++i) { - for (int j = 0; j < d; j += VEC_ELEM_NUM) { - int start = i * d; - if constexpr (is_gated) { - start *= 2; - } + int start = i * d; + if constexpr (is_gated) { + start *= 2; + } + for (int j = 0; j < d_vec; j += VEC_ELEM_NUM) { const scalar_vec_t x(input + start + j); const vec_op::FP32Vec8 f32_x(x); vec_op::FP32Vec8 f32_ans = func(f32_x); @@ -31,6 +35,36 @@ void activation_kernel(int num_tokens, int d, scalar_t* __restrict__ input, const scalar_vec_t result(f32_ans); result.save(output + i * d + j); } + + // Scalar tail: copy the remaining < VEC_ELEM_NUM elements through a + // zero-padded vector to reuse the same activation function. Cheap + // because tail runs at most once per token. + if (tail > 0) { + alignas(64) scalar_t x_buf[VEC_ELEM_NUM] = {}; + for (int k = 0; k < tail; ++k) { + x_buf[k] = input[start + d_vec + k]; + } + const scalar_vec_t x(x_buf); + const vec_op::FP32Vec8 f32_x(x); + vec_op::FP32Vec8 f32_ans = func(f32_x); + + if constexpr (is_gated) { + alignas(64) scalar_t y_buf[VEC_ELEM_NUM] = {}; + for (int k = 0; k < tail; ++k) { + y_buf[k] = input[start + d + d_vec + k]; + } + const scalar_vec_t y(y_buf); + const vec_op::FP32Vec8 f32_y(y); + f32_ans = f32_y * f32_ans; + } + + alignas(64) scalar_t out_buf[VEC_ELEM_NUM]; + const scalar_vec_t result(f32_ans); + result.save(out_buf); + for (int k = 0; k < tail; ++k) { + output[i * d + d_vec + k] = out_buf[k]; + } + } } } diff --git a/csrc/cpu/cpu_types_x86.hpp b/csrc/cpu/cpu_types_x86.hpp index 396b9b7e041f..0d508bfcac73 100644 --- a/csrc/cpu/cpu_types_x86.hpp +++ b/csrc/cpu/cpu_types_x86.hpp @@ -99,7 +99,7 @@ struct FP16Vec8 : public Vec { explicit FP16Vec8(const FP32Vec8&); - void save(void* ptr) const { *reinterpret_cast<__m128i*>(ptr) = reg; } + void save(void* ptr) const { _mm_storeu_si128((__m128i*)ptr, reg); } }; struct FP16Vec16 : public Vec { @@ -146,7 +146,7 @@ struct BF16Vec8 : public Vec { explicit BF16Vec8(const FP32Vec8&); - void save(void* ptr) const { *reinterpret_cast<__m128i*>(ptr) = reg; } + void save(void* ptr) const { _mm_storeu_si128((__m128i*)ptr, reg); } }; struct BF16Vec16 : public Vec { diff --git a/tests/kernels/core/test_cpu_activation.py b/tests/kernels/core/test_cpu_activation.py index 40b5f0454683..09f4ec19169a 100644 --- a/tests/kernels/core/test_cpu_activation.py +++ b/tests/kernels/core/test_cpu_activation.py @@ -23,7 +23,10 @@ DTYPES = [torch.bfloat16, torch.float32] NUM_TOKENS = [7, 83] -D = [512, 2048] +# 513 / 2049 exercise the scalar tail when d is not a multiple of the SIMD +# width (regression check for the original TORCH_CHECK that fired on Qwen2.5 +# with non-aligned hidden sizes). +D = [512, 513, 2048, 2049] SEEDS = [0] diff --git a/vllm/model_executor/layers/activation.py b/vllm/model_executor/layers/activation.py index ddad6801adc4..453191b60794 100644 --- a/vllm/model_executor/layers/activation.py +++ b/vllm/model_executor/layers/activation.py @@ -129,10 +129,12 @@ class SiluAndMul(CustomOp): def __init__(self, *, compile_native: bool = True): super().__init__(compile_native=compile_native) - if current_platform.is_cuda_alike() or current_platform.is_xpu(): + if ( + current_platform.is_cuda_alike() + or current_platform.is_xpu() + or current_platform.is_cpu() + ): self.op = torch.ops._C.silu_and_mul - elif current_platform.is_cpu(): - self._forward_method = self.forward_native @staticmethod def forward_native(x: torch.Tensor) -> torch.Tensor: @@ -147,6 +149,9 @@ def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: self.op(out, x) return out + def forward_cpu(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_cuda(x) + def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: return self.forward_cuda(x)