register xpu custom op as torch.ops.vllm_xpu.xxx

Signed-off-by: Xinyu Chen <xinyu1.chen@intel.com>

Author: xinyu-intel

vllm/kernels/xpu_ops.py

@@ -1,10 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import functools
+
 import torch
 from torch import Tensor
+from torch.library import Library
 
 from vllm import ir
 from vllm.platforms import current_platform
+from vllm.utils.torch_utils import direct_register_custom_op
 
 current_platform.import_kernels()
 
@@ -18,6 +22,15 @@ def is_xpu_kernels_found() -> bool:
 XPU_KERNELS_SUPPORTED = is_xpu_kernels_found()
 """Kernels in this file are supported if vLLM XPU kernels are installed."""
 
+xpu_kernels_lib = Library("vllm_xpu", "FRAGMENT")
+"""
+This library holds torch ops for vLLM XPU kernels.
+"""
+direct_register_xpu_op = functools.partial(
+    direct_register_custom_op, target_lib=xpu_kernels_lib
+)
+"""Syntactic sugar for registering XPU custom ops."""
+
 rms_no_var = lambda x, weight, epsilon, variance_size=None: variance_size is None
 
 
@@ -31,6 +44,19 @@ def rms_norm(
         # Kernel requires weight tensor, pass ones
         weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
     assert variance_size is None
+    return torch.ops.vllm_xpu.rms_norm(x, weight, epsilon)
+
+
+def _rms_norm_impl(x: Tensor, weight: Tensor, epsilon: float) -> Tensor:
     output = torch.empty(x.shape, device=x.device, dtype=x.dtype)
     torch.ops._C.rms_norm(output, x, weight, epsilon)
     return output
+
+
+def _rms_norm_fake(x: Tensor, weight: Tensor, epsilon: float) -> Tensor:
+    return torch.empty_like(x)
+
+
+direct_register_xpu_op(
+    op_name="rms_norm", op_func=_rms_norm_impl, fake_impl=_rms_norm_fake
+)

vllm/platforms/xpu.py

@@ -285,7 +285,7 @@ def get_default_ir_op_priority(
         # use fused kernels where available when no codegen
         cc = vllm_config.compilation_config
         using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE
-        default = ["native"] if using_inductor else ["xpu_kernels", "vllm_c", "native"]
+        default = ["native"] if using_inductor else ["xpu_kernels", "native"]
 
         return IrOpPriorityConfig.with_default(default)