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, 17:45 UTC
[Bugfix] Fix Broken ModelOpt NVFP4 MoE (#31742)
Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com>
Robert Shaw committed
Jan 5, 2026 at 18:18 UTC
f6c0009afa366e4238cb118882415cd6c448c16b
4 files changed
+32
-15
vllm/model_executor/layers/fused_moe/all2all_utils.py
+5
-13
@@ -15,9 +15,6 @@ from vllm.model_executor.layers.fused_moe.config import (
15
from vllm.model_executor.layers.fused_moe.modular_kernel import (
16
FusedMoEPrepareAndFinalize,
17
)
18
-from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
19
- build_flashinfer_fp8_cutlass_moe_prepare_finalize,
20
-)
18
from vllm.platforms import current_platform
19
from vllm.utils.import_utils import has_deep_ep, has_pplx
20
@@ -80,17 +77,12 @@ def maybe_make_prepare_finalize(
77
78
prepare_finalize: FusedMoEPrepareAndFinalize | None = None
79
83
- if moe.use_flashinfer_cutlass_kernels:
84
- assert quant_config is not None
85
- use_deepseek_fp8_block_scale = (
86
- quant_config is not None and quant_config.is_block_quantized
87
- )
88
- prepare_finalize = build_flashinfer_fp8_cutlass_moe_prepare_finalize(
89
- moe=moe,
90
- use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale,
91
- )
80
+ # TODO(rob): update this as part of the MoE refactor.
81
+ assert not moe.use_flashinfer_cutlass_kernels, (
82
+ "Must be created in modelopt.py or fp8.py"
83
+ )
84
93
- elif moe.use_pplx_kernels:
85
+ if moe.use_pplx_kernels:
86
assert quant_config is not None
87
88
hidden_dim_bytes, hidden_scale_bytes = pplx_hidden_dim_scale_bytes(
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
+3
-1
@@ -241,7 +241,9 @@ def flashinfer_cutlass_moe_fp4(
241
apply_router_weight_on_input: bool = False,
242
) -> torch.Tensor:
243
fused_experts = mk.FusedMoEModularKernel(
244
- create_flashinfer_prepare_finalize(use_dp=False),
244
+ create_flashinfer_prepare_finalize(
245
+ use_dp=False, use_nvfp4=True, enable_alltoallv=False
246
+ ),
247
FlashInferExperts(
248
out_dtype=hidden_states.dtype,
249
quant_config=quant_config,
vllm/model_executor/layers/quantization/fp8.py
+9
-1
@@ -48,6 +48,7 @@ from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
48
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
49
FlashinferMoeBackend,
50
apply_flashinfer_per_tensor_scale_fp8,
51
+ build_flashinfer_fp8_cutlass_moe_prepare_finalize,
52
get_flashinfer_moe_backend,
53
register_moe_scaling_factors,
54
rotate_flashinfer_fp8_moe_weights,
@@ -149,7 +150,7 @@ def get_fp8_moe_backend(
150
if block_quant and current_platform.is_device_capability_family(100):
151
raise ValueError(
152
"FlashInfer FP8 MoE throughput backend does not "
152
- "support block quantization. Please use "
153
+ "support block quantization on SM100. Please use "
154
"VLLM_FLASHINFER_MOE_BACKEND=latency "
155
"instead."
156
)
@@ -1102,6 +1103,13 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1103
or self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM
1104
):
1105
return None
1106
+ elif self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:
1107
+ prepare_finalize = build_flashinfer_fp8_cutlass_moe_prepare_finalize(
1108
+ self.moe,
1109
+ use_deepseek_fp8_block_scale=self.block_quant,
1110
+ )
1111
+ logger.debug_once("%s", prepare_finalize.__class__.__name__)
1112
+ return prepare_finalize
1113
return super().maybe_make_prepare_finalize(routing_tables)
1114
1115
def select_gemm_impl(
vllm/model_executor/layers/quantization/modelopt.py
+15
@@ -46,6 +46,7 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
46
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
47
FlashinferMoeBackend,
48
apply_flashinfer_per_tensor_scale_fp8,
49
+ build_flashinfer_fp8_cutlass_moe_prepare_finalize,
50
flashinfer_cutlass_moe_fp8,
51
get_flashinfer_moe_backend,
52
is_flashinfer_supporting_global_sf,
@@ -750,6 +751,17 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
751
# TRT LLM not supported with all2all yet.
752
if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:
753
return None
754
+ elif self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:
755
+ # TP case: avoid convert to ModularKernelMethod - to be refactored.
756
+ if self.moe.dp_size == 1:
757
+ return None
758
+
759
+ prepare_finalize = build_flashinfer_fp8_cutlass_moe_prepare_finalize(
760
+ self.moe,
761
+ use_deepseek_fp8_block_scale=False,
762
+ )
763
+ logger.debug_once("%s", prepare_finalize.__class__.__name__)
764
+ return prepare_finalize
765
return super().maybe_make_prepare_finalize(routing_tables)
766
767
def select_gemm_impl(
@@ -1444,6 +1456,9 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1456
self.allow_flashinfer
1457
and self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS
1458
):
1459
+ # TP case: avoid convert to ModularKernelMethod - to be refactored.
1460
+ if self.moe.dp_size == 1:
1461
+ return None
1462
# For now, fp4 moe only works with the flashinfer dispatcher.
1463
prepare_finalize = build_flashinfer_fp4_cutlass_moe_prepare_finalize(
1464
self.moe