[CPU] Refine CPU kernel dispatch

cmake/cpu_extension.cmake

@@ -345,7 +345,7 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
             FetchContent_Declare(
                 oneDNN
                 GIT_REPOSITORY https://github.com/oneapi-src/oneDNN.git
-                GIT_TAG        v3.10
+                GIT_TAG        v3.13
                 GIT_PROGRESS   TRUE
                 GIT_SHALLOW    TRUE
             )

docs/getting_started/installation/cpu.md

@@ -152,7 +152,6 @@ VLLM_USE_PRECOMPILED=1 VLLM_PRECOMPILED_WHEEL_VARIANT=cpu VLLM_TARGET_DEVICE=cpu
 - `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.
 - `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`.
 - `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.
-- `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.
 - `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).
 
 ## FAQ

vllm/envs.py

@@ -51,7 +51,6 @@
     VLLM_CPU_KVCACHE_SPACE: int | None = 0
     VLLM_CPU_OMP_THREADS_BIND: str = "auto"
     VLLM_CPU_NUM_OF_RESERVED_CPU: int | None = None
-    VLLM_CPU_SGL_KERNEL: bool = False
     VLLM_CPU_ATTN_SPLIT_KV: bool = True
     VLLM_ZENTORCH_WEIGHT_PREPACK: bool = True
     VLLM_CPU_INT4_W4A8: bool = True
@@ -868,8 +867,6 @@ def _resolve_rust_cli_path() -> str | None:
         if "VLLM_CPU_NUM_OF_RESERVED_CPU" in os.environ
         else None
     ),
-    # (CPU backend only) whether to use SGL kernels, optimized for small batch.
-    "VLLM_CPU_SGL_KERNEL": lambda: bool(int(os.getenv("VLLM_CPU_SGL_KERNEL", "0"))),
     # (CPU backend only) whether to enable attention spilt KV.
     "VLLM_CPU_ATTN_SPLIT_KV": lambda: bool(
         int(os.getenv("VLLM_CPU_ATTN_SPLIT_KV", "1"))

vllm/model_executor/kernels/linear/scaled_mm/cpu.py

@@ -5,7 +5,6 @@
 import torch
 
 from vllm import _custom_ops as ops
-from vllm import envs
 from vllm.model_executor.layers.quantization.utils import replace_parameter
 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
     convert_to_channelwise,
@@ -44,7 +43,6 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         N, K = weight.size()
         if (
             current_platform.get_cpu_architecture() == CpuArchEnum.X86
-            and envs.VLLM_CPU_SGL_KERNEL
             and self.config.input_symmetric
             and check_cpu_sgl_kernel(N, K, dtype)
         ):

vllm/model_executor/layers/fused_moe/experts/cpu_moe.py

@@ -508,7 +508,11 @@ def activation_format() -> mk.FusedMoEActivationFormat:
 
     @staticmethod
     def _supports_current_device() -> bool:
-        return current_platform.is_cpu()
+        return (
+            current_platform.is_cpu()
+            and current_platform.get_cpu_architecture() == CpuArchEnum.X86
+            and torch.cpu._is_amx_tile_supported()
+        )
 
     @staticmethod
     def _supports_no_act_and_mul() -> bool:
@@ -662,7 +666,11 @@ def activation_format() -> mk.FusedMoEActivationFormat:
 
     @staticmethod
     def _supports_current_device() -> bool:
-        return current_platform.is_cpu()
+        return (
+            current_platform.is_cpu()
+            and current_platform.get_cpu_architecture() == CpuArchEnum.X86
+            and torch.cpu._is_amx_tile_supported()
+        )
 
     @staticmethod
     def _supports_no_act_and_mul() -> bool:
@@ -865,7 +873,11 @@ def activation_format() -> mk.FusedMoEActivationFormat:
 
     @staticmethod
     def _supports_current_device() -> bool:
-        return current_platform.is_cpu()
+        return (
+            current_platform.is_cpu()
+            and current_platform.get_cpu_architecture() == CpuArchEnum.X86
+            and torch.cpu._is_amx_tile_supported()
+        )
 
     @staticmethod
     def _supports_no_act_and_mul() -> bool:
@@ -1008,7 +1020,36 @@ def _supports_current_device() -> bool:
         return (
             current_platform.is_cpu()
             and current_platform.get_cpu_architecture() == CpuArchEnum.X86
+            and torch.cpu._is_amx_tile_supported()
+        )
+
+    @staticmethod
+    def is_supported_config(
+        cls: type[mk.FusedMoEExperts],
+        moe_config: FusedMoEConfig,
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+        activation_format: mk.FusedMoEActivationFormat,
+    ) -> tuple[bool, str | None]:
+        supported, reason = mk.FusedMoEExperts.is_supported_config(
+            cls,
+            moe_config,
+            weight_key,
+            activation_key,
+            activation_format,
         )
+        if not supported:
+            return supported, reason
+        # convert_weight_packed (shared VNNI prepack) requires the w13
+        # OC/IC and w2 OC/IC to be multiples of TILE_N=16/TILE_K=32; the
+        # w1 gate-up kernel additionally requires the intermediate size
+        # itself (not 2x) to be a multiple of 32 (moe_int8.cpp), which
+        # dominates. Net effect: both dims must be multiples of 32.
+        if moe_config.hidden_dim % 32 != 0:
+            return False, "kernel requires hidden dim divisible by 32"
+        if moe_config.intermediate_size_per_partition % 32 != 0:
+            return False, "kernel requires intermediate dim divisible by 32"
+        return True, None
 
     @staticmethod
     def _supports_no_act_and_mul() -> bool:

vllm/model_executor/layers/utils.py

@@ -225,7 +225,7 @@ def rocm_unquantized_gemm(
 def check_cpu_sgl_kernel(n: int, k: int, dtype: torch.dtype) -> bool:
     return (
         torch.cpu._is_amx_tile_supported()
-        and (dtype in (torch.bfloat16, torch.int8))
+        and (dtype in (torch.bfloat16, torch.float16, torch.int8))
         and k % 32 == 0
         and n % 16 == 0
     )
@@ -263,9 +263,6 @@ def dispatch_cpu_unquantized_gemm(
             layer.weight.data = ops.causal_conv1d_weight_pack(unpacked)
         return
 
-    N, K = layer.weight.size()
-    dtype = layer.weight.dtype
-
     # Zen CPU path: zentorch_linear_unary with optional eager weight prepacking.
     if current_platform.is_zen_cpu() and hasattr(
         torch.ops.zentorch, "zentorch_linear_unary"
@@ -294,22 +291,7 @@ def dispatch_cpu_unquantized_gemm(
         )
         return
 
-    if envs.VLLM_CPU_SGL_KERNEL and check_cpu_sgl_kernel(N, K, dtype):
-        packed_weight = torch.ops._C.convert_weight_packed(layer.weight)
-        if getattr(layer, "bias", None) is not None:
-            bias_f32 = layer.bias.to(torch.float32)
-        else:
-            bias_f32 = None
-        layer.cpu_linear = lambda x, weight, bias: torch.ops._C.weight_packed_linear(
-            x, packed_weight, bias_f32 if bias is not None else None, True
-        )
-        if remove_weight:
-            layer.weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
-        logger.debug_once(
-            "CPU unquantized GEMM dispatch: using sgl-kernel weight_packed_linear"
-        )
-        return
-    elif (
+    if (
         ops._supports_onednn
         and current_platform.get_cpu_architecture() != CpuArchEnum.POWERPC
     ):