Deprecate old FP8 online MoE quantization class (#44514)
Signed-off-by: Yan Ma <yan.ma@intel.com>
Yan Ma committed
Jun 24, 2026 at 02:53 UTC
acce57d8dd8bed25543f8ece834a3757208c677c
3 files changed
+19
-137
tests/quantization/test_fp8.py
+12
-3
@@ -14,6 +14,9 @@ import torch
14
from tests.quantization.utils import is_quant_method_supported
15
from vllm import _custom_ops as ops
16
from vllm.config.model import ModelConfig
17
+from vllm.model_executor.kernels.linear.scaled_mm import (
18
+ MarlinFP8ScaledMMLinearKernel,
19
+)
20
from vllm.model_executor.layers.fused_moe import FusedMoE
21
from vllm.model_executor.layers.quantization.fp8 import (
22
Fp8Config,
@@ -46,7 +49,7 @@ MODELS = [
49
)
50
@pytest.mark.parametrize("model_id", MODELS)
51
@pytest.mark.parametrize(
49
- "force_marlin", [False] if current_platform.is_rocm() else [False, True]
52
+ "force_marlin", [True, False] if current_platform.is_cuda() else [False]
53
)
54
@pytest.mark.parametrize(
55
"use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
@@ -73,7 +76,7 @@ def test_model_load_and_run(
76
)
77
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
78
@pytest.mark.parametrize(
76
- "force_marlin", [False] if current_platform.is_rocm() else [False, True]
79
+ "force_marlin", [True, False] if current_platform.is_cuda() else [False]
80
)
81
@pytest.mark.parametrize(
82
"use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
@@ -110,14 +113,20 @@ def test_online_quantization(
113
assert attn._k_scale == 1.0
114
assert attn._v_scale == 1.0
115
113
- if current_platform.is_cuda():
116
+ if current_platform.is_cuda() or current_platform.is_xpu():
117
if current_platform.supports_fp8() and not force_marlin:
118
# For GPUs with hardware support, we keep weights in fp8
119
assert fc1.weight.dtype == torch.float8_e4m3fn
120
+ assert not isinstance(
121
+ fc1.quant_method.fp8_linear, MarlinFP8ScaledMMLinearKernel
122
+ )
123
else:
124
# For GPUs without hardware support, we pack the fp8 weights
125
# for weight-only quantization using Marlin kernels
126
assert fc1.weight.dtype == torch.int32
127
+ assert isinstance(
128
+ fc1.quant_method.fp8_linear, MarlinFP8ScaledMMLinearKernel
129
+ )
130
elif current_platform.is_rocm():
131
if current_platform.supports_fp8() and not force_marlin:
132
# For GPUs with hardware support, we keep weights in fp8
vllm/model_executor/layers/quantization/fp8.py
+6
-133
@@ -8,7 +8,6 @@ from torch.utils._python_dispatch import TorchDispatchMode
8
9
import vllm.envs as envs
10
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11
-from vllm import _custom_ops as ops
11
from vllm.config import get_current_vllm_config
12
from vllm.distributed import get_tensor_model_parallel_world_size
13
from vllm.logger import init_logger
@@ -75,9 +74,6 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
74
cutlass_fp8_supported,
75
normalize_e4m3fn_to_e4m3fnuz,
76
)
78
-from vllm.model_executor.model_loader.reload.layerwise import (
79
- initialize_online_processing,
80
-)
77
from vllm.model_executor.parameter import (
78
BlockQuantScaleParameter,
79
PerTensorScaleParameter,
@@ -212,10 +208,13 @@ class Fp8Config(QuantizationConfig):
208
209
return Mxfp4MoEMethod(layer.moe_config)
210
if self.is_checkpoint_fp8_serialized:
215
- moe_quant_method = Fp8MoEMethod(self, layer)
211
+ return Fp8MoEMethod(self, layer)
212
else:
217
- moe_quant_method = Fp8OnlineMoEMethod(self, layer)
218
- return moe_quant_method
213
+ from vllm.model_executor.layers.quantization.online.fp8 import (
214
+ Fp8PerTensorOnlineMoEMethod,
215
+ )
216
+
217
+ return Fp8PerTensorOnlineMoEMethod(layer=layer)
218
elif isinstance(layer, Attention):
219
return Fp8KVCacheMethod(self)
220
return None
@@ -854,132 +853,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):
853
)
854
855
857
-# TODO(future PR): remove this class in favor of
858
-# online/fp8.py::Fp8PerTensorOnlineMoEMethod
859
-class Fp8OnlineMoEMethod(Fp8MoEMethod):
860
- """MoE method for online FP8 quantization.
861
- Supports loading quantized FP16/BF16 model checkpoints with dynamic
862
- activation scaling. The weight scaling factor will be initialized after
863
- the model weights are loaded.
864
-
865
- Args:
866
- quant_config: The quantization config.
867
- """
868
-
869
- uses_meta_device: bool = True
870
-
871
- def __init__(self, quant_config: Fp8Config, layer: RoutedExperts):
872
- super().__init__(quant_config, layer)
873
- assert not quant_config.is_checkpoint_fp8_serialized
874
- assert quant_config.activation_scheme == "dynamic"
875
- assert quant_config.weight_block_size is None
876
-
877
- def create_weights(
878
- self,
879
- layer: RoutedExperts,
880
- num_experts: int,
881
- hidden_size: int,
882
- intermediate_size_per_partition: int,
883
- params_dtype: torch.dtype,
884
- **extra_weight_attrs,
885
- ):
886
- layer.num_experts = num_experts
887
- layer.orig_dtype = params_dtype
888
- layer.weight_block_size = None
889
-
890
- # WEIGHTS
891
- w13_weight = torch.nn.Parameter(
892
- torch.empty(
893
- num_experts,
894
- 2 * intermediate_size_per_partition,
895
- hidden_size,
896
- device="meta",
897
- dtype=params_dtype,
898
- ),
899
- requires_grad=False,
900
- )
901
- layer.register_parameter("w13_weight", w13_weight)
902
- set_weight_attrs(w13_weight, extra_weight_attrs)
903
-
904
- w2_weight = torch.nn.Parameter(
905
- torch.empty(
906
- num_experts,
907
- hidden_size,
908
- intermediate_size_per_partition,
909
- device="meta", # materialized and processed during loading
910
- dtype=params_dtype,
911
- ),
912
- requires_grad=False,
913
- )
914
- layer.register_parameter("w2_weight", w2_weight)
915
- set_weight_attrs(w2_weight, extra_weight_attrs)
916
-
917
- # BIASES (for models like GPT-OSS that have biased MoE)
918
- if self.moe.has_bias:
919
- w13_bias = torch.nn.Parameter(
920
- torch.zeros(
921
- num_experts,
922
- 2 * intermediate_size_per_partition,
923
- device="meta", # materialized and processed during loading
924
- dtype=layer.orig_dtype,
925
- ),
926
- requires_grad=False,
927
- )
928
- layer.register_parameter("w13_bias", w13_bias)
929
- set_weight_attrs(w13_bias, extra_weight_attrs)
930
-
931
- w2_bias = torch.nn.Parameter(
932
- torch.zeros(
933
- num_experts,
934
- hidden_size,
935
- device="meta", # materialized and processed during loading
936
- dtype=layer.orig_dtype,
937
- ),
938
- requires_grad=False,
939
- )
940
- layer.register_parameter("w2_bias", w2_bias)
941
- set_weight_attrs(w2_bias, extra_weight_attrs)
942
-
943
- initialize_online_processing(layer)
944
-
945
- def process_weights_after_loading(self, layer: RoutedExperts) -> None:
946
- # TODO(@ksayers): inplace fp8 quant kernel, initialize scales with ones
947
- if getattr(layer, "_already_called_process_weights_after_loading", False):
948
- return
949
-
950
- fp8_dtype = current_platform.fp8_dtype()
951
- w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)
952
- w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)
953
- w13_scale = torch.ones(
954
- layer.num_experts, device=w13.device, dtype=torch.float32
955
- )
956
- w2_scale = torch.ones(layer.num_experts, device=w2.device, dtype=torch.float32)
957
- layer.w13_input_scale = None
958
- layer.w2_input_scale = None
959
-
960
- for expert in range(layer.local_num_experts):
961
- w13[expert, :, :], w13_scale[expert] = ops.scaled_fp8_quant(
962
- layer.w13_weight[expert, :, :]
963
- )
964
- w2[expert, :, :], w2_scale[expert] = ops.scaled_fp8_quant(
965
- layer.w2_weight[expert, :, :]
966
- )
967
-
968
- # Shuffle weights to runtime format and setup kernel.
969
- self._setup_kernel(
970
- layer,
971
- w13,
972
- w2,
973
- w13_scale,
974
- w2_scale,
975
- w13_input_scale=layer.w13_input_scale,
976
- w2_input_scale=layer.w2_input_scale,
977
- )
978
-
979
- # Prevent duplicate processing (e.g., during weight reload)
980
- layer._already_called_process_weights_after_loading = True
981
-
982
-
856
class Fp8KVCacheMethod(BaseKVCacheMethod):
857
"""
858
Supports loading kv-cache scaling factors from FP8 checkpoints.
vllm/model_executor/model_loader/base_loader.py
+1
-1
@@ -65,7 +65,7 @@ class BaseModelLoader(ABC):
65
66
# Log peak GPU memory after loading weights. This is needed
67
# to have test coverage on peak memory for online quantization.
68
- if current_platform.is_cuda_alike():
68
+ if current_platform.is_cuda_alike() or current_platform.is_xpu():
69
peak_memory = torch.accelerator.max_memory_allocated()
70
logger.debug_once(
71
"Peak GPU memory after loading weights: %s GiB",