[XPU] Route activation CustomOps to SYCL kernels (#53734)
Signed-off-by: Marceli Fylcek <marceli.fylcek@intel.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Marceli Fylcek committed
Sep 1, 2026 at 03:19 UTC
4c21d4173358b45500769a95776500a34bf04589
1 file changed
+22
-4
vllm/model_executor/layers/activation.py
+22
-4
@@ -87,7 +87,7 @@ class FatreluAndMul(CustomOp):
87
def __init__(self, threshold: float = 0.0):
88
super().__init__()
89
self.threshold = threshold
90
- if current_platform.is_cuda_alike():
90
+ if current_platform.is_cuda_alike() or current_platform.is_xpu():
91
self.op = torch.ops._C.fatrelu_and_mul
92
elif current_platform.is_cpu():
93
self._forward_method = self.forward_native
@@ -106,6 +106,9 @@ class FatreluAndMul(CustomOp):
106
self.op(out, x, self.threshold)
107
return out
108
109
+ def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
110
+ return self.forward_cuda(x)
111
+
112
113
# --8<-- [start:silu_and_mul]
114
@CustomOp.register("silu_and_mul")
@@ -171,7 +174,7 @@ class SituAndMul(CustomOp):
174
super().__init__(compile_native=compile_native)
175
self.beta = float(beta)
176
self.linear_beta = None if linear_beta is None else float(linear_beta)
174
- if current_platform.is_cuda_alike():
177
+ if current_platform.is_cuda_alike() or current_platform.is_xpu():
178
self.op = torch.ops._C.situ_and_mul
179
180
def forward_native(self, x: torch.Tensor) -> torch.Tensor:
@@ -184,7 +187,7 @@ class SituAndMul(CustomOp):
187
return (gate * up).to(x.dtype)
188
189
def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
187
- # Fused CUDA kernel: writes straight to `out`, no fp32 temporaries.
190
+ # Fused CUDA/XPU kernel: writes straight to `out`, no fp32 temporaries.
191
# linear_beta<=0 signals "unset" to the kernel (up passed through).
192
d = x.shape[-1] // 2
193
out = torch.empty(x.shape[:-1] + (d,), dtype=x.dtype, device=x.device)
@@ -194,7 +197,7 @@ class SituAndMul(CustomOp):
197
return out
198
199
def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
197
- return self.forward_native(x)
200
+ return self.forward_cuda(x)
201
202
203
@CustomOp.register("silu_and_mul_with_clamp")
@@ -503,6 +506,9 @@ class SwigluOAIAndMul(CustomOp):
506
torch.ops._C.swigluoai_and_mul(out, x, self.alpha, self.limit)
507
return out
508
509
+ def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
510
+ return self.forward_cuda(x)
511
+
512
def extra_repr(self) -> str:
513
return f"alpha={repr(self.alpha)}, limit={repr(self.limit)}"
514
@@ -541,6 +547,13 @@ class SwigluStepAndMul(CustomOp):
547
swiglustep_and_mul_triton(out, x, self.limit)
548
return out
549
550
+ def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
551
+ d = x.shape[-1] // 2
552
+ output_shape = x.shape[:-1] + (d,)
553
+ out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
554
+ torch.ops._C.swiglustep_and_mul(out, x, self.limit)
555
+ return out
556
+
557
def extra_repr(self) -> str:
558
return f"limit={repr(self.limit)}"
559
@@ -656,6 +669,8 @@ class ReLUSquaredActivation(CustomOp):
669
super().__init__(compile_native=compile_native)
670
if current_platform.is_cuda_alike():
671
self.op = torch.ops._C.relu_squared
672
+ elif current_platform.is_xpu():
673
+ self.op = torch.ops._C.relu2_no_mul
674
675
def forward_native(self, x: torch.Tensor) -> torch.Tensor:
676
"""PyTorch-native implementation equivalent to forward()."""
@@ -666,6 +681,9 @@ class ReLUSquaredActivation(CustomOp):
681
self.op(out, x)
682
return out
683
684
+ def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
685
+ return self.forward_cuda(x)
686
+
687
688
# --8<-- [start:xielu]
689
@CustomOp.register("xielu")