[XPU] Unify XPU RMSNorm kernels with vllm_c and drop redundant XPU-specific implementation (#46981)
Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Chaojun Zhang committed
Jul 31, 2026 at 17:28 UTC
482cfc20f05fd7d08401f646ce44ba1f7086dde3
7 files changed
+19
-74
.buildkite/intel_jobs/engine_intel.yaml
+3
-1
@@ -56,11 +56,13 @@ steps:
56
source_file_dependencies:
57
- vllm/v1/engine/
58
- tests/v1/engine/
59
+ - tests/test_config/
60
commands:
61
- >-
62
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
63
'cd tests &&
63
- pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py'
64
+ pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py &&
65
+ VLLM_XPU_ENABLE_XPU_GRAPH=1 pytest -v -s test_config.py'
66
67
- label: V1 e2e (2 GPUs)
68
timeout_in_minutes: 30
.buildkite/intel_jobs/kernels_intel.yaml
+1
@@ -22,4 +22,5 @@ steps:
22
- >-
23
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
24
'cd tests &&
25
+ pytest -v -s ir &&
26
pytest -v -s kernels/ir'
tests/kernels/ir/test_layernorm.py
+9
-9
@@ -18,20 +18,21 @@ from vllm.platforms import current_platform
18
19
rms_norm_native = ir.ops.rms_norm.impls["native"].impl_fn
20
21
+IS_GPGPU_DEVICE = current_platform.is_cuda_alike() or current_platform.is_xpu()
22
+
23
24
@pytest.mark.skipif(
23
- not current_platform.is_cuda_alike() and not current_platform.is_xpu(),
25
+ not IS_GPGPU_DEVICE,
26
reason="Currently only kernels on CUDA, ROCm and XPU",
27
)
28
def test_rms_norm_registration():
29
expected = {
30
"native": True,
29
- "vllm_c": current_platform.is_cuda_alike(),
31
+ "vllm_c": IS_GPGPU_DEVICE,
32
"aiter": current_platform.is_rocm(),
33
"oink": current_platform.has_device_capability(100)
34
and hasattr(torch.ops, "oink")
35
and hasattr(torch.ops.oink, "rmsnorm"),
34
- "xpu_kernels": current_platform.is_xpu(),
36
}
37
38
actual = {
@@ -46,7 +47,7 @@ def test_rms_norm_registration():
47
@pytest.mark.parametrize("hidden_size", COMMON_HIDDEN_SIZES)
48
@pytest.mark.parametrize("epsilon", [1e-6, 1e-5])
49
@pytest.mark.skipif(
49
- not current_platform.is_cuda_alike() and not current_platform.is_xpu(),
50
+ not IS_GPGPU_DEVICE,
51
reason="Currently only kernels on CUDA, ROCm and XPU",
52
)
53
class TestRMSNorm:
@@ -117,7 +118,7 @@ class TestRMSNorm:
118
out_unit_weight = impl.impl_fn(x, torch.ones_like(weight), eps)
119
assert_close(ir.ops.rms_norm, out_no_weight, out_unit_weight)
120
120
- @pytest.mark.parametrize("provider", ["vllm_c", "aiter", "xpu_kernels", "native"])
121
+ @pytest.mark.parametrize("provider", ["vllm_c", "aiter", "native"])
122
def test_torch_opcheck(self, dtype, n_tokens, hidden_size, epsilon, provider):
123
if not ir.ops.rms_norm.impls[provider].supported:
124
pytest.skip(f"{provider} impl not supported on this platform")
@@ -203,18 +204,17 @@ fused_add_rms_norm_native = ir.ops.fused_add_rms_norm.impls["native"].impl_fn
204
205
206
@pytest.mark.skipif(
206
- not current_platform.is_cuda_alike() and not current_platform.is_xpu(),
207
+ not IS_GPGPU_DEVICE,
208
reason="Currently only kernels on CUDA, ROCm and XPU",
209
)
210
def test_fused_add_rms_norm_registration():
211
expected = {
212
"native": True,
212
- "vllm_c": current_platform.is_cuda_alike(),
213
+ "vllm_c": IS_GPGPU_DEVICE,
214
"aiter": current_platform.is_rocm(),
215
"oink": current_platform.has_device_capability(100)
216
and hasattr(torch.ops, "oink")
217
and hasattr(torch.ops.oink, "fused_add_rms_norm"),
217
- "xpu_kernels": current_platform.is_xpu(),
218
}
219
220
actual = {
@@ -261,7 +261,7 @@ def test_vllm_c_fused_add_rms_norm_accepts_nd_input():
261
@pytest.mark.parametrize("hidden_size", COMMON_HIDDEN_SIZES)
262
@pytest.mark.parametrize("epsilon", [1e-6, 1e-5])
263
@pytest.mark.skipif(
264
- not current_platform.is_cuda_alike() and not current_platform.is_xpu(),
264
+ not IS_GPGPU_DEVICE,
265
reason="Currently only kernels on CUDA, ROCm and XPU",
266
)
267
class TestFusedAddRMSNorm:
vllm/kernels/__init__.py
+2
-2
@@ -2,6 +2,6 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Kernel implementations for vLLM."""
4
5
-from . import aiter_ops, oink_ops, vllm_c, xpu_ops
5
+from . import aiter_ops, oink_ops, vllm_c
6
7
-__all__ = ["vllm_c", "aiter_ops", "oink_ops", "xpu_ops"]
7
+__all__ = ["vllm_c", "aiter_ops", "oink_ops"]
vllm/kernels/vllm_c.py
+3
-2
@@ -12,6 +12,7 @@ CUDA_ALIKE = current_platform.is_cuda_alike()
12
"""Most kernels in this file are supported on all CUDA-alike platforms."""
13
IS_ROCM = current_platform.is_rocm()
14
"""ROCm needs shape normalization before calling some vLLM C kernels."""
15
+GPGPU_DEVICE = CUDA_ALIKE or current_platform.is_xpu()
16
17
rms_no_var_size = lambda x, weight, epsilon, variance_size=None: (
18
variance_size is None and (weight is None or weight.dtype == x.dtype)
@@ -20,7 +21,7 @@ rms_no_var_size = lambda x, weight, epsilon, variance_size=None: (
21
22
23
@ir.ops.rms_norm.register_impl(
23
- "vllm_c", supports_args=rms_no_var_size, supported=CUDA_ALIKE
24
+ "vllm_c", supports_args=rms_no_var_size, supported=GPGPU_DEVICE
25
)
26
def rms_norm(
27
x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None
@@ -53,7 +54,7 @@ matching input/weight dtype."""
54
@ir.ops.fused_add_rms_norm.register_impl(
55
"vllm_c",
56
supports_args=rms_add_no_var_size,
56
- supported=CUDA_ALIKE,
57
+ supported=GPGPU_DEVICE,
58
inplace=True,
59
)
60
def fused_add_rms_norm(
vllm/kernels/xpu_ops.py
deleted
-59
@@ -1,59 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import torch
4
-from torch import Tensor
5
-
6
-from vllm import ir
7
-from vllm.platforms import current_platform
8
-
9
-current_platform.import_kernels()
10
-
11
-
12
-def is_xpu_kernels_found() -> bool:
13
- from importlib.util import find_spec
14
-
15
- return find_spec("vllm_xpu_kernels") is not None
16
-
17
-
18
-XPU_KERNELS_SUPPORTED = is_xpu_kernels_found()
19
-"""Kernels in this file are supported if vLLM XPU kernels are installed."""
20
-
21
-rms_no_var = lambda x, weight, epsilon, variance_size=None: variance_size is None and (
22
- weight is None or weight.dtype == x.dtype
23
-)
24
-
25
-
26
-@ir.ops.rms_norm.register_impl(
27
- "xpu_kernels", supports_args=rms_no_var, supported=XPU_KERNELS_SUPPORTED
28
-)
29
-def rms_norm(
30
- x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None
31
-) -> Tensor:
32
- assert variance_size is None
33
- output = torch.empty(x.shape, device=x.device, dtype=x.dtype)
34
- torch.ops._C.rms_norm(output, x, weight, epsilon)
35
- return output
36
-
37
-
38
-rms_add_no_var_size = (
39
- lambda x, x_residual, weight, epsilon, variance_size=None: variance_size is None
40
- and (weight is None or weight.dtype == x.dtype)
41
-)
42
-
43
-
44
-@ir.ops.fused_add_rms_norm.register_impl(
45
- "xpu_kernels",
46
- supports_args=rms_add_no_var_size,
47
- supported=XPU_KERNELS_SUPPORTED,
48
- inplace=True,
49
-)
50
-def fused_add_rms_norm(
51
- x: Tensor,
52
- x_residual: Tensor,
53
- weight: Tensor | None,
54
- epsilon: float,
55
- variance_size: int | None = None,
56
-) -> tuple[Tensor, Tensor]:
57
- assert variance_size is None
58
- torch.ops._C.fused_add_rms_norm(x, x_residual, weight, epsilon)
59
- return x, x_residual
vllm/platforms/xpu.py
+1
-1
@@ -465,7 +465,7 @@ class XPUPlatform(Platform):
465
# use fused kernels where available when no codegen
466
cc = vllm_config.compilation_config
467
using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE
468
- default = ["native"] if using_inductor else ["xpu_kernels", "native"]
468
+ default = ["native"] if using_inductor else ["vllm_c", "native"]
469
470
return IrOpPriorityConfig.with_default(default)
471