[XPU] Add W8A8 FP8 linear kernel with multi-granularity quant support (#43645)
Signed-off-by: Chaojun,Zhang <chaojun.zhang@intel.com> Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Chaojun Zhang committed
Jul 4, 2026 at 18:10 UTC
1a308c449c93d25b11dc79a199ed67db185a18df
6 files changed
+188
-12
.buildkite/intel_jobs/test-intel.yaml
+48
-1
@@ -49,6 +49,30 @@ steps:
49
VLLM_XPU_FUSED_MOE_USE_REF=1 python3 examples/basic/offline_inference/generate.py --model Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 --enforce-eager -tp 2 --max-model-len 8192 &&
50
python3 examples/basic/offline_inference/generate.py --model INCModel/Qwen3-30B-A3B-Instruct-2507-MXFP4-LLMC --enforce-eager -tp 2 --max-model-len 8192
51
'
52
+ - label: "XPU W8A8 FP8 Linear Examples"
53
+ depends_on:
54
+ - image-build-xpu
55
+ timeout_in_minutes: 60
56
+ device: intel_gpu
57
+ agent_tags:
58
+ label: production
59
+ gpu: 1+
60
+ mem: 24+
61
+ no_plugin: true
62
+ env:
63
+ REGISTRY: "public.ecr.aws/q9t5s3a7"
64
+ REPO: "vllm-ci-test-repo"
65
+ VLLM_TEST_DEVICE: "xpu"
66
+ source_file_dependencies:
67
+ - vllm/
68
+ - .buildkite/intel_jobs/test-intel.yaml
69
+ commands:
70
+ - >-
71
+ bash .buildkite/scripts/hardware_ci/run-intel-test.sh
72
+ 'python3 examples/basic/offline_inference/generate.py --linear-backend xpu --model RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8 --enforce-eager --max-model-len 4096 &&
73
+ python3 examples/basic/offline_inference/generate.py --linear-backend xpu --model neuralmagic/Llama-3.2-1B-Instruct-FP8-dynamic --enforce-eager --max-model-len 4096 &&
74
+ python3 examples/basic/offline_inference/generate.py --linear-backend xpu --model meta-llama/Llama-3.2-1B-Instruct --quantization fp8 --enforce-eager --max-model-len 4096
75
+ '
76
- label: "XPU V1 test"
77
depends_on:
78
- image-build-xpu
@@ -120,4 +144,27 @@ steps:
144
- >-
145
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
146
'cd tests &&
123
- pytest -v -s quantization/test_auto_round.py'
\ No newline at end of file
147
+ pytest -v -s quantization/test_auto_round.py'
148
+ - label: "XPU compressed tensors FP8 test"
149
+ depends_on:
150
+ - image-build-xpu
151
+ timeout_in_minutes: 60
152
+ device: intel_gpu
153
+ agent_tags:
154
+ label: production
155
+ gpu: 1+
156
+ mem: 16+
157
+ no_plugin: true
158
+ env:
159
+ REGISTRY: "public.ecr.aws/q9t5s3a7"
160
+ REPO: "vllm-ci-test-repo"
161
+ VLLM_TEST_DEVICE: "xpu"
162
+ source_file_dependencies:
163
+ - vllm/
164
+ - tests/quantization/test_compressed_tensors.py
165
+ - .buildkite/intel_jobs/test-intel.yaml
166
+ commands:
167
+ - >-
168
+ bash .buildkite/scripts/hardware_ci/run-intel-test.sh
169
+ 'cd tests &&
170
+ pytest -v -s quantization/test_compressed_tensors.py::test_compressed_tensors_fp8'
\ No newline at end of file
docs/features/quantization/online.md
+2
@@ -62,6 +62,8 @@ weight name. Unset fields fall back to the `--quantization` shorthand's
62
defaults, or for already-quantized checkpoints to whatever the checkpoint
63
declares.
64
65
+On XPU, non-block FP8 scaled-mm linear layers default to W8A16; setting `--linear-backend xpu` forces W8A8. Use `--linear-backend xpu_woq` to explicitly select weight-only quantization (W8A16).
66
+
67
The CLI accepts the same shape as JSON or as dotted keys:
68
69
```bash
vllm/config/kernel.py
+6
-1
@@ -156,6 +156,8 @@ LinearBackend = Literal[
156
"conch",
157
"exllama",
158
"emulation",
159
+ "xpu",
160
+ "xpu_woq",
161
]
162
163
@@ -217,7 +219,10 @@ class KernelConfig:
219
- "fbgemm": Use FBGEMM kernels
220
- "conch": Use Conch mixed-precision kernels
221
- "exllama": Use Exllama mixed-precision kernels
220
- - "emulation": Use slow dequant-to-BF16 emulation (for testing only)"""
222
+ - "emulation": Use slow dequant-to-BF16 emulation (for testing only)
223
+ - "xpu": Use XPU kernels
224
+ - "xpu_woq": Use XPU kernels for weight-only quantization (e.g. W8A16)
225
+ """
226
227
@field_validator("moe_backend", mode="before")
228
@classmethod
vllm/model_executor/kernels/linear/__init__.py
+12
-3
@@ -171,7 +171,8 @@ from vllm.model_executor.kernels.linear.scaled_mm.triton import (
171
)
172
from vllm.model_executor.kernels.linear.scaled_mm.xpu import (
173
XPUFp8BlockScaledMMKernel,
174
- XPUFP8ScaledMMLinearKernel,
174
+ XPUW8A8FP8LinearKernel,
175
+ XPUW8A16FP8LinearKernel,
176
)
177
from vllm.model_executor.kernels.linear.scaled_mm.zentorch import (
178
ZentorchInt8ScaledMMLinearKernel,
@@ -266,6 +267,13 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = {
267
EmulationMxfp8LinearKernel,
268
EmulationNvFp4LinearKernel,
269
},
270
+ "xpu": {
271
+ XPUW8A8FP8LinearKernel,
272
+ XPUFp8BlockScaledMMKernel,
273
+ },
274
+ "xpu_woq": {
275
+ XPUW8A16FP8LinearKernel,
276
+ },
277
}
278
279
@@ -311,7 +319,8 @@ _POSSIBLE_FP8_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] =
319
ChannelWiseTorchFP8ScaledMMLinearKernel,
320
],
321
PlatformEnum.XPU: [
314
- XPUFP8ScaledMMLinearKernel,
322
+ XPUW8A16FP8LinearKernel,
323
+ XPUW8A8FP8LinearKernel,
324
],
325
}
326
@@ -351,7 +360,7 @@ _POSSIBLE_WFP8A16_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]
360
# To be added
361
],
362
PlatformEnum.XPU: [
354
- XPUFP8ScaledMMLinearKernel,
363
+ XPUW8A16FP8LinearKernel,
364
],
365
}
366
vllm/model_executor/kernels/linear/scaled_mm/xpu.py
+107
-4
@@ -6,8 +6,11 @@ from collections.abc import Sequence
6
import torch
7
8
from vllm.model_executor.layers.quantization.utils.quant_utils import (
9
+ kFp8DynamicTensorSym,
10
+ kFp8DynamicTokenSym,
11
kFp8StaticChannelSym,
12
kFp8StaticTensorSym,
13
+ kFp8StaticTokenSym,
14
)
15
from vllm.model_executor.utils import replace_parameter
16
from vllm.platforms import current_platform
@@ -16,13 +19,112 @@ from .BlockScaledMMLinearKernel import Fp8BlockScaledMMLinearKernel
19
from .ScaledMMLinearKernel import FP8ScaledMMLinearKernel, FP8ScaledMMLinearLayerConfig
20
21
19
-class XPUFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
22
+class XPUW8A8FP8LinearKernel(FP8ScaledMMLinearKernel):
23
+ _SUPPORTED_ACT_QUANT_KEYS = {
24
+ kFp8DynamicTensorSym,
25
+ kFp8DynamicTokenSym,
26
+ kFp8StaticTensorSym,
27
+ kFp8StaticTokenSym,
28
+ }
29
+ _SUPPORTED_WEIGHT_QUANT_KEYS = {
30
+ kFp8StaticChannelSym,
31
+ kFp8StaticTensorSym,
32
+ }
33
+
34
+ @classmethod
35
+ def is_supported(
36
+ cls, compute_capability: int | None = None
37
+ ) -> tuple[bool, str | None]:
38
+ if not current_platform.is_xpu():
39
+ return False, "XPUW8A8FP8Linear only support on XPU"
40
+ return True, None
41
+
42
+ @classmethod
43
+ def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
44
+ if c.weight_quant_key not in cls._SUPPORTED_WEIGHT_QUANT_KEYS:
45
+ return (
46
+ False,
47
+ "XPUW8A8FP8Linear only support per-channel and per-tensor quantization",
48
+ )
49
+ if c.activation_quant_key not in cls._SUPPORTED_ACT_QUANT_KEYS:
50
+ return (
51
+ False,
52
+ "XPUW8A8FP8Linear only support per-tensor and per-token activation "
53
+ "quantization",
54
+ )
55
+ if c.weight_quant_key.dtype not in {torch.float8_e5m2, torch.float8_e4m3fn}:
56
+ return False, "XPUW8A8FP8Linear only support FP8 weight dtype"
57
+ if c.activation_quant_key.dtype not in {
58
+ torch.float8_e5m2,
59
+ torch.float8_e4m3fn,
60
+ }:
61
+ return False, "XPUW8A8FP8Linear only support FP8 activation dtype"
62
+ return True, None
63
+
64
+ def __init__(
65
+ self, c: FP8ScaledMMLinearLayerConfig, layer_param_names: Sequence[str]
66
+ ) -> None:
67
+ super().__init__(c, layer_param_names)
68
+
69
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
70
+ """Ensure weight is stored as C-contiguous [K, N] (KN layout).
71
+
72
+ Checkpoints store weight as [N, K]; fp8_gemm requires [K, N],
73
+ C-contiguous. Three incoming layouts are possible:
74
+ • [N, K] C-contiguous ← direct checkpoint → .t().contiguous()
75
+ • [K, N] Fortran-order ← fp8.py's weight.t() → .contiguous()
76
+ • [K, N] C-contiguous ← already correct → no-op
77
+
78
+ For square weights (K == N) the shape is ambiguous; contiguity is used
79
+ as a proxy: C-contiguous ≡ checkpoint [N, K] (needs transpose);
80
+ Fortran-order ≡ fp8.py already transposed (needs only contiguous).
81
+ """
82
+ K = getattr(layer, "input_size_per_partition", self.config.weight_shape[1])
83
+ N = getattr(layer, "output_size_per_partition", self.config.weight_shape[0])
84
+ w = layer.weight
85
+
86
+ if w.shape not in {(K, N), (N, K)}:
87
+ raise ValueError(
88
+ f"XPUFP8ScaledMM expects weight shape ({K},{N}) or ({N},{K}), "
89
+ f"but got {tuple(w.shape)}"
90
+ )
91
+
92
+ needs_transpose = w.shape == (N, K) if K != N else w.is_contiguous()
93
+ layer_weight = w.t() if needs_transpose else w
94
+ replace_parameter(layer, "weight", layer_weight)
95
+ ws = layer.weight_scale
96
+ if ws.numel() == 1:
97
+ replace_parameter(layer, "weight_scale", ws.reshape(1))
98
+
99
+ def apply_scaled_mm(
100
+ self,
101
+ *,
102
+ A: torch.Tensor,
103
+ B: torch.Tensor,
104
+ out_dtype: torch.dtype,
105
+ As: torch.Tensor,
106
+ Bs: torch.Tensor,
107
+ bias: torch.Tensor | None,
108
+ output_shape: list,
109
+ ) -> torch.Tensor:
110
+ # B is C-contiguous [K, N] from process_weights_after_loading.
111
+ # fp8_gemm routes on scale dtype (float32) and numel:
112
+ # As [1] → per-tensor (numel==1 branch)
113
+ # As [M,1] → per-token (group={1,K} branch, broadcast across K)
114
+ # Bs [1] → per-tensor
115
+ # Bs [N] → per-channel (mask=bit1 branch)
116
+ # No shape manipulation needed here.
117
+ output = torch.ops._xpu_C.fp8_gemm(A, B, out_dtype, As, Bs, bias)
118
+ return output.view(*output_shape)
119
+
120
+
121
+class XPUW8A16FP8LinearKernel(FP8ScaledMMLinearKernel):
122
@classmethod
123
def is_supported(
124
cls, compute_capability: int | None = None
125
) -> tuple[bool, str | None]:
126
if not current_platform.is_xpu():
25
- return False, "XPUFP8ScaledMM only support on XPU"
127
+ return False, "XPUW8A16FP8Linear only support on XPU"
128
return True, None
129
130
@classmethod
@@ -30,10 +132,11 @@ class XPUFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
132
if c.weight_quant_key not in {kFp8StaticChannelSym, kFp8StaticTensorSym}:
133
return (
134
False,
33
- "XPUFP8ScaledMM only support per-channel and per-tensor quantization",
135
+ "XPUW8A16FP8Linear only support per-channel and per-tensor "
136
+ "quantization",
137
)
138
if c.weight_quant_key.dtype not in {torch.float8_e5m2, torch.float8_e4m3fn}:
36
- return False, "XPUFP8ScaledMM only support FP8 weight dtype"
139
+ return False, "XPUW8A16FP8Linear only support FP8 weight dtype"
140
return True, None
141
142
def __init__(
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+13
-3
@@ -13,6 +13,7 @@ from compressed_tensors.quantization import (
13
)
14
from compressed_tensors.transform import TransformConfig
15
16
+from vllm.config import get_current_vllm_config_or_none
17
from vllm.distributed import (
18
get_tensor_model_parallel_rank,
19
get_tensor_model_parallel_world_size,
@@ -747,9 +748,18 @@ class CompressedTensorsConfig(QuantizationConfig):
748
act_quant_format = is_activation_quantization_format(format)
749
if act_quant_format:
750
if self._is_fp8_w8a8(weight_quant, input_quant):
750
- is_fp8_w8a8_supported = self._check_scheme_supported(
751
- CompressedTensorsW8A8Fp8.get_min_capability(), error=False
752
- )
751
+ if current_platform.is_xpu():
752
+ # On XPU, --linear-backend xpu opts into W8A8 FP8
753
+ # linear kernel; otherwise default to W8A16.
754
+ config = get_current_vllm_config_or_none()
755
+ is_fp8_w8a8_supported = (
756
+ config is not None
757
+ and config.kernel_config.linear_backend == "xpu"
758
+ )
759
+ else:
760
+ is_fp8_w8a8_supported = self._check_scheme_supported(
761
+ CompressedTensorsW8A8Fp8.get_min_capability(), error=False
762
+ )
763
if is_fp8_w8a8_supported:
764
return CompressedTensorsW8A8Fp8(
765
weight_quant=weight_quant,