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, 11:00 UTC
[Bugfix] Fix FP8 Marlin MoE and enable for compressed-tensors models (#18026)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
May 13, 2025 at 22:48 UTC
9a2a6357de8aa112692dab93f7d40b2d7e75ac67
3 files changed
+53
-6
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+48
-6
@@ -9,6 +9,7 @@ from compressed_tensors import CompressionFormat
9
from compressed_tensors.quantization import (ActivationOrdering,
10
QuantizationStrategy)
11
12
+import vllm.envs as envs
13
import vllm.model_executor.layers.fused_moe # noqa
14
from vllm import _custom_ops as ops
15
from vllm.logger import init_logger
@@ -20,10 +21,13 @@ 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_fp8 import (
25
+ prepare_moe_fp8_layer_for_marlin)
26
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
27
all_close_1d, normalize_e4m3fn_to_e4m3fnuz, per_tensor_dequantize)
28
from vllm.model_executor.utils import set_weight_attrs
29
from vllm.platforms import current_platform
30
+from vllm.scalar_type import scalar_types
31
32
logger = init_logger(__name__)
33
@@ -114,10 +118,24 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
118
"For FP8 Fused MoE layer, we require either per tensor or "
119
"channelwise, dynamic per token quantization.")
120
121
+ # For GPUs that lack FP8 hardware support, we can leverage the Marlin
122
+ # kernel for fast weight-only FP8 quantization
123
+ self.use_marlin = (not current_platform.has_device_capability(89)
124
+ or envs.VLLM_TEST_FORCE_FP8_MARLIN)
125
+ # Disable marlin for rocm
126
+ if current_platform.is_rocm():
127
+ self.use_marlin = False
128
+
129
def create_weights(self, layer: torch.nn.Module, num_experts: int,
130
hidden_size: int, intermediate_size_per_partition: int,
131
params_dtype: torch.dtype, **extra_weight_attrs):
132
133
+ layer.intermediate_size_per_partition = intermediate_size_per_partition
134
+ layer.hidden_size = hidden_size
135
+ layer.num_experts = num_experts
136
+ layer.orig_dtype = params_dtype
137
+ layer.weight_block_size = None
138
+
139
params_dtype = torch.float8_e4m3fn
140
141
# WEIGHTS
@@ -280,6 +298,12 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
298
from vllm.model_executor.layers.fused_moe import fused_experts
299
self.fused_experts_func = fused_experts
300
301
+ if self.use_marlin:
302
+ prepare_moe_fp8_layer_for_marlin(layer, False)
303
+ # Activations not quantized for marlin.
304
+ del layer.w13_input_scale
305
+ del layer.w2_input_scale
306
+
307
def apply(
308
self,
309
layer: torch.nn.Module,
@@ -311,6 +335,24 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
335
scoring_func=scoring_func,
336
e_score_correction_bias=e_score_correction_bias)
337
338
+ if self.use_marlin:
339
+ assert activation == "silu", (
340
+ f"{activation} not supported for Marlin MoE.")
341
+ assert not apply_router_weight_on_input, (
342
+ "Apply router weight on input not supported for Marlin MoE.")
343
+ return torch.ops.vllm.fused_marlin_moe(
344
+ x,
345
+ layer.w13_weight,
346
+ layer.w2_weight,
347
+ layer.w13_weight_scale,
348
+ layer.w2_weight_scale,
349
+ router_logits,
350
+ topk_weights,
351
+ topk_ids,
352
+ quant_type_id=scalar_types.float8_e4m3fn.id,
353
+ global_num_experts=global_num_experts,
354
+ expert_map=expert_map)
355
+
356
return self.fused_experts_func(
357
hidden_states=x,
358
w1=layer.w13_weight,
@@ -517,7 +559,8 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
559
activation: str = "silu",
560
) -> torch.Tensor:
561
520
- assert activation == "silu"
562
+ assert activation == "silu", (
563
+ f"{activation} not supported for Cutlass MoE.")
564
565
topk_weights, topk_ids = FusedMoE.select_experts(
566
hidden_states=x,
@@ -942,11 +985,10 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
985
apply_router_weight_on_input: bool = False,
986
activation: str = "silu",
987
) -> torch.Tensor:
945
- assert activation == "silu", "Only SiLU activation is supported."
946
- if apply_router_weight_on_input:
947
- raise NotImplementedError(
948
- "Apply router weight on input is not supported for "
949
- "fused Marlin MoE method.")
988
+ assert activation == "silu", (
989
+ f"{activation} not supported for Marlin MoE.")
990
+ assert not apply_router_weight_on_input, (
991
+ "Apply router weight on input not supported for Marlin MoE.")
992
993
topk_weights, topk_ids = FusedMoE.select_experts(
994
hidden_states=x,
vllm/model_executor/layers/quantization/fp8.py
+4
@@ -811,6 +811,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):
811
)
812
813
if self.use_marlin:
814
+ assert activation == "silu", (
815
+ f"{activation} not supported for Marlin MoE.")
816
+ assert not apply_router_weight_on_input, (
817
+ "Apply router weight on input not supported for Marlin MoE.")
818
return torch.ops.vllm.fused_marlin_moe(
819
x,
820
layer.w13_weight,
vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py
+1
@@ -268,6 +268,7 @@ def prepare_moe_fp8_layer_for_marlin(layer: torch.nn.Module,
268
tensor_list.append(marlin_scales)
269
270
scales = torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
271
+ scales = fp8_fused_exponent_bias_into_scales(scales)
272
scales = torch.nn.Parameter(scales, requires_grad=False)
273
274
setattr(layer, name + "_weight_scale", scales)