A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 02:15 UTC
[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>
Xiaochang Wu committed
Jun 3, 2026 at 20:16 UTC
e5232679a349fc8a8f09ca8434b9bceb7b829642
4 files changed
+45
-5
vllm/model_executor/kernels/linear/__init__.py
+2
@@ -161,6 +161,7 @@ from vllm.model_executor.kernels.linear.scaled_mm.triton import (
161
TritonInt8ScaledMMLinearKernel,
162
)
163
from vllm.model_executor.kernels.linear.scaled_mm.xpu import (
164
+ XPUFp8BlockScaledMMKernel,
165
XPUFP8ScaledMMLinearKernel,
166
)
167
from vllm.model_executor.kernels.linear.scaled_mm.zentorch import (
@@ -317,6 +318,7 @@ _POSSIBLE_FP8_BLOCK_KERNELS: dict[
318
CPUFp8BlockScaledMMKernel,
319
],
320
PlatformEnum.XPU: [
321
+ XPUFp8BlockScaledMMKernel,
322
TritonFp8BlockScaledMMKernel,
323
],
324
}
vllm/model_executor/kernels/linear/scaled_mm/__init__.py
+4
@@ -39,6 +39,9 @@ from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
39
from vllm.model_executor.kernels.linear.scaled_mm.triton import (
40
TritonInt8ScaledMMLinearKernel,
41
)
42
+from vllm.model_executor.kernels.linear.scaled_mm.xpu import (
43
+ XPUFp8BlockScaledMMKernel,
44
+)
45
from vllm.model_executor.kernels.linear.scaled_mm.zentorch import (
46
ZentorchInt8ScaledMMLinearKernel,
47
)
@@ -64,4 +67,5 @@ __all__ = [
67
"ZentorchInt8ScaledMMLinearKernel",
68
"Fp8BlockScaledMMLinearKernel",
69
"CPUFp8BlockScaledMMKernel",
70
+ "XPUFp8BlockScaledMMKernel",
71
]
vllm/model_executor/kernels/linear/scaled_mm/triton.py
+1
-1
@@ -160,7 +160,7 @@ class TritonFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
160
@classmethod
161
def is_supported(cls, compute_capability=None):
162
if not (current_platform.is_cuda_alike() or current_platform.is_xpu()):
163
- return False, "only cuda-like and xpu devices are supported."
163
+ return False, "only CUDA-alike and XPU devices are supported."
164
return True, None
165
166
def apply_block_scaled_mm(
vllm/model_executor/kernels/linear/scaled_mm/xpu.py
+38
-4
@@ -5,10 +5,6 @@ from collections.abc import Sequence
5
6
import torch
7
8
-from vllm.model_executor.kernels.linear import ( # noqa: E501
9
- FP8ScaledMMLinearKernel,
10
- FP8ScaledMMLinearLayerConfig,
11
-)
8
from vllm.model_executor.layers.quantization.utils.quant_utils import (
9
kFp8StaticChannelSym,
10
kFp8StaticTensorSym,
@@ -16,6 +12,9 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
12
from vllm.model_executor.utils import replace_parameter
13
from vllm.platforms import current_platform
14
15
+from .BlockScaledMMLinearKernel import Fp8BlockScaledMMLinearKernel
16
+from .ScaledMMLinearKernel import FP8ScaledMMLinearKernel, FP8ScaledMMLinearLayerConfig
17
+
18
19
class XPUFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
20
@classmethod
@@ -84,3 +83,38 @@ class XPUFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
83
output_shape: list,
84
) -> torch.Tensor:
85
pass
86
+
87
+
88
+class XPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
89
+ @classmethod
90
+ def is_supported(
91
+ cls, compute_capability: int | None = None
92
+ ) -> tuple[bool, str | None]:
93
+ if not current_platform.is_xpu():
94
+ return False, "XPUFp8BlockScaledMM only support on XPU"
95
+ return True, None
96
+
97
+ def process_weights_after_loading(self, layer: torch.nn.Module):
98
+ super().process_weights_after_loading(layer)
99
+ scale_attr = (
100
+ "weight_scale_inv" if hasattr(layer, "weight_scale_inv") else "weight_scale"
101
+ )
102
+ scale = getattr(layer, scale_attr)
103
+ replace_parameter(layer, scale_attr, scale.data.t().contiguous())
104
+
105
+ def apply_block_scaled_mm(
106
+ self,
107
+ A: torch.Tensor,
108
+ B: torch.Tensor,
109
+ As: torch.Tensor,
110
+ Bs: torch.Tensor,
111
+ ) -> torch.Tensor:
112
+ # Weight is [N, K]. Use .t() to create a [K, N] view without copying.
113
+ return torch.ops._xpu_C.fp8_gemm(
114
+ A,
115
+ B.t(),
116
+ self.config.out_dtype,
117
+ As,
118
+ Bs,
119
+ torch.Tensor(),
120
+ )