diff --git a/vllm/_xpu_ops.py b/vllm/_xpu_ops.py index 4cdba1e745..288dbdae9b 100644 --- a/vllm/_xpu_ops.py +++ b/vllm/_xpu_ops.py @@ -56,6 +56,24 @@ if hasattr(torch.ops._xpu_C, "fp8_gemm_w8a16"): return torch.empty((M, N), dtype=input.dtype, device=input.device) +if hasattr(torch.ops._xpu_C, "fp4_gemm"): + + @register_fake("_xpu_C::fp4_gemm") + def _fp4_gemm_fake( + input: torch.Tensor, + weight: torch.Tensor, + scale_act: torch.Tensor, + scale_wei: torch.Tensor, + out_dtype: torch.dtype | None = None, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + input_2d = input.view(-1, input.shape[-1]) + M = input_2d.size(0) + N = weight.size(-1) + result_dtype = out_dtype if out_dtype is not None else torch.get_default_dtype() + return torch.empty((M, N), dtype=result_dtype, device=input.device) + + if hasattr(torch.ops._xpu_C, "int4_gemm_w4a8"): @register_fake("_xpu_C::int4_gemm_w4a8")