@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 268b1c55ad

Fork 0

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, 19:00 UTC

[MoE Refactor][13/N] Convert FI to Use PFNoEP (#31533)

Signed-off-by: Robert Shaw <robshaw@redhat.com> Signed-off-by: Robert Shaw <robertgshaw2@gmail.com> Signed-off-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>

Robert Shaw committed Jan 3, 2026 at 15:26 UTC 268b1c55ade396f9b4b7c811d66725a4c9b74590

7 files changed +78 -168

vllm/model_executor/layers/fused_moe/all2all_utils.py

+13 -3

@@ -15,6 +15,9 @@ 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 +)

21 from vllm.platforms import current_platform

22 from vllm.utils.import_utils import has_deep_ep, has_pplx

@@ -77,10 +80,17 @@ def maybe_make_prepare_finalize(

81 prepare_finalize: FusedMoEPrepareAndFinalize | None = None

80 - # TODO: could allow this now

81 - assert not moe.use_flashinfer_cutlass_kernels, "Must be created in modelopt.py"

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 + )

83 - if moe.use_pplx_kernels:

93 + elif moe.use_pplx_kernels:

94 assert quant_config is not None

96 hidden_dim_bytes, hidden_scale_bytes = pplx_hidden_dim_scale_bytes(

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py

+17 -5

@@ -10,6 +10,9 @@ from vllm.distributed.device_communicators.base_device_communicator import (

10 )

11 from vllm.forward_context import get_forward_context

12 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

13 +from vllm.model_executor.layers.fused_moe.prepare_finalize import (

14 + MoEPrepareAndFinalizeNoEP,

15 +)

16 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

17 TopKWeightAndReduceNoOP,

18 )

@@ -349,14 +352,23 @@ def create_flashinfer_prepare_finalize(

352 use_nvfp4: bool = False,

353 enable_alltoallv: bool = False,

354 use_deepseek_fp8_block_scale: bool = False,

352 -) -> FlashInferCutlassMoEPrepareAndFinalize:

355 +) -> FlashInferCutlassMoEPrepareAndFinalize | MoEPrepareAndFinalizeNoEP:

356 """Factory function to create the appropriate FlashInfer implementation."""

357 + # TODO(rob): migrate non-DP cases to MoEPrepareAndFinalizeNoEP

358 + # once we complete the FP8 refactor.

359 if use_nvfp4:

360 if enable_alltoallv:

361 return FlashInferAllToAllMoEPrepareAndFinalize(use_dp)

362 else:

363 return FlashInferAllGatherMoEPrepareAndFinalize(use_dp)

359 - # FP8 path currently supported via AllGather; optionally enable block-scale

360 - return FlashInferAllGatherMoEPrepareAndFinalize(

361 - use_dp=use_dp, use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale

362 - )

364 +

365 + # FP8 DP path currently supported via AllGather.

366 + if use_dp:

367 + return FlashInferAllGatherMoEPrepareAndFinalize(

368 + use_dp=True,

369 + use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale,

370 + )

371 + else:

372 + # NOTE(rob): CUTLASS FP8 block quant executes the input

373 + # quantzation and grouped gemm in a single kernel.

374 + return MoEPrepareAndFinalizeNoEP(defer_input_quant=use_deepseek_fp8_block_scale)

vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py

-1

@@ -49,7 +49,6 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):

49 prepare_finalize,

50 old_quant_method.select_gemm_impl(prepare_finalize, moe_layer),

51 shared_experts,

52 - getattr(moe_layer, "shared_experts_stream", None),

52 moe_parallel_config=moe_layer.moe_parallel_config,

53 ),

54 )

vllm/model_executor/layers/fused_moe/layer.py

+2 -2

@@ -356,14 +356,14 @@ class FusedMoE(CustomOp):

356 # TODO: Remove this after more extensive testings with TP/DP

357 # and other execution modes

358 if envs.VLLM_DISABLE_SHARED_EXPERTS_STREAM:

359 - logger.info_once("Disabling MoE shared_experts cuda stream")

359 + logger.debug_once("Disabling MoE shared_experts cuda stream", scope="local")

360 self.shared_experts_stream = None

361 else:

362 # TODO(rob): enable shared expert overlap with non-cuda-alike.

363 # aux_stream() returns None on non-cuda-alike platforms.

364 self.shared_experts_stream = aux_stream()

365 if self.shared_experts_stream is not None:

366 - logger.info_once(

366 + logger.debug_once(

367 "Enabled separate cuda stream for MoE shared_experts", scope="local"

368 )

369

vllm/model_executor/layers/fused_moe/modular_kernel.py

+5 -75

@@ -21,7 +21,6 @@ from vllm.model_executor.layers.fused_moe.utils import (

21 count_expert_num_tokens,

22 disable_inplace,

23 )

24 -from vllm.platforms import current_platform

24 from vllm.utils.math_utils import cdiv

25 from vllm.v1.worker.ubatching import (

26 dbo_enabled,

@@ -682,14 +681,12 @@ class FusedMoEModularKernel(torch.nn.Module):

681 prepare_finalize: FusedMoEPrepareAndFinalize,

682 fused_experts: FusedMoEPermuteExpertsUnpermute,

683 shared_experts: torch.nn.Module | None = None,

685 - shared_experts_stream: torch.cuda.Stream | None = None,

684 moe_parallel_config: FusedMoEParallelConfig | None = None,

685 ):

686 super().__init__()

687 self.prepare_finalize = prepare_finalize

688 self.fused_experts = fused_experts

689 self.shared_experts = shared_experts

692 - self.shared_experts_stream = shared_experts_stream

690

691 # prefer an explicit FusedMoEParallelConfig when available (from

692 # FusedMoE layers / tests).

@@ -904,34 +901,6 @@ class FusedMoEModularKernel(torch.nn.Module):

901 expert_num_tokens_cpu=c_expert_num_tokens_cpu,

902 )

903

907 - def _maybe_setup_shared_experts_stream(

908 - self, hidden_states: torch.Tensor

909 - ) -> tuple[bool, torch.Tensor | None]:

910 - # decide whether to run shared experts on a separate CUDA stream to

911 - # overlap with the main fused MoE kernel.

912 - use_shared_experts_stream = (

913 - self.shared_experts is not None

914 - and self.shared_experts_stream is not None

915 - and hidden_states.is_cuda

916 - and (

917 - hidden_states.shape[0]

918 - <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD

919 - )

920 - )

921 -

922 - hidden_states_clone: torch.Tensor | None = None

923 - if use_shared_experts_stream and self.shared_experts_stream is not None:

924 - # TODO: Optimize this (complicated)

925 - # Note: this clone adds overhead but is required

926 - # for correctness with multiple CUDA streams and CUDA graph capture.

927 - hidden_states_clone = hidden_states.clone()

928 - # record that the clone will be used by the separate stream so its

929 - # lifetime is correctly tracked.

930 - hidden_states_clone.record_stream(self.shared_experts_stream)

931 - self.shared_experts_stream.wait_stream(torch.cuda.current_stream())

932 -

933 - return use_shared_experts_stream, hidden_states_clone

934 -

904 def _prepare(

905 self,

906 hidden_states: torch.Tensor,

@@ -1119,30 +1088,12 @@ class FusedMoEModularKernel(torch.nn.Module):

1088 topk_weights: torch.Tensor,

1089 topk_ids: torch.Tensor,

1090 apply_router_weight_on_input: bool,

1122 - hidden_states_clone: torch.Tensor | None = None,

1123 - use_shared_experts_stream: bool = False,

1091 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

1092 """

1093 The _finalize method is a wrapper around self.prepare_finalize.finalize

1094 that handles DBO, async and shared expert overlap.

1095 """

1129 -

1130 - def maybe_run_shared_experts() -> torch.Tensor | None:

1131 - if self.shared_experts is None:

1132 - return None

1133 -

1134 - if (

1135 - not use_shared_experts_stream

1136 - or self.shared_experts_stream is not None

1137 - and (not hidden_states.is_cuda or not torch.cuda.is_available())

1138 - ):

1139 - # fall back to running on the current stream

1140 - return self.shared_experts(hidden_states)

1141 -

1142 - assert hidden_states_clone is not None

1143 - # launch shared experts on the dedicated stream.

1144 - with torch.cuda.stream(self.shared_experts_stream):

1145 - return self.shared_experts(hidden_states_clone)

1096 + shared_output: torch.Tensor | None = None

1097

1098 if not self.prepare_finalize.supports_async():

1099 assert not dbo_enabled()

@@ -1155,7 +1106,8 @@ class FusedMoEModularKernel(torch.nn.Module):

1106 apply_router_weight_on_input,

1107 self.fused_experts.finalize_weight_and_reduce_impl(),

1108 )

1158 - shared_output = maybe_run_shared_experts()

1109 + if self.shared_experts is not None:

1110 + shared_output = self.shared_experts(hidden_states)

1111 else:

1112 finalize_ret = self.prepare_finalize.finalize_async(

1113 output,

@@ -1165,8 +1117,8 @@ class FusedMoEModularKernel(torch.nn.Module):

1117 apply_router_weight_on_input,

1118 self.fused_experts.finalize_weight_and_reduce_impl(),

1119 )

1168 -

1169 - shared_output = maybe_run_shared_experts()

1120 + if self.shared_experts is not None:

1121 + shared_output = self.shared_experts(hidden_states)

1122

1123 # TODO(lucas): refactor this in the alternative schedules followup

1124 # currently unpack if we have hook + receiver pair or just

@@ -1189,28 +1141,12 @@ class FusedMoEModularKernel(torch.nn.Module):

1141

1142 receiver()

1143

1192 - self._wait_for_shared_experts_stream(hidden_states, use_shared_experts_stream)

1193 -

1144 if self.shared_experts is None:

1145 return output

1146 else:

1147 assert shared_output is not None

1148 return shared_output, output

1149

1200 - def _wait_for_shared_experts_stream(

1201 - self, hidden_states: torch.Tensor, use_shared_experts_stream: bool

1202 - ) -> None:

1203 - # ensure that any work enqueued on the shared_experts_stream is

1204 - # completed before the shared_output tensor is consumed

1205 - if (

1206 - self.shared_experts is not None

1207 - and use_shared_experts_stream

1208 - and self.shared_experts_stream is not None

1209 - and hidden_states.is_cuda

1210 - and current_platform.is_cuda()

1211 - ):

1212 - torch.cuda.current_stream().wait_stream(self.shared_experts_stream)

1213 -

1150 def forward(

1151 self,

1152 hidden_states: torch.Tensor,

@@ -1257,10 +1193,6 @@ class FusedMoEModularKernel(torch.nn.Module):

1193 else:

1194 output = torch.zeros_like(hidden_states)

1195

1260 - use_shared_experts_stream, hidden_states_clone = (

1261 - self._maybe_setup_shared_experts_stream(hidden_states)

1262 - )

1263 -

1196 local_num_experts = w1.size(0)

1197 if global_num_experts == -1:

1198 global_num_experts = local_num_experts

@@ -1297,6 +1229,4 @@ class FusedMoEModularKernel(torch.nn.Module):

1229 topk_weights,

1230 topk_ids,

1231 apply_router_weight_on_input,

1300 - hidden_states_clone=hidden_states_clone,

1301 - use_shared_experts_stream=use_shared_experts_stream,

1232 )

vllm/model_executor/layers/quantization/fp8.py

+40 -73

@@ -48,7 +48,6 @@ 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,

51 get_flashinfer_moe_backend,

52 register_moe_scaling_factors,

53 rotate_flashinfer_fp8_moe_weights,

@@ -973,27 +972,33 @@ class Fp8MoEMethod(FusedMoEMethodBase):

972 # done, then we will initialzie the TP case and DP/EP case

973 # via the same code path (i.e. via maybe_init_modular_kernel).

974 # NOTE(rob): in progress migrating all into this format.

976 - if self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:

977 - from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (

978 - FlashInferExperts,

979 - )

980 - from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize import ( # noqa: E501

981 - FlashInferAllGatherMoEPrepareAndFinalize,

982 - )

975

984 - config = self.get_fused_moe_quant_config(layer)

985 - assert config is not None

986 - self.moe_quant_config = config

976 + from vllm.model_executor.layers.fused_moe import (

977 + TritonOrDeepGemmExperts,

978 + )

979 + from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (

980 + FlashInferExperts,

981 + )

982 + from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (

983 + MarlinExperts,

984 + )

985 + from vllm.model_executor.layers.fused_moe.prepare_finalize import (

986 + MoEPrepareAndFinalizeNoEP,

987 + )

988 + from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (

989 + AiterExperts,

990 + )

991 +

992 + self.moe_quant_config = self.get_fused_moe_quant_config(layer)

993 + assert self.moe_quant_config is not None

994 + self.use_inplace = True

995

996 + if self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:

997 self.kernel = mk.FusedMoEModularKernel(

989 - # TODO(rob): we can use the generic MoEPrepareAndFinalizeNoEP

990 - # with the changes to defer input quantization

991 - FlashInferAllGatherMoEPrepareAndFinalize(

992 - use_dp=(self.moe.dp_size > 1),

993 - use_deepseek_fp8_block_scale=self.block_quant,

994 - ),

998 + # TODO: make defer_input_quant an attr of the FlashInferExperts

999 + MoEPrepareAndFinalizeNoEP(defer_input_quant=self.block_quant),

1000 FlashInferExperts(

996 - out_dtype=torch.get_default_dtype(),

1001 + out_dtype=layer.orig_dtype,

1002 quant_config=self.moe_quant_config,

1003 ep_rank=self.moe.ep_rank,

1004 ep_size=self.moe.ep_size,

@@ -1005,50 +1010,26 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1010 )

1011 self.use_inplace = False

1012

1008 - elif self.fp8_backend in [

1009 - Fp8MoeBackend.DEEPGEMM,

1010 - Fp8MoeBackend.TRITON,

1011 - Fp8MoeBackend.MARLIN,

1012 - Fp8MoeBackend.AITER,

1013 - ]:

1014 - from vllm.model_executor.layers.fused_moe import (

1015 - TritonOrDeepGemmExperts,

1016 - )

1017 - from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (

1018 - MarlinExperts,

1013 + elif self.fp8_backend == Fp8MoeBackend.AITER:

1014 + self.kernel = mk.FusedMoEModularKernel(

1015 + # TODO: make defer_input_quant an attr of the AiterExperts

1016 + MoEPrepareAndFinalizeNoEP(defer_input_quant=True),

1017 + AiterExperts(quant_config=self.moe_quant_config),

1018 )

1020 - from vllm.model_executor.layers.fused_moe.prepare_finalize import (

1021 - MoEPrepareAndFinalizeNoEP,

1019 + elif self.fp8_backend == Fp8MoeBackend.MARLIN:

1020 + self.kernel = mk.FusedMoEModularKernel(

1021 + MoEPrepareAndFinalizeNoEP(),

1022 + MarlinExperts(quant_config=self.moe_quant_config),

1023 )

1023 - from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (

1024 - AiterExperts,

1024 + else:

1025 + self.kernel = mk.FusedMoEModularKernel(

1026 + MoEPrepareAndFinalizeNoEP(),

1027 + TritonOrDeepGemmExperts(

1028 + quant_config=self.moe_quant_config,

1029 + allow_deep_gemm=(self.fp8_backend == Fp8MoeBackend.DEEPGEMM),

1030 + ),

1031 )

1032

1027 - config = self.get_fused_moe_quant_config(layer)

1028 - assert config is not None

1029 - self.moe_quant_config = config

1030 -

1031 - if self.fp8_backend == Fp8MoeBackend.AITER:

1032 - self.kernel = mk.FusedMoEModularKernel(

1033 - # TODO: make defer_input_quant an attr of the AiterExperts

1034 - MoEPrepareAndFinalizeNoEP(defer_input_quant=True),

1035 - AiterExperts(quant_config=self.moe_quant_config),

1036 - )

1037 - elif self.fp8_backend == Fp8MoeBackend.MARLIN:

1038 - self.kernel = mk.FusedMoEModularKernel(

1039 - MoEPrepareAndFinalizeNoEP(),

1040 - MarlinExperts(quant_config=self.moe_quant_config),

1041 - )

1042 - else:

1043 - self.kernel = mk.FusedMoEModularKernel(

1044 - MoEPrepareAndFinalizeNoEP(),

1045 - TritonOrDeepGemmExperts(

1046 - quant_config=self.moe_quant_config,

1047 - allow_deep_gemm=(self.fp8_backend == Fp8MoeBackend.DEEPGEMM),

1048 - ),

1049 - )

1050 - self.use_inplace = True

1051 -

1033 def process_weights_after_loading(self, layer: Module) -> None:

1034 if getattr(layer, "_already_called_process_weights_after_loading", False):

1035 return

@@ -1121,21 +1102,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1102 or self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

1103 ):

1104 return None

1124 - elif self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:

1125 - if self.block_quant:

1126 - assert self.weight_block_size == [128, 128], (

1127 - f"Only support weight_block_size == [128, 128], "

1128 - f"got {self.weight_block_size}"

1129 - )

1130 - # Wire block-scale flag through prepare/finalize when using CUTLASS

1131 - prepare_finalize = build_flashinfer_fp8_cutlass_moe_prepare_finalize(

1132 - self.moe,

1133 - use_deepseek_fp8_block_scale=self.block_quant,

1134 - )

1135 - logger.debug_once("%s", prepare_finalize.__class__.__name__)

1136 - return prepare_finalize

1137 - else:

1138 - return super().maybe_make_prepare_finalize(routing_tables)

1105 + return super().maybe_make_prepare_finalize(routing_tables)

1106

1107 def select_gemm_impl(

1108 self,

vllm/model_executor/layers/quantization/modelopt.py

+1 -9

@@ -46,7 +46,6 @@ 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,

49 flashinfer_cutlass_moe_fp8,

50 get_flashinfer_moe_backend,

51 is_flashinfer_supporting_global_sf,

@@ -751,14 +750,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

750 # TRT LLM not supported with all2all yet.

751 if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:

752 return None

754 - elif self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:

755 - prepare_finalize = build_flashinfer_fp8_cutlass_moe_prepare_finalize(

756 - self.moe

757 - )

758 - logger.debug_once("%s", prepare_finalize.__class__.__name__)

759 - return prepare_finalize

760 - else:

761 - return super().maybe_make_prepare_finalize(routing_tables)

753 + return super().maybe_make_prepare_finalize(routing_tables)

754

755 def select_gemm_impl(

756 self,