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 14, 05:30 UTC
[Bugfix] Fix Marlin NVFP4 for modelopt (#23659)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Aug 27, 2025 at 17:48 UTC
f9ca2b40a0357d98e3fb8bd951745dfaceab459e
1 file changed
+12
-13
vllm/model_executor/layers/quantization/modelopt.py
+12
-13
@@ -891,7 +891,11 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
891
assert (layer.weight_scale.dtype == torch.float8_e4m3fn), (
892
"Weight Block scale must be represented as FP8-E4M3")
893
894
- if self.backend == "flashinfer-trtllm":
894
+ if self.backend == "marlin":
895
+ prepare_fp4_layer_for_marlin(layer)
896
+ del layer.alpha
897
+ del layer.input_scale
898
+ elif self.backend == "flashinfer-trtllm":
899
# FlashInfer TRTLLM FP4 GEMM requires a different weight layout.
900
# FlashInfer provides nvfp4_quantize to quantize + shuffle the
901
# layout but we use our own quantization so we have to call
@@ -916,11 +920,6 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
920
requires_grad=False)
921
layer.weight = Parameter(layer.weight.data, requires_grad=False)
922
919
- if self.backend == "marlin":
920
- prepare_fp4_layer_for_marlin(layer)
921
- del layer.alpha
922
- del layer.input_scale
923
-
923
def apply(
924
self,
925
layer: torch.nn.Module,
@@ -1312,6 +1311,13 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1311
del layer.w2_weight_scale
1312
del layer.w13_weight
1313
del layer.w13_weight_scale
1314
+ elif self.use_marlin:
1315
+ # Marlin processing
1316
+ prepare_moe_fp4_layer_for_marlin(layer)
1317
+ del layer.g1_alphas
1318
+ del layer.g2_alphas
1319
+ del layer.w13_input_scale_quant
1320
+ del layer.w2_input_scale_quant
1321
else:
1322
# Non-TRT-LLM processing (Cutlass or non-flashinfer)
1323
assert (layer.w13_weight_scale.shape[2] % 16 == 0), (
@@ -1333,13 +1339,6 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1339
layer.w2_weight = Parameter(layer.w2_weight.data,
1340
requires_grad=False)
1341
1336
- if self.use_marlin:
1337
- prepare_moe_fp4_layer_for_marlin(layer)
1338
- del layer.g1_alphas
1339
- del layer.g2_alphas
1340
- del layer.w13_input_scale_quant
1341
- del layer.w2_input_scale_quant
1342
-
1342
def apply(
1343
self,
1344
layer: torch.nn.Module,