[CPU] Restore linear dispatch for small unquantized GEMMs (#51379)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
Li, Jiang committed
Aug 10, 2026 at 11:29 UTC
7ce84b99cecc40e94fcee5ee21114fbf1474ad65
1 file changed
+65
-6
vllm/model_executor/layers/utils.py
+65
-6
@@ -216,13 +216,46 @@ direct_register_custom_op(
216
)
217
218
219
+# Above this weight size, oneDNN's onednn_mm consistently matches or beats
220
+# the SGL AMX kernel once M grows past decode-sized batches, and is within
221
+# noise of it at decode-sized M -- so larger weights default to oneDNN
222
+# rather than SGL. 1 MiB comfortably covers MoE router/gate weights (e.g.
223
+# (2048, 128) .. (2880, 32) bf16/fp16, 180-720 KiB) while staying well below
224
+# any dense qkv/o_proj/gate_up/down/lm_head projection in practice. This
225
+# threshold is derived from bf16/fp16 unquantized dense-GEMM benchmarks only,
226
+# so it does not apply to the int8 scaled_mm path below.
227
+_CPU_SGL_GEMM_MAX_WEIGHT_BYTES = 1 * 1024 * 1024
228
+
229
+
230
def check_cpu_sgl_kernel(n: int, k: int, dtype: torch.dtype) -> bool:
220
- return (
221
- torch.cpu._is_amx_tile_supported()
222
- and (dtype in (torch.bfloat16, torch.float16, torch.int8))
223
- and k % 32 == 0
224
- and n % 16 == 0
225
- )
231
+ if not torch.cpu._is_amx_tile_supported() or dtype not in (
232
+ torch.bfloat16,
233
+ torch.float16,
234
+ torch.int8,
235
+ ):
236
+ return False
237
+ if dtype == torch.float16 and not torch.cpu._is_amx_fp16_supported():
238
+ # AMX-BF16/INT8 (amx_tile) and AMX-FP16 are separate CPU ISA
239
+ # extensions -- e.g. Sapphire/Emerald Rapids expose the former but
240
+ # not the latter -- and can_use_brgemm<at::Half> (gemm.h) always
241
+ # attempts brgemm for fp16 regardless of M, so this needs its own
242
+ # capability check rather than piggybacking on amx_tile.
243
+ return False
244
+ if dtype == torch.int8:
245
+ # int8_scaled_mm_with_quant requires the packed weight to stay int8
246
+ # (gemm_int8.cpp); convert_weight_packed's N < TILE_N fallback
247
+ # returns a float32 tensor instead (gemm.cpp), which would trip
248
+ # that check, so N must be a full TILE_N tile here.
249
+ return k % 32 == 0 and n % 16 == 0
250
+ if n * k * dtype.itemsize > _CPU_SGL_GEMM_MAX_WEIGHT_BYTES:
251
+ return False
252
+ if n < 16:
253
+ # convert_weight_packed transposes to fp32 instead of VNNI-packing
254
+ # when N < TILE_N (gemm.cpp), and weight_packed_linear detects that
255
+ # (via the packed weight's dtype) and routes to its fp32/brgemm
256
+ # fallback kernel -- no N/K alignment required in that regime.
257
+ return True
258
+ return k % 32 == 0 and n % 16 == 0
259
260
261
def dispatch_cpu_unquantized_gemm(
@@ -257,6 +290,9 @@ def dispatch_cpu_unquantized_gemm(
290
layer.weight.data = ops.causal_conv1d_weight_pack(unpacked)
291
return
292
293
+ N, K = layer.weight.size()
294
+ dtype = layer.weight.dtype
295
+
296
# Zen CPU path: zentorch_linear_unary with optional eager weight prepacking.
297
if current_platform.is_zen_cpu() and hasattr(
298
torch.ops.zentorch, "zentorch_linear_unary"
@@ -285,6 +321,29 @@ def dispatch_cpu_unquantized_gemm(
321
)
322
return
323
324
+ # Small weights (e.g. MoE router/gate projections, where N is the expert
325
+ # count rather than a hidden-size-scaled dimension) never reach oneDNN's
326
+ # compute-bound regime, no matter how large the batch gets: SGL's lower
327
+ # per-call dispatch overhead wins consistently across the full measured
328
+ # M range. Larger dense projections (qkv/o_proj/gate_up/down/lm_head)
329
+ # cross over to favoring oneDNN once batch size grows past decode-sized
330
+ # M, so they keep using oneDNN below.
331
+ if check_cpu_sgl_kernel(N, K, dtype):
332
+ packed_weight = torch.ops._C.convert_weight_packed(layer.weight)
333
+ if getattr(layer, "bias", None) is not None:
334
+ bias_f32 = layer.bias.to(torch.float32)
335
+ else:
336
+ bias_f32 = None
337
+ layer.cpu_linear = lambda x, weight, bias: torch.ops._C.weight_packed_linear(
338
+ x, packed_weight, bias_f32 if bias is not None else None, True
339
+ )
340
+ if remove_weight:
341
+ layer.weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
342
+ logger.debug_once(
343
+ "CPU unquantized GEMM dispatch: using sgl-kernel weight_packed_linear"
344
+ )
345
+ return
346
+
347
if (
348
ops._supports_onednn
349
and current_platform.get_cpu_architecture() != CpuArchEnum.POWERPC