fix act fun for xpu (#23809)

This commit is contained in:
sushil Dubey
2026-05-21 14:02:21 +08:00
committed by GitHub
parent 8fa56a0ab1
commit c4f14650b9
2 changed files with 13 additions and 1 deletions
@@ -16,9 +16,11 @@ from sglang.multimodal_gen.runtime.platforms import current_platform
_is_cuda = current_platform.is_cuda()
_is_hip = current_platform.is_hip()
_is_npu = current_platform.is_npu()
_is_xpu = current_platform.is_xpu()
if _is_cuda:
from sglang.jit_kernel.activation import silu_and_mul
elif _is_hip:
elif _is_hip or _is_xpu:
from sgl_kernel import silu_and_mul
@@ -61,6 +63,13 @@ class SiluAndMul(CustomOp):
def forward_musa(self, x: torch.Tensor) -> torch.Tensor:
return nn.SwishGLU()(x)
def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
d = x.shape[-1] // 2
output_shape = x.shape[:-1] + (d,)
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
silu_and_mul(x, out)
return out
@CustomOp.register("gelu_and_mul")
class GeluAndMul(CustomOp):
@@ -61,6 +61,9 @@ class RotaryEmbedding(CustomOp):
def forward_cuda(self, *args, **kwargs):
return self.forward_native(*args, **kwargs)
def forward_xpu(self, *args, **kwargs):
return self.forward_native(*args, **kwargs)
def forward_native(
self,
positions: torch.Tensor,