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
[Quantization] Add compressed-tensors NVFP4 MoE Support (#19990)
Signed-off-by: Dipika Sikka <dipikasikka1@gmail.com> Signed-off-by: Dipika <dipikasikka1@gmail.com>
Dipika Sikka committed
Jun 30, 2025 at 00:05 UTC
6f2f53a82dc9703abb389761bf9931e3c9a5a75b
6 files changed
+295
-22
tests/quantization/test_compressed_tensors.py
+3
-3
@@ -17,7 +17,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso
17
CompressedTensorsW4A4Fp4, CompressedTensorsW4A16Fp4,
18
CompressedTensorsW4A16Sparse24, CompressedTensorsW8A8Fp8,
19
CompressedTensorsW8A8Int8, CompressedTensorsW8A16Fp8,
20
- CompressedTensorsWNA16)
20
+ CompressedTensorsWNA16, cutlass_fp4_supported)
21
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
22
sparse_cutlass_supported)
23
from vllm.platforms import current_platform
@@ -668,8 +668,8 @@ def test_compressed_tensors_nvfp4(vllm_runner, args):
668
assert isinstance(qkv_proj.quant_method,
669
CompressedTensorsLinearMethod)
670
if isinstance(qkv_proj.scheme, scheme) or isinstance(
671
- qkv_proj.scheme, CompressedTensorsW4A16Fp4
672
- ) and not CompressedTensorsW4A4Fp4.cutlass_fp4_supported():
671
+ qkv_proj.scheme,
672
+ CompressedTensorsW4A16Fp4) and not cutlass_fp4_supported():
673
assert True
674
else:
675
raise AssertionError("FP4 Scheme Mismatch")
vllm/model_executor/layers/fused_moe/layer.py
+3
-1
@@ -1246,6 +1246,7 @@ class FusedMoE(torch.nn.Module):
1246
param.materialize(final_shape, dtype=loaded_weight.dtype)
1247
1248
expert_data = param.data if full_load else param.data[expert_id]
1249
+
1250
# Case input scale: input_scale loading is only supported for fp8
1251
if "input_scale" in weight_name:
1252
# this is needed for compressed-tensors only
@@ -1273,6 +1274,7 @@ class FusedMoE(torch.nn.Module):
1274
tp_rank=self.tp_rank)
1275
return True if return_success else None
1276
1277
+ # TODO @dsikka: ModelOpt should follow the proper MoE loading pattern
1278
if "ModelOpt" in quant_method_name:
1279
if ('weight_scale_2' in weight_name
1280
or 'input_scale' in weight_name):
@@ -1289,7 +1291,7 @@ class FusedMoE(torch.nn.Module):
1291
tp_rank=self.tp_rank)
1292
return True if return_success else None
1293
1292
- # Case weight scales, zero_points and offset
1294
+ # Case weight scales, zero_points and offset, weight/input global scales
1295
if ("scale" in weight_name or "zero" in weight_name
1296
or "offset" in weight_name):
1297
# load the weight scales and zp based on the quantization scheme
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+3
-1
@@ -33,6 +33,8 @@ from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
33
find_matched_target, is_activation_quantization_format,
34
should_ignore_layer)
35
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
36
+from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( # noqa: E501
37
+ cutlass_fp4_supported)
38
from vllm.platforms import current_platform
39
40
logger = init_logger(__name__)
@@ -375,7 +377,7 @@ class CompressedTensorsConfig(QuantizationConfig):
377
378
if is_activation_quantization_format(self.quant_format):
379
if self._is_fp4a4_nvfp4(weight_quant, input_quant):
378
- if CompressedTensorsW4A4Fp4.cutlass_fp4_supported(
380
+ if cutlass_fp4_supported(
381
) or envs.VLLM_USE_NVFP4_CT_EMULATIONS:
382
return CompressedTensorsW4A4Fp4()
383
else:
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+271
-4
@@ -21,8 +21,12 @@ from vllm.model_executor.layers.quantization.utils import replace_parameter
21
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
22
check_moe_marlin_supports_layer, marlin_make_workspace_new,
23
marlin_moe_permute_scales)
24
+from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
25
+ prepare_moe_fp4_layer_for_marlin)
26
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
27
prepare_moe_fp8_layer_for_marlin)
28
+from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( # noqa: E501
29
+ cutlass_fp4_supported)
30
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
31
all_close_1d, normalize_e4m3fn_to_e4m3fnuz, per_tensor_dequantize)
32
from vllm.model_executor.utils import set_weight_attrs
@@ -46,12 +50,11 @@ class GPTQMarlinState(Enum):
50
51
52
__all__ = [
49
- "CompressedTensorsMoEMethod",
50
- "CompressedTensorsW8A8Fp8MoEMethod",
53
+ "CompressedTensorsMoEMethod", "CompressedTensorsW8A8Fp8MoEMethod",
54
"CompressedTensorsW8A8Fp8MoECutlassMethod",
55
"CompressedTensorsW8A8Int8MoEMethod",
53
- "CompressedTensorsWNA16MarlinMoEMethod",
54
- "CompressedTensorsWNA16MoEMethod",
56
+ "CompressedTensorsWNA16MarlinMoEMethod", "CompressedTensorsWNA16MoEMethod",
57
+ "CompressedTensorsW4A4MoeMethod"
58
]
59
60
@@ -84,6 +87,8 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
87
else:
88
logger.info_once("Using CompressedTensorsWNA16MarlinMoEMethod")
89
return CompressedTensorsWNA16MarlinMoEMethod(quant_config)
90
+ elif quant_config._is_fp4a4_nvfp4(weight_quant, input_quant):
91
+ return CompressedTensorsW4A4MoeMethod()
92
elif quant_config._is_fp8_w8a8_sm90(weight_quant, input_quant):
93
return CompressedTensorsW8A8Fp8MoECutlassMethod(quant_config)
94
elif quant_config._is_fp8_w8a8(weight_quant, input_quant):
@@ -95,6 +100,268 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
100
f"Unsupported FusedMoe scheme: {weight_quant}, {input_quant}")
101
102
103
+class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
104
+
105
+ def __init__(self):
106
+ self.use_marlin = not cutlass_fp4_supported()
107
+ self.group_size = 16
108
+
109
+ def create_weights(self, layer: torch.nn.Module, num_experts: int,
110
+ hidden_size: int, intermediate_size_per_partition: int,
111
+ params_dtype: torch.dtype, **extra_weight_attrs):
112
+
113
+ layer.num_experts = num_experts
114
+ layer.params_dtype = params_dtype
115
+
116
+ w13_weight = torch.nn.Parameter(
117
+ torch.empty(
118
+ num_experts,
119
+ 2 * intermediate_size_per_partition,
120
+ # 2 fp4 items are packed in the input dimension
121
+ hidden_size // 2,
122
+ requires_grad=False,
123
+ dtype=torch.uint8),
124
+ requires_grad=False)
125
+ layer.register_parameter("w13_weight_packed", w13_weight)
126
+ set_weight_attrs(w13_weight, extra_weight_attrs)
127
+
128
+ w2_weight = torch.nn.Parameter(
129
+ torch.empty(
130
+ num_experts,
131
+ hidden_size,
132
+ # 2 fp4 items are packed in the input dimension
133
+ intermediate_size_per_partition // 2,
134
+ dtype=torch.uint8),
135
+ requires_grad=False)
136
+ layer.register_parameter("w2_weight_packed", w2_weight)
137
+ set_weight_attrs(w2_weight, extra_weight_attrs)
138
+
139
+ # Weight Scales
140
+ w13_weight_scale = torch.nn.Parameter(
141
+ torch.empty(
142
+ num_experts,
143
+ 2 * intermediate_size_per_partition,
144
+ # 2 fp4 items are packed in the input dimension
145
+ hidden_size // self.group_size,
146
+ dtype=torch.float8_e4m3fn),
147
+ requires_grad=False)
148
+ layer.register_parameter("w13_weight_scale", w13_weight_scale)
149
+ extra_weight_attrs.update(
150
+ {"quant_method": FusedMoeWeightScaleSupported.GROUP.value})
151
+ set_weight_attrs(w13_weight_scale, extra_weight_attrs)
152
+
153
+ w2_weight_scale = torch.nn.Parameter(
154
+ torch.empty(
155
+ num_experts,
156
+ hidden_size,
157
+ # 2 fp4 items are packed in the input dimension
158
+ intermediate_size_per_partition // self.group_size,
159
+ dtype=torch.float8_e4m3fn),
160
+ requires_grad=False)
161
+ layer.register_parameter("w2_weight_scale", w2_weight_scale)
162
+ extra_weight_attrs.update(
163
+ {"quant_method": FusedMoeWeightScaleSupported.GROUP.value})
164
+ set_weight_attrs(w2_weight_scale, extra_weight_attrs)
165
+
166
+ # Weight Global Scales
167
+ w13_weight_scale_2 = torch.nn.Parameter(torch.empty(
168
+ num_experts, 2, dtype=torch.float32),
169
+ requires_grad=False)
170
+ layer.register_parameter("w13_weight_global_scale", w13_weight_scale_2)
171
+ extra_weight_attrs.update(
172
+ {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value})
173
+ set_weight_attrs(w13_weight_scale_2, extra_weight_attrs)
174
+
175
+ w2_weight_scale_2 = torch.nn.Parameter(torch.empty(
176
+ num_experts, dtype=torch.float32),
177
+ requires_grad=False)
178
+ layer.register_parameter("w2_weight_global_scale", w2_weight_scale_2)
179
+ extra_weight_attrs.update(
180
+ {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value})
181
+ set_weight_attrs(w2_weight_scale_2, extra_weight_attrs)
182
+
183
+ # Input Global Scales
184
+ w13_input_scale = torch.nn.Parameter(torch.empty(num_experts,
185
+ 2,
186
+ dtype=torch.float32),
187
+ requires_grad=False)
188
+ layer.register_parameter("w13_input_global_scale", w13_input_scale)
189
+ extra_weight_attrs.update(
190
+ {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value})
191
+ set_weight_attrs(w13_input_scale, extra_weight_attrs)
192
+
193
+ w2_input_scale = torch.nn.Parameter(torch.empty(num_experts,
194
+ dtype=torch.float32),
195
+ requires_grad=False)
196
+ layer.register_parameter("w2_input_global_scale", w2_input_scale)
197
+ extra_weight_attrs.update(
198
+ {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value})
199
+ set_weight_attrs(w2_input_scale, extra_weight_attrs)
200
+
201
+ def swizzle_blockscale(self, scale: torch.tensor):
202
+ assert (scale.dtype == torch.float8_e4m3fn)
203
+ # Pad and blockwise interleave weight_scale
204
+ scale_ndim = scale.ndim
205
+ if scale.ndim == 2:
206
+ scale = scale.unsqueeze(0)
207
+ assert scale.ndim == 3
208
+ B, M, K = scale.shape
209
+ round_up_multiple = lambda x, m: (x + m - 1) // m * m
210
+ M_padded = round_up_multiple(M, 128)
211
+ K_padded = round_up_multiple(K, 4)
212
+ padded_scale = torch.zeros((B, M_padded, K_padded), dtype=scale.dtype)
213
+ padded_scale[:B, :M, :K] = scale
214
+ batches, rows, cols = padded_scale.shape
215
+ assert rows % 128 == 0
216
+ assert cols % 4 == 0
217
+ padded_scale = padded_scale.reshape(batches, rows // 128, 4, 32,
218
+ cols // 4, 4)
219
+ swizzled_scale = padded_scale.permute((0, 1, 4, 3, 2, 5))
220
+ swizzled_scale = swizzled_scale.contiguous().cuda()
221
+ return (swizzled_scale.reshape(M, K)
222
+ if scale_ndim == 2 else swizzled_scale.reshape(B, M, K))
223
+
224
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
225
+
226
+ # From packed to weight
227
+ layer.w13_weight = torch.nn.Parameter(layer.w13_weight_packed.data,
228
+ requires_grad=False)
229
+
230
+ layer.w2_weight = torch.nn.Parameter(layer.w2_weight_packed.data,
231
+ requires_grad=False)
232
+
233
+ if not torch.allclose(layer.w13_weight_global_scale[:, 0],
234
+ layer.w13_weight_global_scale[:, 1]):
235
+ logger.warning_once(
236
+ "w1_weight_global_scale must match w3_weight_global_scale. "
237
+ "Accuracy may be affected.")
238
+
239
+ # Take inverse of global scale saved to disk
240
+ layer.w13_weight_scale_2 = torch.nn.Parameter(
241
+ 1 / layer.w13_weight_global_scale[:, 0], requires_grad=False)
242
+
243
+ layer.w2_weight_scale_2 = torch.nn.Parameter(
244
+ 1 / layer.w2_weight_global_scale.data, requires_grad=False)
245
+
246
+ if self.use_marlin:
247
+ prepare_moe_fp4_layer_for_marlin(layer)
248
+ return
249
+
250
+ # swizzle weight scales
251
+ layer.w13_blockscale_swizzled = torch.nn.Parameter(
252
+ self.swizzle_blockscale(layer.w13_weight_scale),
253
+ requires_grad=False)
254
+
255
+ layer.w2_blockscale_swizzled = torch.nn.Parameter(
256
+ self.swizzle_blockscale(layer.w2_weight_scale),
257
+ requires_grad=False)
258
+
259
+ # w13
260
+ w13_input_global_scale = layer.w13_input_global_scale.max(
261
+ dim=1).values.to(torch.float32)
262
+
263
+ layer.g1_alphas = torch.nn.Parameter(
264
+ ((1 / w13_input_global_scale) * layer.w13_weight_scale_2),
265
+ requires_grad=False)
266
+
267
+ layer.w13_input_scale_quant = torch.nn.Parameter(
268
+ (w13_input_global_scale), requires_grad=False)
269
+
270
+ # w2
271
+ layer.g2_alphas = torch.nn.Parameter(
272
+ ((1 / layer.w2_input_global_scale) * layer.w2_weight_scale_2).to(
273
+ torch.float32),
274
+ requires_grad=False)
275
+
276
+ layer.w2_input_scale_quant = torch.nn.Parameter(
277
+ (layer.w2_input_global_scale), requires_grad=False)
278
+
279
+ def apply(
280
+ self,
281
+ layer: torch.nn.Module,
282
+ x: torch.Tensor,
283
+ router_logits: torch.Tensor,
284
+ top_k: int,
285
+ renormalize: bool,
286
+ use_grouped_topk: bool = False,
287
+ topk_group: Optional[int] = None,
288
+ num_expert_group: Optional[int] = None,
289
+ global_num_experts: int = -1,
290
+ expert_map: Optional[torch.Tensor] = None,
291
+ custom_routing_function: Optional[Callable] = None,
292
+ scoring_func: str = "softmax",
293
+ e_score_correction_bias: Optional[torch.Tensor] = None,
294
+ apply_router_weight_on_input: bool = False,
295
+ activation: str = "silu",
296
+ enable_eplb: bool = False,
297
+ expert_load_view: Optional[torch.Tensor] = None,
298
+ logical_to_physical_map: Optional[torch.Tensor] = None,
299
+ logical_replica_count: Optional[torch.Tensor] = None,
300
+ ) -> torch.Tensor:
301
+ if enable_eplb:
302
+ raise NotImplementedError("EPLB not supported for "
303
+ "`CompressedTensorsW4A4MoeMethod` yet.")
304
+
305
+ topk_weights, topk_ids = FusedMoE.select_experts(
306
+ hidden_states=x,
307
+ router_logits=router_logits,
308
+ use_grouped_topk=use_grouped_topk,
309
+ top_k=top_k,
310
+ renormalize=renormalize,
311
+ topk_group=topk_group,
312
+ num_expert_group=num_expert_group,
313
+ custom_routing_function=custom_routing_function,
314
+ scoring_func=scoring_func,
315
+ e_score_correction_bias=e_score_correction_bias,
316
+ )
317
+
318
+ if self.use_marlin:
319
+ return torch.ops.vllm.fused_marlin_moe(
320
+ x,
321
+ layer.w13_weight,
322
+ layer.w2_weight,
323
+ layer.w13_weight_scale,
324
+ layer.w2_weight_scale,
325
+ router_logits,
326
+ topk_weights,
327
+ topk_ids,
328
+ global_scale1=layer.w13_weight_scale_2,
329
+ global_scale2=layer.w2_weight_scale_2,
330
+ quant_type_id=scalar_types.float4_e2m1f.id,
331
+ global_num_experts=global_num_experts,
332
+ expert_map=expert_map)
333
+
334
+ assert activation == "silu", "Only SiLU activation is supported."
335
+ assert not apply_router_weight_on_input, (
336
+ "Router weight on input is not "
337
+ "supported for CompressedTensorsW4A4MoeMethod.")
338
+ assert expert_map is None, ("Expert Parallelism / expert_map "
339
+ "is currently not supported for "
340
+ "CompressedTensorsW4A4MoeMethod.")
341
+
342
+ from vllm.model_executor.layers.fused_moe.cutlass_moe import (
343
+ cutlass_moe_fp4)
344
+
345
+ # Cutlass moe takes in activations in BF16/Half precision
346
+ # and fp4 quantized weights loaded from the checkpoint
347
+ return cutlass_moe_fp4(a=x,
348
+ w1_fp4=layer.w13_weight,
349
+ w1_blockscale=layer.w13_blockscale_swizzled,
350
+ w1_alphas=layer.g1_alphas,
351
+ w2_fp4=layer.w2_weight,
352
+ w2_blockscale=layer.w2_blockscale_swizzled,
353
+ w2_alphas=layer.g2_alphas,
354
+ topk_weights=topk_weights,
355
+ topk_ids=topk_ids,
356
+ m=x.shape[0],
357
+ n=layer.w2_weight.shape[2] * 2,
358
+ k=x.shape[1],
359
+ e=layer.w13_weight.shape[0],
360
+ a1_gscale=layer.w13_input_scale_quant,
361
+ a2_gscale=layer.w2_input_scale_quant,
362
+ device=x.device).to(x.dtype)
363
+
364
+
365
class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
366
367
def __init__(
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py
+1
-12
@@ -5,8 +5,7 @@ import torch
5
from torch.nn.parameter import Parameter
6
7
import vllm.envs as envs
8
-from vllm._custom_ops import (cutlass_scaled_fp4_mm,
9
- cutlass_scaled_mm_supports_fp4, scaled_fp4_quant)
8
+from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
9
from vllm.logger import init_logger
10
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
11
CompressedTensorsScheme)
@@ -15,7 +14,6 @@ from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import
14
from vllm.model_executor.parameter import (GroupQuantScaleParameter,
15
ModelWeightParameter,
16
PerTensorScaleParameter)
18
-from vllm.platforms import current_platform
17
18
logger = init_logger(__name__)
19
@@ -33,15 +31,6 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
31
return 80
32
return 100
33
36
- @classmethod
37
- def cutlass_fp4_supported(cls) -> bool:
38
- if not current_platform.is_cuda():
39
- return False
40
- capability_tuple = current_platform.get_device_capability()
41
- capability = -1 if capability_tuple is None else capability_tuple.to_int( # noqa: E501
42
- )
43
- return cutlass_scaled_mm_supports_fp4(capability)
44
-
34
def create_weights(self, layer: torch.nn.Module,
35
output_partition_sizes: list[int],
36
input_size_per_partition: int,
vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py
+14
-1
@@ -2,9 +2,14 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import torch
4
5
+from vllm._custom_ops import cutlass_scaled_mm_supports_fp4
6
+from vllm.platforms import current_platform
7
from vllm.scalar_type import scalar_types
8
7
-__all__ = ["break_fp4_bytes", "dequantize_to_dtype", "ref_nvfp4_quant"]
9
+__all__ = [
10
+ "break_fp4_bytes", "dequantize_to_dtype", "ref_nvfp4_quant",
11
+ "cutlass_fp4_supported"
12
+]
13
14
FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max()
15
@@ -12,6 +17,14 @@ kE2M1ToFloat = torch.tensor([0., 0.5, 1., 1.5, 2., 3., 4., 6.],
17
dtype=torch.float32)
18
19
20
+def cutlass_fp4_supported() -> bool:
21
+ if not current_platform.is_cuda():
22
+ return False
23
+ capability_tuple = current_platform.get_device_capability()
24
+ capability = -1 if capability_tuple is None else capability_tuple.to_int()
25
+ return cutlass_scaled_mm_supports_fp4(capability)
26
+
27
+
28
def break_fp4_bytes(a, dtype):
29
assert a.dtype == torch.uint8
30
m, n = a.shape