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 5, 06:00 UTC
[MXFP8] [XPU] add a new compressed tensor schema and add a xpu mxfp8 gemm kernel (#38707)
Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com>
zofia committed
Apr 13, 2026 at 16:59 UTC
1b19bd758936496751432eccabf8adb7b5d8936a
2 files changed
+59
vllm/model_executor/kernels/linear/__init__.py
+8
@@ -71,6 +71,9 @@ from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
71
from vllm.model_executor.kernels.linear.mxfp8.marlin import (
72
MarlinMxfp8LinearKernel,
73
)
74
+from vllm.model_executor.kernels.linear.mxfp8.xpu import (
75
+ XPUMxFp8LinearKernel,
76
+)
77
from vllm.model_executor.kernels.linear.nvfp4 import (
78
NvFp4LinearKernel,
79
NvFp4LinearLayerConfig,
@@ -243,6 +246,10 @@ _POSSIBLE_MXFP8_KERNELS: dict[PlatformEnum, list[type[Mxfp8LinearKernel]]] = {
246
PlatformEnum.ROCM: [
247
EmulationMxfp8LinearKernel,
248
],
249
+ PlatformEnum.XPU: [
250
+ XPUMxFp8LinearKernel,
251
+ EmulationMxfp8LinearKernel,
252
+ ],
253
}
254
255
_POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = {
@@ -742,6 +749,7 @@ __all__ = [
749
"Mxfp8LinearLayerConfig",
750
"FlashInferCutlassMxfp8LinearKernel",
751
"MarlinMxfp8LinearKernel",
752
+ "XPUMxFp8LinearKernel",
753
"EmulationMxfp8LinearKernel",
754
"CutlassNvFp4LinearKernel",
755
"EmulationNvFp4LinearKernel",
vllm/model_executor/kernels/linear/mxfp8/xpu.py
new
+51
@@ -0,0 +1,51 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
7
+ xpu_mxfp8_quantize as quant_mxfp8,
8
+)
9
+from vllm.model_executor.utils import replace_parameter
10
+from vllm.platforms import current_platform
11
+
12
+from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
13
+
14
+
15
+class XPUMxFp8LinearKernel(Mxfp8LinearKernel):
16
+ """MXFP8 W8A8 GEMM on XPU."""
17
+
18
+ @classmethod
19
+ def is_supported(
20
+ cls, compute_capability: int | None = None
21
+ ) -> tuple[bool, str | None]:
22
+ if not current_platform.is_xpu():
23
+ return False, "XPUMxFp8 only support on XPU"
24
+ return True, None
25
+
26
+ @classmethod
27
+ def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
28
+ return True, None
29
+
30
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
31
+ weight_scale = layer.weight_scale.view(torch.float8_e8m0fnu)
32
+ weight_scale = weight_scale.t().contiguous()
33
+ replace_parameter(layer, "weight", layer.weight.t())
34
+ replace_parameter(layer, "weight_scale", weight_scale.data)
35
+
36
+ def apply_weights(
37
+ self,
38
+ layer: torch.nn.Module,
39
+ x: torch.Tensor,
40
+ bias: torch.Tensor | None = None,
41
+ ) -> torch.Tensor:
42
+ out_dtype = x.dtype
43
+ x_fp8, x_scale = quant_mxfp8(x)
44
+ return torch.ops._xpu_C.fp8_gemm(
45
+ x_fp8,
46
+ layer.weight,
47
+ out_dtype,
48
+ x_scale,
49
+ layer.weight_scale,
50
+ bias,
51
+ )