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 3, 19:00 UTC
[XPU] [CT] Enable CT W4A4MxFp4 path and add xpu kernel (#38896)
Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com> Signed-off-by: zofia <110436990+zufangzhu@users.noreply.github.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
zofia committed
May 13, 2026 at 21:43 UTC
0ddaf6dffa0def89e2e734b02d9b69ffc78fd0ae
2 files changed
+59
vllm/model_executor/kernels/linear/__init__.py
+6
@@ -68,6 +68,9 @@ from vllm.model_executor.kernels.linear.mxfp4.flashinfer import (
68
from vllm.model_executor.kernels.linear.mxfp4.marlin import (
69
MarlinMxFp4LinearKernel,
70
)
71
+from vllm.model_executor.kernels.linear.mxfp4.xpu import (
72
+ XPUMxFp4LinearKernel,
73
+)
74
from vllm.model_executor.kernels.linear.mxfp8 import (
75
Mxfp8LinearKernel,
76
Mxfp8LinearLayerConfig,
@@ -291,6 +294,9 @@ _POSSIBLE_MXFP4_KERNELS: dict[PlatformEnum, list[type[MxFp4LinearKernel]]] = {
294
FlashInferMxFp4LinearKernel,
295
MarlinMxFp4LinearKernel,
296
],
297
+ PlatformEnum.XPU: [
298
+ XPUMxFp4LinearKernel,
299
+ ],
300
}
301
302
# TODO make all kernels inherit from MMLinearKernel
vllm/model_executor/kernels/linear/mxfp4/xpu.py
new
+53
@@ -0,0 +1,53 @@
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.mxfp4_utils import (
7
+ xpu_mxfp4_quantize as quant_mxfp4,
8
+)
9
+from vllm.model_executor.utils import replace_parameter
10
+from vllm.platforms import current_platform
11
+
12
+from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig
13
+
14
+
15
+class XPUMxFp4LinearKernel(MxFp4LinearKernel):
16
+ """MXFP4 W4A4 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, "XPUMxFp4 only support on XPU"
24
+ return True, None
25
+
26
+ @classmethod
27
+ def can_implement(cls, c: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
28
+ return True, None
29
+
30
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
31
+ weight = layer.weight.view(torch.float4_e2m1fn_x2)
32
+ replace_parameter(layer, "weight", weight.data.t())
33
+
34
+ weight_scale = layer.weight_scale.view(torch.float8_e8m0fnu)
35
+ weight_scale = weight_scale.t().contiguous()
36
+ replace_parameter(layer, "weight_scale", weight_scale.data)
37
+
38
+ def apply_weights(
39
+ self,
40
+ layer: torch.nn.Module,
41
+ x: torch.Tensor,
42
+ bias: torch.Tensor | None = None,
43
+ ) -> torch.Tensor:
44
+ out_dtype = x.dtype
45
+ x_fp4, x_blockscale = quant_mxfp4(x)
46
+ return torch.ops._xpu_C.fp4_gemm(
47
+ x_fp4,
48
+ layer.weight,
49
+ x_blockscale,
50
+ layer.weight_scale,
51
+ out_dtype,
52
+ bias,
53
+ )