[CPU] Refine CPU kernel dispatch by bigPYJ1151 · Pull Request #50801 · vllm-project/vllm

bigPYJ1151 · 2026-08-03T03:33:40Z

Purpose

Make the existing SGLang-derived AMX kernel for INT8 W8A8 linear (symmetric
weight quant + shape-aligned) unconditional instead of gated behind the
experimental VLLM_CPU_SGL_KERNEL opt-in flag. Benchmarking shows no
downside to always enabling it within its existing eligibility window.
oneDNN remains the fallback for every ineligible case (asymmetric quant,
misaligned shapes, non-AMX x86) and is untouched for unquantized
FP32/FP16/BF16 GEMM, which continues to use oneDNN exclusively, exactly as
before.
Fix check_cpu_sgl_kernel's eligibility gate, which excluded
torch.float16 even though the underlying vendored kernels
(weight_packed_linear, int8_scaled_mm_with_quant, fused_experts_cpu
in csrc/cpu/sgl-kernels/) all dispatch via
AT_DISPATCH_REDUCED_FLOATING_TYPES, which covers both BFloat16 and
Half with no additional dtype restriction. Confirmed against the
vendored kernel source, which is byte-identical to sglang's latest
upstream main as of this PR.
Fix two device/shape-selection bugs found while auditing every SGLang
GEMM/MoE selection condition against the vendored kernels:
- CPUExpertsFp8, CPUExpertsMxfp4, and CPUExpertsInt4's
  _supports_current_device() only checked current_platform.is_cpu()
  (no x86 or AMX check at all), and CPUExpertsInt8's checked x86 but not
  AMX. All four unconditionally call fused_experts_cpu /
  torch.ops._C.convert_weight_packed in apply(), which per
  cmake/cpu_extension.cmake's VLLM_EXT_SRC_SGL are compiled only into
  the AMX-tier _C extension -- not _C_AVX512/_C_AVX2, not ARM.
  Selecting one of these classes on non-AMX x86 hardware (or, for the
  three without any arch check, on ARM) would crash with an
  AttributeError from inside the kernel call instead of the oracle
  cleanly reporting no supported backend. Added the same x86 +
  torch.cpu._is_amx_tile_supported() check already used by the linear
  counterpart (CPUFp8BlockScaledMMKernel.is_supported()) and by
  check_cpu_sgl_kernel.
- CPUExpertsInt8 had no shape-alignment check at the oracle-selection
  level. Verified concretely (not assumed): convert_weight_packed (the
  shared VNNI prepack all four quantized experts call) requires weight
  OC % TILE_N(16) == 0 and IC % TILE_K(32) == 0
  (csrc/cpu/sgl-kernels/gemm.cpp); moe_int8.cpp additionally
  hard-requires the intermediate size itself to be a multiple of 32 via
  its own TORCH_CHECK. Combined across w13/w2, both hidden_dim and
  intermediate_size_per_partition must be multiples of 32 -- exactly
  what ArmCPUExpertsInt8 already validates for its NEON kernel. Without
  this, a misaligned model would pass oracle selection and only fail once
  process_weights_after_loading calls convert_weight_packed, crashing
  with a generic TORCH_CHECK message instead of the oracle cleanly
  rejecting the backend up front. Verified the gate behaves correctly via
  a direct unit check (misaligned shapes rejected with a clear reason,
  aligned 512/512 accepted).
Bump the vendored x86 oneDNN dependency from v3.10 to v3.13 to pick up
upstream fixes and improvements. ARM keeps its separately pinned commit,
unaffected.

VLLM_CPU_SGL_KERNEL is now fully removed. #50133 merged ahead of this
PR and deleted the unquantized-MoE SGL dispatch block (SGLFusedMOE and its
flag check) as part of migrating unquantized CPU MoE to the modular-kernel
experts structure -- the flag's only remaining use after that was the linear
INT8 W8A8 path this PR already made unconditional, so this PR now removes
the envs.py declaration and the doc entry outright. This branch has been
rebased onto latest main (post-#50133); the environment/MoE-dispatch/doc
changes that were being deferred are no longer needed since #50133 already
handles them.

Not fixed here (flagged, not addressed): CPUExpertsFp8, CPUExpertsMxfp4,
and CPUExpertsInt4 still have no shape-alignment gate. convert_weight_packed's
IC accounting doubles for byte-packed uint8 storage (mxfp4/int4), and
neither moe_fp8.cpp nor moe_int4.cpp has any runtime shape TORCH_CHECK
at all to cross-validate a threshold against, unlike moe_int8.cpp. Getting
the modulus wrong for these would either falsely reject working configs or
still let bad ones through, so left as a follow-up rather than guessed at.

Changes

vllm/envs.py: remove VLLM_CPU_SGL_KERNEL declaration and parsing.
vllm/model_executor/kernels/linear/scaled_mm/cpu.py: drop the flag
check from the INT8 SGL-vs-oneDNN dispatch condition.
vllm/model_executor/layers/utils.py: remove the now-dead SGL branch
from unquantized FP32/FP16/BF16 GEMM dispatch (oneDNN-only even with the
flag on, since only the INT8/MoE paths ever used SGL by default); add
torch.float16 to check_cpu_sgl_kernel's dtype gate.
vllm/model_executor/layers/fused_moe/experts/cpu_moe.py: add
x86 + AMX check to CPUExpertsFp8/CPUExpertsMxfp4/CPUExpertsInt4/
CPUExpertsInt8's _supports_current_device(); add a shape-alignment
is_supported_config override to CPUExpertsInt8.
docs/getting_started/installation/cpu.md: remove the flag's doc entry.
cmake/cpu_extension.cmake: bump x86 oneDNN FetchContent pin, v3.10 -> v3.13.

Test Plan

.venv/bin/python -m pytest tests/kernels/test_onednn.py \
  tests/model_executor/test_cpu_unquantized_gemm_dispatch.py \
  tests/kernels/quantization/test_scaled_mm_kernel_selection.py \
  tests/kernels/quantization/test_cpu_fp8_scaled_mm.py \
  tests/quantization/test_cpu_w8a8.py \
  tests/kernels/moe/test_cpu_quant_fused_moe.py \
  tests/kernels/moe/test_cpu_int4_moe.py \
  tests/kernels/moe/test_cpu_fused_moe.py

Offline-inference smoke test (examples/basic/offline_inference/generate.py)
across meta-llama/Llama-3.1-8B-Instruct, google/gemma-7b, Qwen/Qwen3-8B,
and RedHatAI/Meta-Llama-3.1-8B-quantized.w8a8.

Serving benchmark (vllm serve + vllm bench serve, 1024/1024 input/output
length) comparing oneDNN v3.10 vs. v3.13 on an AMX-capable x86 host, for the
three BF16 models (6 prompts each), plus a separate before/after comparison
of the linear INT8 flag removal on the W8A8 model (15 prompts).

Test Result

Rebased onto latest main (post-#50133) and rebuilt from a clean incremental
build; 651 tests pass, 289 skipped (non-AMX-specific), 0 failed:
test_onednn.py, CPU unquantized GEMM dispatch, scaled_mm kernel selection,
CPU FP8 scaled_mm, W8A8 e2e, and the full CPU MoE suite
(test_cpu_quant_fused_moe.py, test_cpu_int4_moe.py, test_cpu_fused_moe.py).

All smoke-test models produce coherent output; Llama and Qwen3 are
byte-identical between oneDNN v3.10/v3.13, gemma-7b diverges on one of four
prompts (expected BF16 rounding-order sensitivity between GEMM library
versions, still coherent).

The MoE device-check and shape-gate fixes only change behavior for
configurations not exercisable in this environment (non-AMX x86, ARM,
misaligned shapes -- the existing test suite's shapes are all already
32-aligned); on this AMX-capable host all four expert classes still resolve
to supported=True as before, which the 0-failed MoE test run confirms (no
regression for the eligible case). The shape gate's rejection path was
additionally verified with a direct unit check.

oneDNN v3.10 -> v3.13, serving benchmark (1024/1024, 6 prompts):

Model	TTFT (mean)	TPOT (mean)	total tok/s
Llama-3.1-8B-Instruct	2750 -> 2898 ms	103.3 -> 100.1 ms	113.2 -> 116.6
gemma-7b	3255 -> 3092 ms	133.3 -> 130.4 ms	87.9 -> 89.9
Qwen3-8B	2812 -> 2806 ms	106.3 -> 102.4 ms	110.0 -> 114.1

No regression; TPOT/throughput improve slightly (~2-3%) across all three
models, TTFT is within run-to-run noise for a 6-prompt sample.

Linear INT8 flag removal, serving benchmark on
RedHatAI/Meta-Llama-3.1-8B-quantized.w8a8 (1024/1024, 15 prompts), oneDNN
(flag off, previous default) vs. SGL now unconditional:

Metric	oneDNN	SGL (unconditional)
Mean TTFT	4136.7 ms	3721.5 ms
Mean TPOT	79.15 ms	74.09 ms
Total tok/s	359.8	385.1

claude

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

mergify · 2026-08-03T03:34:17Z

Documentation preview: https://vllm--50801.org.readthedocs.build/en/50801/

The SGLang-derived AMX kernels for INT8 W8A8 linear (symmetric, shape-aligned) and unquantized MoE were gated behind an experimental opt-in flag. Benchmarking shows no downside to enabling them unconditionally in their existing eligibility window, so drop the flag and always use them when eligible. oneDNN remains the fallback for all ineligible cases (asymmetric quant, misaligned shapes, non-AMX x86) and is untouched for unquantized FP32/FP16/BF16 GEMM, which continues to use oneDNN exclusively as before. Tested: 279 tests pass (test_onednn.py, CPU unquantized GEMM dispatch, scaled_mm kernel selection, CPU FP8 scaled_mm, W8A8 e2e model test). Offline-inference smoke test across Llama-3.1-8B-Instruct, gemma-7b, Qwen3-8B, and RedHatAI/Meta-Llama-3.1-8B-quantized.w8a8 all produce coherent output. Benchmarked the W8A8 model (1024/1024 in/out, 15 prompts) with the flag removed vs. oneDNN baseline (flag off): ~10% lower mean TTFT (4136.7ms -> 3721.5ms), ~6% lower mean TPOT (79.15ms -> 74.09ms), ~7% higher total throughput (359.8 -> 385.1 tok/s). Signed-off-by: jiang1.li <jiang1.li@intel.com>

Move the x86 FetchContent pin from v3.10 to v3.13 to pick up upstream GEMM/threading fixes and improvements. ARM keeps its separately pinned commit, unchanged. Tested: 279 tests pass (test_onednn.py, CPU unquantized GEMM dispatch, scaled_mm kernel selection, CPU FP8 scaled_mm, W8A8 e2e model test). Offline-inference smoke test across Llama-3.1-8B-Instruct, gemma-7b, and Qwen3-8B is coherent on both versions (Llama and Qwen3 produce byte-identical greedy output vs. v3.10; gemma-7b diverges on one of four canned prompts, expected BF16 rounding-order sensitivity, still coherent). Benchmarked serving (1024/1024 in/out, 6 prompts) v3.10 vs. v3.13: TTFT (mean) TPOT (mean) total tok/s Llama-3.1-8B 2750->2898 ms 103.3->100.1 ms 113.2->116.6 gemma-7b 3255->3092 ms 133.3->130.4 ms 87.9->89.9 Qwen3-8B 2812->2806 ms 106.3->102.4 ms 110.0->114.1 No regression; TPOT and throughput improve slightly (~2-3%) across all three models, TTFT is within run-to-run noise for a 6-prompt sample. Signed-off-by: jiang1.li <jiang1.li@intel.com>

The gate excluded torch.float16, but the underlying kernels (weight_packed_linear, int8_scaled_mm_with_quant, fused_experts_cpu in csrc/cpu/sgl-kernels/) all dispatch via AT_DISPATCH_REDUCED_FLOATING_TYPES, which covers both BFloat16 and Half with no additional dtype restriction in any of their TORCH_CHECK guards. Confirmed against the actual vendored kernel source (byte-identical to sglang's latest upstream main as of this commit) rather than assumed. Concretely this means FP16 weights/activations were being routed to the slower fallback path even on AMX hardware where the SGL kernel already handles them correctly. Signed-off-by: jiang1.li <jiang1.li@intel.com>

CPUExpertsFp8, CPUExpertsMxfp4, and CPUExpertsInt4's _supports_current_device() only checked current_platform.is_cpu(), with no x86 or AMX check; CPUExpertsInt8 checked x86 but not AMX either. All four unconditionally call fused_experts_cpu / torch.ops._C.convert_weight_packed in apply(), which per cmake/cpu_extension.cmake's VLLM_EXT_SRC_SGL are compiled only into the AMX-tier _C extension, not _C_AVX512/_C_AVX2 or ARM. Selecting one of these classes on non-AMX x86 hardware (or, for the three without any arch check, on ARM) would crash with an AttributeError from inside the kernel call instead of the oracle cleanly rejecting the backend and reporting no supported experts. Add the same x86 + torch.cpu._is_amx_tile_supported() check already used by the linear counterpart (CPUFp8BlockScaledMMKernel.is_supported()) and by check_cpu_sgl_kernel. Tested: tests/kernels/moe/test_cpu_quant_fused_moe.py, tests/kernels/moe/test_cpu_int4_moe.py, tests/kernels/moe/test_cpu_fused_moe.py -- 357 passed, 289 skipped (non-AMX-specific skips), 0 failed on this AMX-capable host, confirming no regression for the eligible case. Signed-off-by: jiang1.li <jiang1.li@intel.com>

Verified concretely rather than assumed: convert_weight_packed (the shared VNNI prepack all four quantized CPU MoE experts call) requires weight OC % TILE_N(16) == 0 and IC % TILE_K(32) == 0 in csrc/cpu/sgl-kernels/gemm.cpp; moe_int8.cpp additionally hard-requires the intermediate size itself (not 2x) to be a multiple of 32 via its own TORCH_CHECK. Combined across w13 ([E, 2*intermediate, hidden]) and w2 ([E, hidden, intermediate]), both hidden_dim and intermediate_size_per_partition must be multiples of 32 -- exactly the same requirement ArmCPUExpertsInt8 already validates for its NEON kernel, now added to the x86 counterpart. Without this, a misaligned model would pass oracle selection and only fail once process_weights_after_loading calls convert_weight_packed, crashing with a generic "invalid weight out/in features" TORCH_CHECK instead of the oracle cleanly rejecting the backend up front. FP8/MXFP4/INT4 experts are not touched here: convert_weight_packed's own IC accounting doubles for byte-packed uint8 storage (mxfp4/int4), and neither moe_fp8.cpp nor moe_int4.cpp has any runtime shape TORCH_CHECK at all (confirmed by inspection) to cross-validate a threshold against -- getting the modulus wrong for these would either falsely reject working configs or still let bad ones through. Left as a flagged follow-up rather than guessed at. Verified the gate behaves correctly via direct unit check (misaligned hidden_dim/intermediate_size -> rejected with a clear reason; aligned 512/512 -> accepted) and via tests/kernels/moe/test_cpu_quant_fused_moe.py + test_cpu_int4_moe.py + test_cpu_fused_moe.py -- 357 passed, 289 skipped, 0 failed (all existing test shapes are already 32-aligned, so this changes no existing test outcome). Signed-off-by: jiang1.li <jiang1.li@intel.com>

mergify · 2026-08-03T07:31:03Z

Documentation preview: https://vllm--50801.org.readthedocs.build/en/50801/

bigPYJ1151 · 2026-08-03T07:33:03Z

/ci run

github-actions · 2026-08-03T07:40:23Z

✅ Triggered Buildkite CI #81934 for commit 601e59635061.

PR vllm-project#50801 removed the SGLang AMX weight_packed_linear path from the unquantized (bf16/fp16) linear dispatch, leaving oneDNN's onednn_mm as the only CPU kernel. For small weights -- most notably MoE router/gate projections, where N is the expert count rather than a hidden-size-scaled dimension -- oneDNN never reaches its compute-bound regime no matter how large the batch gets, so SGL's lower per-call dispatch overhead wins across the full measured M range. Restore that dispatch, gated by a weight-size threshold so larger dense projections (qkv/o_proj/gate_up/ down/lm_head) keep using oneDNN once batch size grows past decode-sized M. Also add a runtime AMX-FP16 capability check: AMX-BF16/INT8 (amx_tile) and AMX-FP16 are separate CPU ISA extensions (e.g. Sapphire/Emerald Rapids expose the former but not the latter), and SGL's can_use_brgemm<at::Half> always attempts brgemm for fp16 regardless of M, so fp16 needs its own capability gate rather than piggybacking on amx_tile. Signed-off-by: jiang1.li <jiang1.li@intel.com>

…nflicts Dropped the legacy TYPE_CHECKING block and environment_variables dict wholesale, then ported main's delta across 10 main-side commits: Additions: VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4 (vllm-project#50582), VLLM_USE_RUST_BENCH (vllm-project#50081), VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT (vllm-project#50656), VLLM_RAISE_ON_LOGIT_NANS (vllm-project#50323), VLLM_ENABLE_COHERE_API (vllm-project#47189). Modifications: VLLM_COMPUTE_NANS_IN_LOGITS is now implied by VLLM_RAISE_ON_LOGIT_NANS (cross-field, so a model_validator); _resolve_rust_frontend_path -> _resolve_rust_cli_path, resolving on either VLLM_USE_RUST_FRONTEND or VLLM_USE_RUST_BENCH. Deletions: VLLM_CPU_SGL_KERNEL (vllm-project#50801), Q_/K_/V_SCALE_CONSTANT (vllm-project#49389 -- main deleted the dict entries but left the TYPE_CHECKING annotations; followed the PR's intent). tests: ported VLLMValidationError assertions; adapted test_rust_bench_auto_path_missing_fails_fast to construct ServerSettings directly. Dropped the Q_SCALE_CONSTANT case from test_envs_pydantic.py. AI assistance (Claude) was used for this merge resolution. Co-authored-by: Claude Signed-off-by: Vinay Damodaran <vrdn@hey.com>

bigPYJ1151 requested review from mgoin, pavanimajety and zyongye as code owners August 3, 2026 03:33

claude Bot reviewed Aug 3, 2026

View reviewed changes

mergify Bot added documentation Improvements or additions to documentation ci/build cpu Related to CPU backends labels Aug 3, 2026

bigPYJ1151 changed the title ~~[CPU] Drop VLLM_CPU_SGL_KERNEL flag, bump oneDNN to v3.13~~ [CPU] Make SGL AMX kernels unconditional for linear layers, fix FP16 gate, bump oneDNN to v3.13 Aug 3, 2026

bigPYJ1151 changed the title ~~[CPU] Make SGL AMX kernels unconditional for linear layers, fix FP16 gate, bump oneDNN to v3.13~~ [CPU] Refine CPU kernel dispatch Aug 3, 2026

bigPYJ1151 added 5 commits August 3, 2026 07:13

bigPYJ1151 force-pushed the remove-cpu-sgl-kernel-flag branch from 66c415c to 601e596 Compare August 3, 2026 07:30

jikunshang approved these changes Aug 3, 2026

View reviewed changes

bigPYJ1151 merged commit c8602c7 into vllm-project:main Aug 3, 2026
122 of 123 checks passed

bigPYJ1151 deleted the remove-cpu-sgl-kernel-flag branch August 3, 2026 09:42

bigPYJ1151 mentioned this pull request Aug 7, 2026

[CPU] Restore linear dispatch for small unquantized GEMMs #51379

Merged

Navigation Menu

Uh oh!

Uh oh!

[CPU] Refine CPU kernel dispatch - #50801

[CPU] Refine CPU kernel dispatch#50801
bigPYJ1151 merged 5 commits into
vllm-project:mainfrom
bigPYJ1151:remove-cpu-sgl-kernel-flag

bigPYJ1151 commented Aug 3, 2026 •

edited

Loading

Uh oh!

claude Bot left a comment

Uh oh!

mergify Bot commented Aug 3, 2026

Uh oh!

mergify Bot commented Aug 3, 2026

Uh oh!

bigPYJ1151 commented Aug 3, 2026

Uh oh!

github-actions Bot commented Aug 3, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Uh oh!

Conversation

bigPYJ1151 commented Aug 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Changes

Test Plan

Test Result

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

mergify Bot commented Aug 3, 2026

Uh oh!

mergify Bot commented Aug 3, 2026

Uh oh!

bigPYJ1151 commented Aug 3, 2026

Uh oh!

github-actions Bot commented Aug 3, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

bigPYJ1151 commented Aug 3, 2026 •

edited

Loading