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 14, 05:30 UTC
[XPU][Feature] fp8 online quantization support for XPU (#23148)
Signed-off-by: Yan Ma <yan.ma@intel.com> Co-authored-by: Qiming Zhang <qiming1.zhang@intel.com>
Yan Ma committed
Sep 2, 2025 at 12:06 UTC
7be0cb8e9e48849060660cf153205b4de1c1c854
4 files changed
+242
-2
vllm/_ipex_ops.py
+55
-1
@@ -1,11 +1,12 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from typing import Optional
4
+from typing import Optional, Union
5
6
import torch
7
8
from vllm.logger import init_logger
9
+from vllm.platforms import current_platform
10
11
logger = init_logger(__name__)
12
@@ -349,3 +350,56 @@ class ipex_ops:
350
def swap_blocks(src: torch.Tensor, dst: torch.Tensor,
351
block_mapping: torch.Tensor) -> None:
352
torch.xpu.swap_blocks(src, dst, block_mapping) # type: ignore
353
+
354
+ @staticmethod
355
+ def scaled_fp8_quant(
356
+ input: torch.Tensor,
357
+ scale: Optional[torch.Tensor] = None,
358
+ num_token_padding: Optional[int] = None,
359
+ scale_ub: Optional[torch.Tensor] = None,
360
+ use_per_token_if_dynamic: bool = False,
361
+ output: Optional[torch.Tensor] = None,
362
+ ) -> tuple[torch.Tensor, torch.Tensor]:
363
+ """
364
+ Quantize input tensor to FP8 and return quantized tensor and scale.
365
+
366
+ This function is designed for both static and dynamic quantization:
367
+ If you provide the scale, it will use static scaling and if you omit
368
+ it, the scale will be determined dynamically. Currently, XPU platform
369
+ only supports dynamic quantization. The function also allows optional
370
+ padding of the output tensors for downstream kernels that will benefit
371
+ from padding.
372
+
373
+ Args:
374
+ input: The input tensor to be quantized to FP8
375
+ scale: Optional scaling factor for the FP8 quantization
376
+ scale_ub: Optional upper bound for scaling factor in dynamic
377
+ per token case
378
+ num_token_padding: If specified, pad the first dimension
379
+ of the output to at least this value.
380
+ use_per_token_if_dynamic: Whether to do per_tensor or per_token
381
+ in the dynamic quantization case.
382
+
383
+ Returns:
384
+ tuple[torch.Tensor, torch.Tensor]: The output tensor in FP8 and
385
+ scaling factor.
386
+ """
387
+ # This code assumes batch_dim and num_tokens are flattened
388
+ assert (input.ndim == 2)
389
+ shape: Union[tuple[int, int], torch.Size] = input.shape
390
+ out_dtype: torch.dtype = current_platform.fp8_dtype()
391
+ if num_token_padding:
392
+ shape = (max(num_token_padding, input.shape[0]), shape[1])
393
+ if output is None:
394
+ output = torch.empty(shape, device=input.device, dtype=out_dtype)
395
+ else:
396
+ assert num_token_padding is None, \
397
+ "padding not supported if output passed in"
398
+ assert output.dtype == out_dtype
399
+ assert scale is None, "only dynamic fp8 quantization supported on XPU"
400
+ assert not use_per_token_if_dynamic, (
401
+ "per token dynamic fp8 quantization not supported on XPU")
402
+ scale = torch.zeros(1, device=input.device, dtype=torch.float32)
403
+ torch.ops.torch_ipex.dynamic_scaled_fp8_quant(output, input, scale)
404
+
405
+ return output, scale
vllm/model_executor/layers/quantization/fp8.py
+25
@@ -137,10 +137,35 @@ class Fp8Config(QuantizationConfig):
137
ignored_layers=ignored_layers,
138
weight_block_size=weight_block_size)
139
140
+ def get_xpu_quant_method(self, layer: torch.nn.Module,
141
+ prefix: str) -> Optional["QuantizeMethodBase"]:
142
+ from vllm.attention.layer import Attention
143
+ from vllm.model_executor.layers.quantization.ipex_quant import (
144
+ XPUFp8LinearMethod, XPUFp8MoEMethod)
145
+ fp8_config = Fp8Config(
146
+ is_checkpoint_fp8_serialized=self.is_checkpoint_fp8_serialized,
147
+ activation_scheme=self.activation_scheme,
148
+ ignored_layers=self.ignored_layers,
149
+ weight_block_size=self.weight_block_size)
150
+
151
+ if isinstance(layer, LinearBase):
152
+ if is_layer_skipped(prefix=prefix,
153
+ ignored_layers=self.ignored_layers,
154
+ fused_mapping=self.packed_modules_mapping):
155
+ return UnquantizedLinearMethod()
156
+ return XPUFp8LinearMethod(fp8_config)
157
+ elif isinstance(layer, FusedMoE):
158
+ return XPUFp8MoEMethod(fp8_config, layer)
159
+ elif isinstance(layer, Attention):
160
+ return Fp8KVCacheMethod(self)
161
+ return None
162
+
163
def get_quant_method(self, layer: torch.nn.Module,
164
prefix: str) -> Optional["QuantizeMethodBase"]:
165
from vllm.attention.layer import Attention # Avoid circular import
166
167
+ if current_platform.is_xpu():
168
+ return self.get_xpu_quant_method(layer, prefix)
169
if isinstance(layer, LinearBase):
170
if is_layer_skipped(prefix=prefix,
171
ignored_layers=self.ignored_layers,
vllm/model_executor/layers/quantization/ipex_quant.py
+158
-1
@@ -1,11 +1,16 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from typing import Any, Optional
4
+from typing import Any, Callable, Optional
5
6
import torch
7
from packaging import version
8
+from torch.nn import Module
9
+from torch.nn.parameter import Parameter
10
11
+from vllm._ipex_ops import ipex_ops as ops
12
+from vllm.model_executor.layers.fused_moe import (FusedMoEMethodBase,
13
+ FusedMoeWeightScaleSupported)
14
from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
15
UnquantizedLinearMethod)
16
from vllm.model_executor.layers.quantization import QuantizationMethods
@@ -13,7 +18,10 @@ from vllm.model_executor.layers.quantization.awq import (AWQLinearMethod,
18
is_layer_skipped_awq)
19
from vllm.model_executor.layers.quantization.base_config import (
20
QuantizationConfig)
21
+from vllm.model_executor.layers.quantization.fp8 import (Fp8Config,
22
+ Fp8LinearMethod)
23
from vllm.model_executor.layers.quantization.gptq import GPTQLinearMethod
24
+from vllm.model_executor.utils import set_weight_attrs
25
from vllm.platforms import current_platform
26
27
MIN_IPEX_VERSION = "2.6.0"
@@ -251,3 +259,152 @@ class IPEXAWQLinearMethod(AWQLinearMethod):
259
reshaped_x = x.reshape(-1, x.shape[-1])
260
out = layer.ipex_qlinear(reshaped_x)
261
return out.reshape(x.shape[:-1] + (layer.ipex_output_size, ))
262
+
263
+
264
+class XPUFp8LinearMethod(Fp8LinearMethod):
265
+
266
+ def __init__(self, quant_config: Fp8Config):
267
+ super().__init__(quant_config)
268
+
269
+ def process_weights_after_loading(self, layer: Module) -> None:
270
+ # If checkpoint not serialized fp8, quantize the weights.
271
+ if not self.quant_config.is_checkpoint_fp8_serialized:
272
+ qweight, weight_scale = ops.scaled_fp8_quant(layer.weight,
273
+ scale=None)
274
+ # Update the layer with the new values.
275
+ layer.weight = Parameter(qweight, requires_grad=False)
276
+ layer.weight_scale = Parameter(weight_scale, requires_grad=False)
277
+ layer.input_scale = None
278
+
279
+ def apply(self,
280
+ layer: torch.nn.Module,
281
+ x: torch.Tensor,
282
+ bias: Optional[torch.Tensor] = None) -> torch.Tensor:
283
+ weight = layer.weight.data
284
+ weight_scale = layer.weight_scale.data
285
+ output = torch.ops.torch_ipex.fp8_gemm_w8a16(x, weight, True,
286
+ weight_scale, bias)
287
+ return output
288
+
289
+
290
+class XPUFp8MoEMethod(FusedMoEMethodBase):
291
+
292
+ def __init__(self, quant_config: Fp8Config, layer: torch.nn.Module):
293
+ super().__init__(layer.moe_config)
294
+ self.quant_config = quant_config
295
+
296
+ def create_weights(self, layer: Module, num_experts: int, hidden_size: int,
297
+ intermediate_size_per_partition: int,
298
+ params_dtype: torch.dtype, **extra_weight_attrs):
299
+
300
+ layer.intermediate_size_per_partition = intermediate_size_per_partition
301
+ layer.hidden_size = hidden_size
302
+ layer.num_experts = num_experts
303
+ layer.orig_dtype = params_dtype
304
+ layer.weight_block_size = None
305
+ # WEIGHTS
306
+ w13_weight = torch.nn.Parameter(torch.empty(
307
+ num_experts,
308
+ 2 * intermediate_size_per_partition,
309
+ hidden_size,
310
+ dtype=params_dtype),
311
+ requires_grad=False)
312
+ layer.register_parameter("w13_weight", w13_weight)
313
+ set_weight_attrs(w13_weight, extra_weight_attrs)
314
+
315
+ w2_weight = torch.nn.Parameter(torch.empty(
316
+ num_experts,
317
+ hidden_size,
318
+ intermediate_size_per_partition,
319
+ dtype=params_dtype),
320
+ requires_grad=False)
321
+ layer.register_parameter("w2_weight", w2_weight)
322
+ set_weight_attrs(w2_weight, extra_weight_attrs)
323
+
324
+ # Allocate 2 scales for w1 and w3 respectively.
325
+ # They will be combined to a single scale after weight loading.
326
+ w13_weight_scale = torch.nn.Parameter(torch.ones(num_experts,
327
+ 2,
328
+ dtype=torch.float32),
329
+ requires_grad=False)
330
+ w2_weight_scale = torch.nn.Parameter(torch.ones(num_experts,
331
+ dtype=torch.float32),
332
+ requires_grad=False)
333
+ layer.register_parameter("w13_weight_scale", w13_weight_scale)
334
+ layer.register_parameter("w2_weight_scale", w2_weight_scale)
335
+
336
+ extra_weight_attrs.update(
337
+ {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value})
338
+ # INPUT_SCALES
339
+ layer.w13_input_scale = None
340
+ layer.w2_input_scale = None
341
+
342
+ def process_weights_after_loading(self, layer: Module) -> None:
343
+ if not self.quant_config.is_checkpoint_fp8_serialized:
344
+ fp8_dtype = current_platform.fp8_dtype()
345
+ w13_weight = torch.empty_like(layer.w13_weight.data,
346
+ dtype=fp8_dtype)
347
+ w2_weight = torch.empty_like(layer.w2_weight.data, dtype=fp8_dtype)
348
+
349
+ # Re-initialize w13_scale because we directly quantize
350
+ # merged w13 weights and generate a single scaling factor.
351
+ layer.w13_weight_scale = torch.nn.Parameter(torch.ones(
352
+ layer.local_num_experts,
353
+ dtype=torch.float32,
354
+ device=w13_weight.device),
355
+ requires_grad=False)
356
+ for expert in range(layer.local_num_experts):
357
+ w13_weight[expert, :, :], layer.w13_weight_scale[
358
+ expert] = ops.scaled_fp8_quant(
359
+ layer.w13_weight.data[expert, :, :])
360
+ w2_weight[expert, :, :], layer.w2_weight_scale[
361
+ expert] = ops.scaled_fp8_quant(
362
+ layer.w2_weight.data[expert, :, :])
363
+ layer.w13_weight = torch.nn.Parameter(w13_weight,
364
+ requires_grad=False)
365
+ layer.w2_weight = torch.nn.Parameter(w2_weight,
366
+ requires_grad=False)
367
+ import intel_extension_for_pytorch as ipex
368
+ layer.ipex_fusion = ipex.llm.modules.GatedMLPMOE(
369
+ layer.w13_weight,
370
+ layer.w2_weight,
371
+ w1_scale_inv=layer.w13_weight_scale,
372
+ w2_scale_inv=layer.w2_weight_scale,
373
+ a1_scale_inv=layer.w13_input_scale,
374
+ a2_scale_inv=layer.w2_input_scale,
375
+ use_prepack=True,
376
+ )
377
+
378
+ def apply(
379
+ self,
380
+ layer: torch.nn.Module,
381
+ x: torch.Tensor,
382
+ router_logits: torch.Tensor,
383
+ top_k: int,
384
+ renormalize: bool,
385
+ use_grouped_topk: bool = False,
386
+ topk_group: Optional[int] = None,
387
+ num_expert_group: Optional[int] = None,
388
+ global_num_experts: int = -1,
389
+ expert_map: Optional[torch.Tensor] = None,
390
+ custom_routing_function: Optional[Callable] = None,
391
+ scoring_func: str = "softmax",
392
+ routed_scaling_factor: float = 1.0,
393
+ e_score_correction_bias: Optional[torch.Tensor] = None,
394
+ apply_router_weight_on_input: bool = False,
395
+ activation: str = "silu",
396
+ enable_eplb: bool = False,
397
+ expert_load_view: Optional[torch.Tensor] = None,
398
+ logical_to_physical_map: Optional[torch.Tensor] = None,
399
+ logical_replica_count: Optional[torch.Tensor] = None,
400
+ ) -> torch.Tensor:
401
+ return layer.ipex_fusion(
402
+ x,
403
+ use_grouped_topk,
404
+ top_k,
405
+ router_logits,
406
+ renormalize,
407
+ topk_group,
408
+ num_expert_group,
409
+ custom_routing_function=custom_routing_function,
410
+ )
vllm/platforms/xpu.py
+4
@@ -148,6 +148,10 @@ class XPUPlatform(Platform):
148
torch.xpu.reset_peak_memory_stats(device)
149
return torch.xpu.max_memory_allocated(device)
150
151
+ @classmethod
152
+ def fp8_dtype(cls) -> torch.dtype:
153
+ return torch.float8_e5m2
154
+
155
@classmethod
156
def is_data_center_gpu(cls) -> bool:
157
device_name = cls.get_device_name().lower()