[Bugfix][Perf][CPU] Restore vectorized silu_and_mul on CPU; handle uneven hidden dim

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<scalar_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];
+      }
+    }
   }
 }
 

csrc/cpu/cpu_types_x86.hpp

@@ -99,7 +99,7 @@ struct FP16Vec8 : public Vec<FP16Vec8> {
 
   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<FP16Vec16> {
@@ -146,7 +146,7 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
 
   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<BF16Vec16> {

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]
 
 

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)