[CPU] Refine CPU kernel dispatch (#50801)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
Li, Jiang committed
Aug 3, 2026 at 17:42 UTC
c8602c79062440074a018c1d5f875a5571eb6881
6 files changed
+47
-30
cmake/cpu_extension.cmake
+1
-1
@@ -345,7 +345,7 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
345
FetchContent_Declare(
346
oneDNN
347
GIT_REPOSITORY https://github.com/oneapi-src/oneDNN.git
348
- GIT_TAG v3.10
348
+ GIT_TAG v3.13
349
GIT_PROGRESS TRUE
350
GIT_SHALLOW TRUE
351
)
docs/getting_started/installation/cpu.md
-1
@@ -152,7 +152,6 @@ VLLM_USE_PRECOMPILED=1 VLLM_PRECOMPILED_WHEEL_VARIANT=cpu VLLM_TARGET_DEVICE=cpu
152
- `VLLM_CPU_OMP_THREADS_BIND`: specify the CPU cores dedicated to the OpenMP threads, can be set as CPU id lists, `auto` (by default), or `nobind` (to disable binding to individual CPU cores and to inherit user-defined OpenMP variables). For example, `VLLM_CPU_OMP_THREADS_BIND=0-31` means there will be 32 OpenMP threads bound on 0-31 CPU cores. `VLLM_CPU_OMP_THREADS_BIND=0-31|32-63` means there will be 2 tensor parallel processes, 32 OpenMP threads of rank0 are bound on 0-31 CPU cores, and the OpenMP threads of rank1 are bound on 32-63 CPU cores. By setting to `auto`, the OpenMP threads of each rank are bound to the CPU cores in each NUMA node respectively. If set to `nobind`, the number of OpenMP threads is determined by the standard `OMP_NUM_THREADS` environment variable.
153
- `VLLM_CPU_NUM_OF_RESERVED_CPU`: specify the number of CPU cores which are not dedicated to the OpenMP threads for each rank. The variable only takes effect when VLLM_CPU_OMP_THREADS_BIND is set to `auto`. Default value is `None`. If the value is not set and use `auto` thread binding, no CPU will be reserved for `world_size == 1`, 1 CPU per rank will be reserved for `world_size > 1`.
154
- `CPU_VISIBLE_MEMORY_NODES`: specify visible NUMA memory nodes for vLLM CPU workers, similar to ```CUDA_VISIBLE_DEVICES```. The variable only takes effect when VLLM_CPU_OMP_THREADS_BIND is set to `auto`. The variable provides more control for the auto thread-binding feature, such as masking nodes and changing nodes binding sequence.
155
-- `VLLM_CPU_SGL_KERNEL` (x86 only, Experimental): whether to use small-batch optimized kernels for the linear layer, especially for low-latency requirements like online serving. The kernels require AMX instruction set, BFloat16 weight type and weight shapes divisible by 32. Default is `0` (False). MoE layers always use the grouped-gemm kernels on x86 and are unaffected by this variable.
155
- `VLLM_ZENTORCH_WEIGHT_PREPACK` (AMD Zen only): when `ZenCpuPlatform` is active, eagerly prepack linear weights into ZenDNN's blocked layout at model load time, eliminating per-inference layout conversion overhead. Default is `1` (enabled). See [AMD Zen optimizations](#amd-zen-optimizations).
156
157
## FAQ
vllm/envs.py
-3
@@ -51,7 +51,6 @@ if TYPE_CHECKING:
51
VLLM_CPU_KVCACHE_SPACE: int | None = 0
52
VLLM_CPU_OMP_THREADS_BIND: str = "auto"
53
VLLM_CPU_NUM_OF_RESERVED_CPU: int | None = None
54
- VLLM_CPU_SGL_KERNEL: bool = False
54
VLLM_CPU_ATTN_SPLIT_KV: bool = True
55
VLLM_ZENTORCH_WEIGHT_PREPACK: bool = True
56
VLLM_CPU_INT4_W4A8: bool = True
@@ -868,8 +867,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
867
if "VLLM_CPU_NUM_OF_RESERVED_CPU" in os.environ
868
else None
869
),
871
- # (CPU backend only) whether to use SGL kernels, optimized for small batch.
872
- "VLLM_CPU_SGL_KERNEL": lambda: bool(int(os.getenv("VLLM_CPU_SGL_KERNEL", "0"))),
870
# (CPU backend only) whether to enable attention spilt KV.
871
"VLLM_CPU_ATTN_SPLIT_KV": lambda: bool(
872
int(os.getenv("VLLM_CPU_ATTN_SPLIT_KV", "1"))
vllm/model_executor/kernels/linear/scaled_mm/cpu.py
-2
@@ -5,7 +5,6 @@
5
import torch
6
7
from vllm import _custom_ops as ops
8
-from vllm import envs
8
from vllm.model_executor.layers.quantization.utils import replace_parameter
9
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
10
convert_to_channelwise,
@@ -44,7 +43,6 @@ class CPUInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel):
43
N, K = weight.size()
44
if (
45
current_platform.get_cpu_architecture() == CpuArchEnum.X86
47
- and envs.VLLM_CPU_SGL_KERNEL
46
and self.config.input_symmetric
47
and check_cpu_sgl_kernel(N, K, dtype)
48
):
vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
+44
-3
@@ -508,7 +508,11 @@ class CPUExpertsFp8(mk.FusedMoEExpertsMonolithic):
508
509
@staticmethod
510
def _supports_current_device() -> bool:
511
- return current_platform.is_cpu()
511
+ return (
512
+ current_platform.is_cpu()
513
+ and current_platform.get_cpu_architecture() == CpuArchEnum.X86
514
+ and torch.cpu._is_amx_tile_supported()
515
+ )
516
517
@staticmethod
518
def _supports_no_act_and_mul() -> bool:
@@ -662,7 +666,11 @@ class CPUExpertsMxfp4(mk.FusedMoEExpertsMonolithic):
666
667
@staticmethod
668
def _supports_current_device() -> bool:
665
- return current_platform.is_cpu()
669
+ return (
670
+ current_platform.is_cpu()
671
+ and current_platform.get_cpu_architecture() == CpuArchEnum.X86
672
+ and torch.cpu._is_amx_tile_supported()
673
+ )
674
675
@staticmethod
676
def _supports_no_act_and_mul() -> bool:
@@ -865,7 +873,11 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic):
873
874
@staticmethod
875
def _supports_current_device() -> bool:
868
- return current_platform.is_cpu()
876
+ return (
877
+ current_platform.is_cpu()
878
+ and current_platform.get_cpu_architecture() == CpuArchEnum.X86
879
+ and torch.cpu._is_amx_tile_supported()
880
+ )
881
882
@staticmethod
883
def _supports_no_act_and_mul() -> bool:
@@ -1008,7 +1020,36 @@ class CPUExpertsInt8(mk.FusedMoEExpertsMonolithic):
1020
return (
1021
current_platform.is_cpu()
1022
and current_platform.get_cpu_architecture() == CpuArchEnum.X86
1023
+ and torch.cpu._is_amx_tile_supported()
1024
+ )
1025
+
1026
+ @staticmethod
1027
+ def is_supported_config(
1028
+ cls: type[mk.FusedMoEExperts],
1029
+ moe_config: FusedMoEConfig,
1030
+ weight_key: QuantKey | None,
1031
+ activation_key: QuantKey | None,
1032
+ activation_format: mk.FusedMoEActivationFormat,
1033
+ ) -> tuple[bool, str | None]:
1034
+ supported, reason = mk.FusedMoEExperts.is_supported_config(
1035
+ cls,
1036
+ moe_config,
1037
+ weight_key,
1038
+ activation_key,
1039
+ activation_format,
1040
)
1041
+ if not supported:
1042
+ return supported, reason
1043
+ # convert_weight_packed (shared VNNI prepack) requires the w13
1044
+ # OC/IC and w2 OC/IC to be multiples of TILE_N=16/TILE_K=32; the
1045
+ # w1 gate-up kernel additionally requires the intermediate size
1046
+ # itself (not 2x) to be a multiple of 32 (moe_int8.cpp), which
1047
+ # dominates. Net effect: both dims must be multiples of 32.
1048
+ if moe_config.hidden_dim % 32 != 0:
1049
+ return False, "kernel requires hidden dim divisible by 32"
1050
+ if moe_config.intermediate_size_per_partition % 32 != 0:
1051
+ return False, "kernel requires intermediate dim divisible by 32"
1052
+ return True, None
1053
1054
@staticmethod
1055
def _supports_no_act_and_mul() -> bool:
vllm/model_executor/layers/utils.py
+2
-20
@@ -225,7 +225,7 @@ direct_register_custom_op(
225
def check_cpu_sgl_kernel(n: int, k: int, dtype: torch.dtype) -> bool:
226
return (
227
torch.cpu._is_amx_tile_supported()
228
- and (dtype in (torch.bfloat16, torch.int8))
228
+ and (dtype in (torch.bfloat16, torch.float16, torch.int8))
229
and k % 32 == 0
230
and n % 16 == 0
231
)
@@ -263,9 +263,6 @@ def dispatch_cpu_unquantized_gemm(
263
layer.weight.data = ops.causal_conv1d_weight_pack(unpacked)
264
return
265
266
- N, K = layer.weight.size()
267
- dtype = layer.weight.dtype
268
-
266
# Zen CPU path: zentorch_linear_unary with optional eager weight prepacking.
267
if current_platform.is_zen_cpu() and hasattr(
268
torch.ops.zentorch, "zentorch_linear_unary"
@@ -294,22 +291,7 @@ def dispatch_cpu_unquantized_gemm(
291
)
292
return
293
297
- if envs.VLLM_CPU_SGL_KERNEL and check_cpu_sgl_kernel(N, K, dtype):
298
- packed_weight = torch.ops._C.convert_weight_packed(layer.weight)
299
- if getattr(layer, "bias", None) is not None:
300
- bias_f32 = layer.bias.to(torch.float32)
301
- else:
302
- bias_f32 = None
303
- layer.cpu_linear = lambda x, weight, bias: torch.ops._C.weight_packed_linear(
304
- x, packed_weight, bias_f32 if bias is not None else None, True
305
- )
306
- if remove_weight:
307
- layer.weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
308
- logger.debug_once(
309
- "CPU unquantized GEMM dispatch: using sgl-kernel weight_packed_linear"
310
- )
311
- return
312
- elif (
294
+ if (
295
ops._supports_onednn
296
and current_platform.get_cpu_architecture() != CpuArchEnum.POWERPC
297
):