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)