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 15, 04:45 UTC
Feat Dynamic Quantization for MoE Layers in GPTQ Marlin Backend (#19395)
Jun-Howie committed
Jun 24, 2025 at 06:23 UTC
dd2ccf8ddead5cdb46613b2adeb2edcc77f426c0
1 file changed
+29
-3
vllm/model_executor/layers/quantization/gptq_marlin.py
+29
-3
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from copy import deepcopy
5
from typing import Any, Callable, Optional, Union
6
7
import torch
@@ -9,7 +10,8 @@ import vllm.model_executor.layers.fused_moe # noqa
10
from vllm import _custom_ops as ops
11
from vllm.logger import init_logger
12
from vllm.model_executor.layers.fused_moe.layer import (
12
- FusedMoE, FusedMoEMethodBase, FusedMoeWeightScaleSupported)
13
+ FusedMoE, FusedMoEMethodBase, FusedMoeWeightScaleSupported,
14
+ UnquantizedFusedMoEMethod)
15
from vllm.model_executor.layers.linear import (LinearMethodBase,
16
set_weight_attrs)
17
from vllm.model_executor.layers.quantization import QuantizationMethods
@@ -19,7 +21,7 @@ from vllm.model_executor.layers.quantization.kernels.mixed_precision import (
21
MPLinearLayerConfig, choose_mp_linear_kernel)
22
from vllm.model_executor.layers.quantization.utils import replace_parameter
23
from vllm.model_executor.layers.quantization.utils.gptq_utils import (
22
- get_linear_quant_method)
24
+ get_dynamic_override, get_linear_quant_method, override_config)
25
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
26
check_marlin_supported, check_moe_marlin_supports_layer,
27
marlin_make_workspace_new, marlin_moe_permute_scales,
@@ -35,6 +37,29 @@ from vllm.scalar_type import scalar_types
37
logger = init_logger(__name__)
38
39
40
+def get_moe_quant_method(
41
+ config: QuantizationConfig,
42
+ layer: torch.nn.Module,
43
+ prefix: str,
44
+ moe_method_cls: type,
45
+):
46
+ cloned_config = deepcopy(config)
47
+
48
+ if isinstance(layer, FusedMoE):
49
+ # False = skip module, None = no override, else = Positive match
50
+ if get_dynamic_override( # noqa: E712
51
+ cloned_config, # noqa: E712
52
+ layer_name=prefix) == False: # noqa: E712
53
+ return UnquantizedFusedMoEMethod(layer.moe_config)
54
+
55
+ if prefix:
56
+ # Dynamic per module/layer rules may override base config
57
+ override_config(cloned_config, prefix=prefix)
58
+
59
+ return moe_method_cls(cloned_config)
60
+ return None
61
+
62
+
63
class GPTQMarlinConfig(QuantizationConfig):
64
"""Config class for GPTQ Marlin"""
65
@@ -163,7 +188,8 @@ class GPTQMarlinConfig(QuantizationConfig):
188
"Falling back to Moe WNA16 kernels.")
189
return MoeWNA16Config.from_config(
190
self.full_config).get_quant_method(layer, prefix)
166
- return GPTQMarlinMoEMethod(self)
191
+ return get_moe_quant_method(self, layer, prefix,
192
+ GPTQMarlinMoEMethod)
193
return get_linear_quant_method(self, layer, prefix,
194
GPTQMarlinLinearMethod)
195