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 11, 03:45 UTC
[XPU] enable fp8 online streaming quantization (#30944)
Signed-off-by: Yan Ma <yan.ma@intel.com>
Yan Ma committed
Dec 20, 2025 at 21:45 UTC
560ae9638c00e761d6f1ca33882249dbfebbe8aa
2 files changed
+29
-107
vllm/model_executor/layers/quantization/fp8.py
+12
-2
@@ -124,11 +124,13 @@ def get_fp8_moe_backend(
124
block_quant: bool,
125
moe_parallel_config: FusedMoEParallelConfig,
126
with_lora_support: bool,
127
-) -> Fp8MoeBackend:
127
+) -> Fp8MoeBackend | None:
128
"""
129
Select the primary FP8 MoE backend
130
Note: Shape-specific fallbacks may still occur at runtime.
131
"""
132
+ if current_platform.is_xpu():
133
+ return None
134
if with_lora_support:
135
return Fp8MoeBackend.TRITON
136
# Prefer FlashInfer backends on supported GPUs; allow SM90 and SM100.
@@ -292,6 +294,13 @@ class Fp8Config(QuantizationConfig):
294
return UnquantizedLinearMethod()
295
return XPUFp8LinearMethod(fp8_config)
296
elif isinstance(layer, FusedMoE):
297
+ if is_layer_skipped(
298
+ prefix=prefix,
299
+ ignored_layers=self.ignored_layers,
300
+ fused_mapping=self.packed_modules_mapping,
301
+ ):
302
+ return UnquantizedFusedMoEMethod(layer.moe_config)
303
+
304
return XPUFp8MoEMethod(fp8_config, layer)
305
elif isinstance(layer, Attention):
306
return Fp8KVCacheMethod(self)
@@ -1107,7 +1116,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1116
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
1117
) -> mk.FusedMoEPrepareAndFinalize | None:
1118
if (
1110
- self.rocm_aiter_moe_enabled
1119
+ current_platform.is_xpu()
1120
+ or self.rocm_aiter_moe_enabled
1121
or self.use_marlin
1122
or self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM
1123
):
vllm/model_executor/layers/quantization/ipex_quant.py
+17
-105
@@ -6,13 +6,8 @@ from typing import Any, Optional
6
import torch
7
from packaging import version
8
from torch.nn import Module
9
-from torch.nn.parameter import Parameter
9
10
from vllm._ipex_ops import ipex_ops as ops
12
-from vllm.model_executor.layers.fused_moe import (
13
- FusedMoEMethodBase,
14
- FusedMoeWeightScaleSupported,
15
-)
11
from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
12
from vllm.model_executor.layers.linear import (
13
LinearBase,
@@ -24,14 +19,14 @@ from vllm.model_executor.layers.quantization import (
19
QuantizationMethods,
20
)
21
from vllm.model_executor.layers.quantization.awq import AWQLinearMethod
27
-from vllm.model_executor.layers.quantization.fp8 import Fp8Config, Fp8LinearMethod
22
+from vllm.model_executor.layers.quantization.fp8 import (
23
+ Fp8Config,
24
+ Fp8LinearMethod,
25
+ Fp8OnlineMoEMethod,
26
+)
27
from vllm.model_executor.layers.quantization.gptq import GPTQLinearMethod
28
from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped
30
-from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
31
- maybe_create_device_identity,
32
-)
33
-from vllm.model_executor.parameter import ModelWeightParameter
34
-from vllm.model_executor.utils import set_weight_attrs
29
+from vllm.model_executor.utils import replace_parameter
30
from vllm.platforms import current_platform
31
32
MIN_IPEX_VERSION = "2.6.0"
@@ -309,44 +304,15 @@ class XPUFp8LinearMethod(Fp8LinearMethod):
304
def __init__(self, quant_config: Fp8Config):
305
super().__init__(quant_config)
306
312
- def create_weights(
313
- self,
314
- layer: torch.nn.Module,
315
- input_size_per_partition: int,
316
- output_partition_sizes: list[int],
317
- input_size: int,
318
- output_size: int,
319
- params_dtype: torch.dtype,
320
- **extra_weight_attrs,
321
- ):
322
- maybe_create_device_identity()
323
-
324
- output_size_per_partition = sum(output_partition_sizes)
325
- weight_loader = extra_weight_attrs.get("weight_loader")
326
- layer.logical_widths = output_partition_sizes
327
- layer.input_size_per_partition = input_size_per_partition
328
- layer.output_size_per_partition = output_size_per_partition
329
- layer.orig_dtype = params_dtype
330
- layer.weight_block_size = None
331
- weight = ModelWeightParameter(
332
- data=torch.empty(
333
- output_size_per_partition,
334
- input_size_per_partition,
335
- dtype=params_dtype,
336
- ),
337
- input_dim=1,
338
- output_dim=0,
339
- weight_loader=weight_loader,
340
- )
341
- layer.register_parameter("weight", weight)
342
-
307
def process_weights_after_loading(self, layer: Module) -> None:
308
+ if getattr(layer, "_already_called_process_weights_after_loading", False):
309
+ return
310
# If checkpoint not serialized fp8, quantize the weights.
311
if not self.quant_config.is_checkpoint_fp8_serialized:
312
qweight, weight_scale = ops.scaled_fp8_quant(layer.weight, scale=None)
313
# Update the layer with the new values.
348
- layer.weight = Parameter(qweight, requires_grad=False)
349
- layer.weight_scale = Parameter(weight_scale, requires_grad=False)
314
+ replace_parameter(layer, "weight", qweight.data)
315
+ replace_parameter(layer, "weight_scale", weight_scale.data)
316
layer.input_scale = None
317
318
def apply(
@@ -363,69 +329,14 @@ class XPUFp8LinearMethod(Fp8LinearMethod):
329
return output
330
331
366
-class XPUFp8MoEMethod(FusedMoEMethodBase):
332
+class XPUFp8MoEMethod(Fp8OnlineMoEMethod):
333
def __init__(self, quant_config: Fp8Config, layer: torch.nn.Module):
368
- super().__init__(layer.moe_config)
334
+ super().__init__(quant_config, layer)
335
self.quant_config = quant_config
336
371
- def create_weights(
372
- self,
373
- layer: Module,
374
- num_experts: int,
375
- hidden_size: int,
376
- intermediate_size_per_partition: int,
377
- params_dtype: torch.dtype,
378
- **extra_weight_attrs,
379
- ):
380
- layer.intermediate_size_per_partition = intermediate_size_per_partition
381
- layer.hidden_size = hidden_size
382
- layer.num_experts = num_experts
383
- layer.orig_dtype = params_dtype
384
- layer.weight_block_size = None
385
- # WEIGHTS
386
- w13_weight = torch.nn.Parameter(
387
- torch.empty(
388
- num_experts,
389
- 2 * intermediate_size_per_partition,
390
- hidden_size,
391
- dtype=params_dtype,
392
- ),
393
- requires_grad=False,
394
- )
395
- layer.register_parameter("w13_weight", w13_weight)
396
- set_weight_attrs(w13_weight, extra_weight_attrs)
397
-
398
- w2_weight = torch.nn.Parameter(
399
- torch.empty(
400
- num_experts,
401
- hidden_size,
402
- intermediate_size_per_partition,
403
- dtype=params_dtype,
404
- ),
405
- requires_grad=False,
406
- )
407
- layer.register_parameter("w2_weight", w2_weight)
408
- set_weight_attrs(w2_weight, extra_weight_attrs)
409
-
410
- # Allocate 2 scales for w1 and w3 respectively.
411
- # They will be combined to a single scale after weight loading.
412
- w13_weight_scale = torch.nn.Parameter(
413
- torch.ones(num_experts, 2, dtype=torch.float32), requires_grad=False
414
- )
415
- w2_weight_scale = torch.nn.Parameter(
416
- torch.ones(num_experts, dtype=torch.float32), requires_grad=False
417
- )
418
- layer.register_parameter("w13_weight_scale", w13_weight_scale)
419
- layer.register_parameter("w2_weight_scale", w2_weight_scale)
420
-
421
- extra_weight_attrs.update(
422
- {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
423
- )
424
- # INPUT_SCALES
425
- layer.w13_input_scale = None
426
- layer.w2_input_scale = None
427
-
337
def process_weights_after_loading(self, layer: Module) -> None:
338
+ if getattr(layer, "_already_called_process_weights_after_loading", False):
339
+ return
340
if not self.quant_config.is_checkpoint_fp8_serialized:
341
fp8_dtype = current_platform.fp8_dtype()
342
w13_weight = torch.empty_like(layer.w13_weight.data, dtype=fp8_dtype)
@@ -448,8 +359,9 @@ class XPUFp8MoEMethod(FusedMoEMethodBase):
359
w2_weight[expert, :, :], layer.w2_weight_scale[expert] = (
360
ops.scaled_fp8_quant(layer.w2_weight.data[expert, :, :])
361
)
451
- layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False)
452
- layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False)
362
+ replace_parameter(layer, "w13_weight", w13_weight)
363
+ replace_parameter(layer, "w2_weight", w2_weight)
364
+
365
import intel_extension_for_pytorch as ipex
366
367
ep_rank_start = self.moe.ep_rank * self.moe.num_local_experts