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

[jaylisde]

Purpose

PR vllm-project#22145 (Aug 2025) routed all CPU SiluAndMul calls through PyTorch native (F.silu(x[..., :d]) * x[..., d:]) to dodge a TORCH_CHECK in csrc/cpu/activation.cpp that fired on Qwen2.5 with non-aligned hidden sizes. The PR claimed "for eager mode, the performance drop will be very slight." On AMD EPYC 9R14 the actual eager-mode drop is 4–8× on prefill.

This PR fixes the underlying kernel and restores C++ dispatch.

Three changes:

1. csrc/cpu/activation.cpp — replace 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 / FP16Vec8 save(void*) were doing aligned MOVDQA via *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. This was latent on main because the TORCH_CHECK always tripped first; once [Bugfix][Perf][CPU] Restore vectorized silu_and_mul on CPU; handle uneven hidden dim #1 lets uneven d through, the unaligned write is exercised. Caught while developing this patch (D=14337, num_tokens=2 → SIGSEGV before the cpu_types_x86.hpp change).
3. vllm/model_executor/layers/activation.py — with the kernel fixed, restore the C++ dispatch on CPU for SiluAndMul.

The same kernel template also serves gelu_and_mul, gelu_tanh_and_mul, gelu_new, gelu_fast, and gelu_quick. Those were never explicitly Python-fallbacked, so they relied on the CustomOp default (forward_cpu → forward_native); any caller that used self.op directly on CPU would have hit the same TORCH_CHECK. This fix makes the kernel correct for them too.

Reproducing on main

import torch, vllm._C
x = torch.randn(1, 2 * 9, dtype=torch.bfloat16)
out = torch.empty(1, 9, dtype=torch.bfloat16)
torch.ops._C.gelu_and_mul(out, x)
# RuntimeError: Expected d % VEC_ELEM_NUM == 0 to be true, but got false.

SiluAndMul itself dodges this since vllm-project#22145 by going through forward_native, but the underlying _C.silu_and_mul, _C.gelu_and_mul, _C.gelu_tanh_and_mul, _C.gelu_new, _C.gelu_fast, _C.gelu_quick are all still affected when called directly.

Perf regression from the Python fallback (this is the visible one for users):

import time, torch, vllm._C
import torch.nn.functional as F
torch.set_num_threads(64)
M, D = 4096, 14336  # Llama-3-8B intermediate size
x = torch.randn(M, 2 * D, dtype=torch.bfloat16)

def native():  # what vLLM CPU does today on main
    d = D
    return F.silu(x[..., :d]) * x[..., d:]

out = torch.empty(M, D, dtype=torch.bfloat16)
def cpp():    # what this PR restores
    torch.ops._C.silu_and_mul(out, x)

for fn in (native, cpp):
    for _ in range(5): fn()
    t0 = time.perf_counter()
    for _ in range(30): fn()
    print(f"{fn.__name__}: {(time.perf_counter()-t0)/30*1e6:.0f} us")
# native: 28000+ us
# cpp:     6000+ us  (~4.7×)

Test Plan

pytest tests/kernels/core/test_cpu_activation.py

D parametrization is extended with 513 and 2049 to exercise the new scalar-tail path. The existing values (512, 2048) only ever covered aligned d.

Test Result

96 passed, 16 skipped (Arm-only paths) in 60s

Microbench, AMD EPYC 9R14 (Zen 4, AVX-512 + BF16, no AMX), 64 threads, bf16:

shape	native (μs)	cpp (μs)	speedup
Llama-3-8B prefill (M=4096)	28387	6068	4.68×
Llama-3-70B prefill (M=4096)	17219	2147	8.02×
D=18944, M=4096	38363	4893	7.84×
D=18943, M=4096 (uneven)	40378	4974	8.12×
D=14329, M=4096 (uneven)	29467	3711	7.94×
Llama-3-8B decode (M=16)	62	49	1.28×
Llama-3-70B decode (M=16)	62	36	1.74×

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).

Notes

The _mm_storeu_si128 change in cpu_types_x86.hpp is a one-line correctness fix on a primitive shared by other CPU kernels. On Zen 4 (and Intel SPR), MOVDQU runs at the same throughput as MOVDQA when the address happens to be aligned, and only MOVDQU is correct when it isn't.

AI assistance was used to draft the patch; the benchmarks and tests were run against a locally rebuilt CPU wheel and manually verified.