[XPU] Add XPU block-scaled W8A8 fp8 path (#39968)

Signed-off-by: Wu, Xiaochang <xiaochang.wu@intel.com>
Signed-off-by: Xiaochang Wu <xiaochang.wu@intel.com>
Co-authored-by: Yuxiang <yuxiang.liang@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>

vllm/model_executor/kernels/linear/__init__.py

@@ -161,6 +161,7 @@ from vllm.model_executor.kernels.linear.scaled_mm.triton import (
     TritonInt8ScaledMMLinearKernel,
 )
 from vllm.model_executor.kernels.linear.scaled_mm.xpu import (
+    XPUFp8BlockScaledMMKernel,
     XPUFP8ScaledMMLinearKernel,
 )
 from vllm.model_executor.kernels.linear.scaled_mm.zentorch import (
@@ -317,6 +318,7 @@ _POSSIBLE_FP8_BLOCK_KERNELS: dict[
         CPUFp8BlockScaledMMKernel,
     ],
     PlatformEnum.XPU: [
+        XPUFp8BlockScaledMMKernel,
         TritonFp8BlockScaledMMKernel,
     ],
 }

vllm/model_executor/kernels/linear/scaled_mm/__init__.py

@@ -39,6 +39,9 @@ from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
 from vllm.model_executor.kernels.linear.scaled_mm.triton import (
     TritonInt8ScaledMMLinearKernel,
 )
+from vllm.model_executor.kernels.linear.scaled_mm.xpu import (
+    XPUFp8BlockScaledMMKernel,
+)
 from vllm.model_executor.kernels.linear.scaled_mm.zentorch import (
     ZentorchInt8ScaledMMLinearKernel,
 )
@@ -64,4 +67,5 @@ __all__ = [
     "ZentorchInt8ScaledMMLinearKernel",
     "Fp8BlockScaledMMLinearKernel",
     "CPUFp8BlockScaledMMKernel",
+    "XPUFp8BlockScaledMMKernel",
 ]

vllm/model_executor/kernels/linear/scaled_mm/triton.py

@@ -160,7 +160,7 @@ class TritonFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
     @classmethod
     def is_supported(cls, compute_capability=None):
         if not (current_platform.is_cuda_alike() or current_platform.is_xpu()):
-            return False, "only cuda-like and xpu devices are supported."
+            return False, "only CUDA-alike and XPU devices are supported."
         return True, None
 
     def apply_block_scaled_mm(

vllm/model_executor/kernels/linear/scaled_mm/xpu.py

@@ -5,10 +5,6 @@ from collections.abc import Sequence
 
 import torch
 
-from vllm.model_executor.kernels.linear import (  # noqa: E501
-    FP8ScaledMMLinearKernel,
-    FP8ScaledMMLinearLayerConfig,
-)
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     kFp8StaticChannelSym,
     kFp8StaticTensorSym,
@@ -16,6 +12,9 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
 from vllm.model_executor.utils import replace_parameter
 from vllm.platforms import current_platform
 
+from .BlockScaledMMLinearKernel import Fp8BlockScaledMMLinearKernel
+from .ScaledMMLinearKernel import FP8ScaledMMLinearKernel, FP8ScaledMMLinearLayerConfig
+
 
 class XPUFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
     @classmethod
@@ -84,3 +83,38 @@ class XPUFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
         output_shape: list,
     ) -> torch.Tensor:
         pass
+
+
+class XPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_xpu():
+            return False, "XPUFp8BlockScaledMM only support on XPU"
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        super().process_weights_after_loading(layer)
+        scale_attr = (
+            "weight_scale_inv" if hasattr(layer, "weight_scale_inv") else "weight_scale"
+        )
+        scale = getattr(layer, scale_attr)
+        replace_parameter(layer, scale_attr, scale.data.t().contiguous())
+
+    def apply_block_scaled_mm(
+        self,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+    ) -> torch.Tensor:
+        # Weight is [N, K]. Use .t() to create a [K, N] view without copying.
+        return torch.ops._xpu_C.fp8_gemm(
+            A,
+            B.t(),
+            self.config.out_dtype,
+            As,
+            Bs,
+            torch.Tensor(),
+        )