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 10, 04:00 UTC
[Docs] Add docs about OOT Quantization Plugins (#32035)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Jan 14, 2026 at 02:25 UTC
6388b50058eaffe357f94ff0ab3dd0734c86b0e4
2 files changed
+158
docs/features/quantization/README.md
+157
@@ -67,3 +67,160 @@ th:not(:first-child) {
67
This compatibility chart is subject to change as vLLM continues to evolve and expand its support for different hardware platforms and quantization methods.
68
69
For the most up-to-date information on hardware support and quantization methods, please refer to [vllm/model_executor/layers/quantization](../../../vllm/model_executor/layers/quantization) or consult with the vLLM development team.
70
+
71
+## Out-of-Tree Quantization Plugins
72
+
73
+vLLM supports registering custom, out-of-tree quantization methods using the `@register_quantization_config` decorator. This allows you to implement and use your own quantization schemes without modifying the vLLM codebase.
74
+
75
+### Registering a Custom Quantization Method
76
+
77
+To register a custom quantization method, create a class that inherits from `QuantizationConfig` and decorate it with `@register_quantization_config`. The `get_quant_method` dispatches to the appropriate quantize method based on the layer type:
78
+
79
+```python
80
+import torch
81
+from vllm.model_executor.layers.quantization import (
82
+ register_quantization_config,
83
+)
84
+from vllm.model_executor.layers.quantization.base_config import (
85
+ QuantizationConfig,
86
+ QuantizeMethodBase,
87
+)
88
+from vllm.model_executor.layers.linear import LinearBase
89
+from vllm.model_executor.layers.fused_moe import FusedMoE
90
+
91
+@register_quantization_config("my_quant")
92
+class MyQuantConfig(QuantizationConfig):
93
+ """Custom quantization config."""
94
+
95
+ def get_name(self) -> str:
96
+ return "my_quant"
97
+
98
+ def get_supported_act_dtypes(self) -> list:
99
+ return [torch.float16, torch.bfloat16]
100
+
101
+ @classmethod
102
+ def get_min_capability(cls) -> int:
103
+ # Minimum GPU compute capability, -1 for no restriction
104
+ return -1
105
+
106
+ @staticmethod
107
+ def get_config_filenames() -> list[str]:
108
+ # Config files to search for in model directory
109
+ return []
110
+
111
+ @classmethod
112
+ def from_config(cls, config: dict) -> "MyQuantConfig":
113
+ # Create config from model's quantization config
114
+ return cls()
115
+
116
+ def get_quant_method(
117
+ self, layer: torch.nn.Module, prefix: str
118
+ ) -> QuantizeMethodBase | None:
119
+ # Dispatch based on layer type
120
+ # NOTE: you only need to implement methods you care about
121
+ if isinstance(layer, LinearBase):
122
+ return MyQuantLinearMethod()
123
+ elif isinstance(layer, FusedMoE):
124
+ return MyQuantMoEMethod(layer.moe_config)
125
+ return None
126
+```
127
+
128
+### Required QuantizationConfig Methods
129
+
130
+Your custom `QuantizationConfig` subclass must implement these abstract methods:
131
+
132
+| Method | Description |
133
+|--------|-------------|
134
+| `get_name()` | Returns the name of the quantization method |
135
+| `get_supported_act_dtypes()` | Returns list of supported activation dtypes (e.g., `torch.float16`) |
136
+| `get_min_capability()` | Returns minimum GPU compute capability (e.g., 80 for Ampere, -1 for no restriction) |
137
+| `get_config_filenames()` | Returns list of config filenames to search for in model directory |
138
+| `from_config(config)` | Class method to create config from model's quantization config dict |
139
+| `get_quant_method(layer, prefix)` | Returns the quantization method for a given layer, or `None` to skip |
140
+
141
+### Implementing a Quantized Linear Method
142
+
143
+For linear layers, return a `QuantizeMethodBase` subclass from `get_quant_method`. You can extend `UnquantizedLinearMethod` as a starting point:
144
+
145
+```python
146
+from vllm.model_executor.layers.linear import UnquantizedLinearMethod
147
+
148
+class MyQuantLinearMethod(UnquantizedLinearMethod):
149
+ """Custom quantization method for linear layers."""
150
+
151
+ def create_weights(
152
+ self, layer: torch.nn.Module, *weight_args, **extra_weight_attrs
153
+ ):
154
+ # Create quantized weights for the layer
155
+ ...
156
+
157
+ def apply(
158
+ self,
159
+ layer: torch.nn.Module,
160
+ x: torch.Tensor,
161
+ bias: torch.Tensor | None = None,
162
+ ) -> torch.Tensor:
163
+ # Apply custom quantization logic here
164
+ ...
165
+```
166
+
167
+### Implementing a Quantized MoE Method
168
+
169
+For Mixture of Experts (MoE) models, return a `FusedMoEMethodBase` subclass from `get_quant_method`. You can use `UnquantizedFusedMoEMethod` to skip MoE quantization:
170
+
171
+```python
172
+from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod
173
+from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
174
+ FusedMoEMethodBase,
175
+)
176
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
177
+
178
+class MyQuantMoEMethod(FusedMoEMethodBase):
179
+ """Custom quantization method for MoE layers."""
180
+
181
+ def create_weights(
182
+ self,
183
+ layer: torch.nn.Module,
184
+ num_experts: int,
185
+ hidden_size: int,
186
+ intermediate_size_per_partition: int,
187
+ params_dtype: torch.dtype,
188
+ **extra_weight_attrs,
189
+ ):
190
+ # Create quantized weights for the MoE layer
191
+ ...
192
+
193
+ def apply(
194
+ self,
195
+ layer: torch.nn.Module,
196
+ router: "FusedMoERouter",
197
+ x: torch.Tensor,
198
+ router_logits: torch.Tensor,
199
+ ) -> torch.Tensor:
200
+ # Apply MoE computation with quantized weights
201
+ ...
202
+
203
+ def get_fused_moe_quant_config(
204
+ self, layer: torch.nn.Module
205
+ ) -> FusedMoEQuantConfig | None:
206
+ # Return the MoE quantization configuration
207
+ ...
208
+```
209
+
210
+See existing implementations like `Fp8MoEMethod` in `vllm/model_executor/layers/quantization/fp8.py` for reference.
211
+
212
+### Using the Plugin
213
+
214
+Once registered, you can use your custom quantization method with vLLM:
215
+
216
+```python
217
+# Register your quantization method (import the module containing your config)
218
+import my_quant_plugin
219
+
220
+from vllm import LLM
221
+
222
+# Use the custom quantization method
223
+llm = LLM(model="your-model", quantization="my_quant")
224
+```
225
+
226
+For more information on the plugin system, see the [Plugin System documentation](../../design/plugin_system.md).
vllm/model_executor/layers/quantization/__init__.py
+1
@@ -188,5 +188,6 @@ __all__ = [
188
"QuantizationConfig",
189
"QuantizationMethods",
190
"get_quantization_config",
191
+ "register_quantization_config",
192
"QUANTIZATION_METHODS",
193
]