@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 5a93b9162b

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 9, 14:45 UTC

[MoE Refactor] Integrate Naive Prepare Finalize into MK (#32567)

Signed-off-by: Robert Shaw <robshaw@redhat.com> Signed-off-by: Amir Klein <203507526+amirkl94@users.noreply.github.com> Co-authored-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: amirkl94 <203507526+amirkl94@users.noreply.github.com>

Robert Shaw committed Jan 26, 2026 at 20:28 UTC 5a93b9162bfef6be6714ba70d3191e1e7ab7d3e4

46 files changed +1017 -875

.buildkite/test-amd.yaml

+1 -1

@@ -1131,7 +1131,7 @@ steps:

1131 - csrc/quantization/cutlass_w8a8/moe/

1132 - vllm/model_executor/layers/fused_moe/cutlass_moe.py

1133 - vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py

1134 - - vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py

1134 + - vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py

1135 - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py

1136 - vllm/v1/attention/backends/flashinfer.py

1137 - vllm/v1/attention/backends/mla/cutlass_mla.py

.buildkite/test-pipeline.yaml

+1 -1

@@ -1017,7 +1017,7 @@ steps:

1017 - csrc/quantization/cutlass_w8a8/moe/

1018 - vllm/model_executor/layers/fused_moe/cutlass_moe.py

1019 - vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py

1020 - - vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py

1020 + - vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py

1021 - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py

1022 - vllm/v1/attention/backends/flashinfer.py

1023 - vllm/v1/attention/backends/mla/cutlass_mla.py

.buildkite/test_areas/kernels.yaml

+1 -1

@@ -85,7 +85,7 @@ steps:

85 - csrc/quantization/cutlass_w8a8/moe/

86 - vllm/model_executor/layers/fused_moe/cutlass_moe.py

87 - vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py

88 - - vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py

88 + - vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py

89 - vllm/model_executor/layers/quantization/utils/flashinfer_utils.py

90 - vllm/v1/attention/backends/flashinfer.py

91 - vllm/v1/attention/backends/mla/cutlass_mla.py

benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py

+2 -2

@@ -197,7 +197,7 @@ def bench_run(

197 )

198

199 kernel = mk.FusedMoEModularKernel(

200 - MoEPrepareAndFinalizeNoEP(defer_input_quant=True),

200 + MoEPrepareAndFinalizeNoEP(),

201 CutlassExpertsFp4(

202 make_dummy_moe_config(),

203 quant_config=quant_config,

@@ -242,7 +242,7 @@ def bench_run(

242 )

243

244 kernel = mk.FusedMoEModularKernel(

245 - MoEPrepareAndFinalizeNoEP(defer_input_quant=True),

245 + MoEPrepareAndFinalizeNoEP(),

246 CutlassExpertsFp4(

247 make_dummy_moe_config(),

248 quant_config=quant_config,

docs/design/moe_kernel_features.md

+1 -2

@@ -36,8 +36,7 @@ th {

36 | pplx | batched | fp8,int8 | G,A,T | Y | Y | [`PplxPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.pplx_prepare_finalize.PplxPrepareAndFinalize] |

37 | deepep_high_throughput | standard | fp8 | G(128),A,T2 | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize.DeepEPLLPrepareAndFinalize] |

38 | deepep_low_latency | batched | fp8 | G(128),A,T3 | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] |

39 -| flashinfer_all2allv | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferAllToAllMoEPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize.FlashInferAllToAllMoEPrepareAndFinalize] |

40 -| flashinfer4 | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferCutlassMoEPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize.FlashInferCutlassMoEPrepareAndFinalize] |

39 +| flashinfer_all2allv | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferA2APrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize.FlashInferA2APrepareAndFinalize] |

40 | MoEPrepareAndFinalizeNoEP5 | standard | fp8,int8 | G,A,T | N | Y | [`MoEPrepareAndFinalizeNoEP`][vllm.model_executor.layers.fused_moe.prepare_finalize.MoEPrepareAndFinalizeNoEP] |

41 | BatchedPrepareAndFinalize5 | batched | fp8,int8 | G,A,T | N | Y | [`BatchedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedPrepareAndFinalize] |

tests/kernels/moe/modular_kernel_tools/common.py

+8 -4

@@ -22,6 +22,9 @@ from vllm.distributed import (

22 )

23 from vllm.forward_context import set_forward_context

24 from vllm.model_executor.layers.fused_moe import fused_topk

25 +from vllm.model_executor.layers.fused_moe.all2all_utils import (

26 + maybe_make_prepare_finalize,

27 +)

28 from vllm.model_executor.layers.fused_moe.config import (

29 FusedMoEConfig,

30 FusedMoEParallelConfig,

@@ -40,7 +43,6 @@ from .mk_objects import (

43 TestMoEQuantConfig,

44 expert_info,

45 make_fused_experts,

43 - make_prepare_finalize,

46 prepare_finalize_info,

47 )

48 from .parallel_utils import ProcessGroupInfo

@@ -603,10 +605,12 @@ def make_modular_kernel(

605 routing_method=RoutingMethodType.DeepSeekV3,

606 )

607

606 - # make modular kernel

607 - prepare_finalize = make_prepare_finalize(

608 - config.prepare_finalize_type, config.all2all_backend(), moe, quant_config

608 + prepare_finalize = maybe_make_prepare_finalize(

609 + moe=moe,

610 + quant_config=quant_config,

611 + allow_new_interface=True,

612 )

613 + assert prepare_finalize is not None

614

615 fused_experts = make_fused_experts(

616 config.fused_experts_type,

tests/kernels/moe/modular_kernel_tools/mk_objects.py

+3 -25

@@ -7,9 +7,6 @@ import torch

7 # Fused experts and PrepareFinalize imports

8 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

9 from vllm.model_executor.layers.fused_moe import TritonExperts

10 -from vllm.model_executor.layers.fused_moe.all2all_utils import (

11 - maybe_make_prepare_finalize,

12 -)

10 from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (

11 BatchedDeepGemmExperts,

12 )

@@ -255,13 +252,12 @@ if has_pplx():

252 )

253

254 if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability(100):

255 + from vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize import ( # noqa: E501

256 + FlashInferCutlassMoEPrepareAndFinalize,

257 + )

258 from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (

259 FlashInferExperts,

260 )

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

262 - FlashInferCutlassMoEPrepareAndFinalize,

263 - create_flashinfer_prepare_finalize,

264 - )

261

262 register_prepare_and_finalize(

263 FlashInferCutlassMoEPrepareAndFinalize,

@@ -429,24 +425,6 @@ if cutlass_fp4_supported() or has_flashinfer_cutlass_fused_moe():

425 ]

426

427

432 -def make_prepare_finalize(

433 - prepare_finalize_type: mk.FusedMoEPrepareAndFinalize,

434 - backend: str | None,

435 - moe: FusedMoEConfig,

436 - quant_config: FusedMoEQuantConfig,

437 -) -> mk.FusedMoEPrepareAndFinalize:

438 - if backend != "naive" and backend is not None:

439 - prepare_finalize = maybe_make_prepare_finalize(moe, quant_config)

440 - assert prepare_finalize is not None

441 - return prepare_finalize

442 - elif prepare_finalize_type == FlashInferCutlassMoEPrepareAndFinalize:

443 - return create_flashinfer_prepare_finalize(

444 - use_dp=moe.moe_parallel_config.dp_size > 1

445 - )

446 - else:

447 - return MoEPrepareAndFinalizeNoEP()

448 -

449 -

428 def _slice(rank: int, num_local_experts: int, t: torch.Tensor) -> torch.Tensor:

429 s = rank * num_local_experts

430 e = s + num_local_experts

tests/kernels/moe/test_flashinfer.py

+1 -6

@@ -294,12 +294,7 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(

294 )

295

296 kernel = mk.FusedMoEModularKernel(

297 - MoEPrepareAndFinalizeNoEP(

298 - defer_input_quant=FlashInferExperts.expects_unquantized_inputs(

299 - moe_config=moe_config,

300 - quant_config=quant_config,

301 - )

302 - ),

297 + MoEPrepareAndFinalizeNoEP(),

298 FlashInferExperts(

299 moe_config=moe_config,

300 quant_config=quant_config,

tests/kernels/moe/test_flashinfer_moe.py

+1 -6

@@ -106,12 +106,7 @@ def test_flashinfer_fp4_moe_no_graph(

106 )

107

108 flashinfer_experts = FusedMoEModularKernel(

109 - MoEPrepareAndFinalizeNoEP(

110 - defer_input_quant=FlashInferExperts.expects_unquantized_inputs(

111 - moe_config=moe_config,

112 - quant_config=quant_config,

113 - )

114 - ),

109 + MoEPrepareAndFinalizeNoEP(),

110 FlashInferExperts(moe_config=moe_config, quant_config=quant_config),

111 )

112

tests/kernels/moe/test_nvfp4_moe.py

+1 -1

@@ -90,7 +90,7 @@ def test_cutlass_fp4_moe_no_graph(

90 )

92 kernel = mk.FusedMoEModularKernel(

93 - MoEPrepareAndFinalizeNoEP(defer_input_quant=True),

93 + MoEPrepareAndFinalizeNoEP(),

94 CutlassExpertsFp4(

95 moe_config=make_dummy_moe_config(),

96 quant_config=quant_config,

vllm/distributed/device_communicators/all2all.py

+98 -4

@@ -59,7 +59,7 @@ class NaiveAll2AllManager(All2AllManagerBase):

60 return buffer

62 - def dispatch(

62 + def dispatch_router_logits(

63 self,

64 hidden_states: torch.Tensor,

65 router_logits: torch.Tensor,

@@ -84,6 +84,34 @@ class NaiveAll2AllManager(All2AllManagerBase):

85 return hidden_states, router_logits

87 + def dispatch(

88 + self,

89 + hidden_states: torch.Tensor,

90 + topk_weights: torch.Tensor,

91 + topk_ids: torch.Tensor,

92 + is_sequence_parallel: bool = False,

93 + extra_tensors: list[torch.Tensor] | None = None,

94 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

95 + if extra_tensors is not None:

96 + raise NotImplementedError(

97 + "extra_tensors is not supported for NaiveAll2AllManager"

98 + )

99 + sp_size = self.tp_group.world_size if is_sequence_parallel else 1

100 + dp_metadata = get_forward_context().dp_metadata

101 + assert dp_metadata is not None

102 + cu_tokens_across_sp_cpu = dp_metadata.cu_tokens_across_sp(sp_size)

103 +

104 + hidden_states = self.naive_multicast(

105 + hidden_states, cu_tokens_across_sp_cpu, is_sequence_parallel

106 + )

107 + topk_weights = self.naive_multicast(

108 + topk_weights, cu_tokens_across_sp_cpu, is_sequence_parallel

109 + )

110 + topk_ids = self.naive_multicast(

111 + topk_ids, cu_tokens_across_sp_cpu, is_sequence_parallel

112 + )

113 + return hidden_states, topk_weights, topk_ids

114 +

115 def combine(

116 self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False

117 ) -> torch.Tensor:

@@ -114,7 +142,7 @@ class AgRsAll2AllManager(All2AllManagerBase):

142 def __init__(self, cpu_group):

143 super().__init__(cpu_group)

144

117 - def dispatch(

145 + def dispatch_router_logits(

146 self,

147 hidden_states: torch.Tensor,

148 router_logits: torch.Tensor,

@@ -148,6 +176,46 @@ class AgRsAll2AllManager(All2AllManagerBase):

176 return (gathered_tensors[0], gathered_tensors[1], gathered_tensors[2:])

177 return gathered_tensors[0], gathered_tensors[1]

178

179 + def dispatch(

180 + self,

181 + hidden_states: torch.Tensor,

182 + topk_weights: torch.Tensor,

183 + topk_ids: torch.Tensor,

184 + is_sequence_parallel: bool = False,

185 + extra_tensors: list[torch.Tensor] | None = None,

186 + ) -> (

187 + tuple[torch.Tensor, torch.Tensor, torch.Tensor]

188 + | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]

189 + ):

190 + """

191 + Gather hidden_states and router_logits from all dp ranks.

192 + """

193 + dp_metadata = get_forward_context().dp_metadata

194 + assert dp_metadata is not None

195 + sizes = dp_metadata.get_chunk_sizes_across_dp_rank()

196 + assert sizes is not None

197 + dist_group = get_ep_group() if is_sequence_parallel else get_dp_group()

198 + assert sizes[dist_group.rank_in_group] == hidden_states.shape[0]

199 +

200 + tensors_to_gather = [hidden_states, topk_weights, topk_ids]

201 + if extra_tensors is not None:

202 + tensors_to_gather.extend(extra_tensors)

203 +

204 + gathered_tensors = dist_group.all_gatherv(

205 + tensors_to_gather,

206 + dim=0,

207 + sizes=sizes,

208 + )

209 +

210 + hidden_states = gathered_tensors[0]

211 + topk_weights = gathered_tensors[1]

212 + topk_ids = gathered_tensors[2]

213 +

214 + if extra_tensors is None:

215 + return hidden_states, topk_weights, topk_ids

216 +

217 + return hidden_states, topk_weights, topk_ids, gathered_tensors[3:]

218 +

219 def combine(

220 self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False

221 ) -> torch.Tensor:

@@ -216,7 +284,7 @@ class PPLXAll2AllManager(All2AllManagerBase):

284 pplx.AllToAll.internode if self.internode else pplx.AllToAll.intranode,

285 )

286

219 - def dispatch(

287 + def dispatch_router_logits(

288 self,

289 hidden_states: torch.Tensor,

290 router_logits: torch.Tensor,

@@ -225,6 +293,19 @@ class PPLXAll2AllManager(All2AllManagerBase):

293 ) -> tuple[torch.Tensor, torch.Tensor]:

294 raise NotImplementedError

295

296 + def dispatch(

297 + self,

298 + hidden_states: torch.Tensor,

299 + topk_weights: torch.Tensor,

300 + topk_ids: torch.Tensor,

301 + is_sequence_parallel: bool = False,

302 + extra_tensors: list[torch.Tensor] | None = None,

303 + ) -> (

304 + tuple[torch.Tensor, torch.Tensor, torch.Tensor]

305 + | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]

306 + ):

307 + raise NotImplementedError

308 +

309 def combine(

310 self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False

311 ) -> torch.Tensor:

@@ -264,7 +345,7 @@ class DeepEPAll2AllManagerBase(All2AllManagerBase):

345 def get_handle(self, kwargs):

346 raise NotImplementedError

347

267 - def dispatch(

348 + def dispatch_router_logits(

349 self,

350 hidden_states: torch.Tensor,

351 router_logits: torch.Tensor,

@@ -273,6 +354,19 @@ class DeepEPAll2AllManagerBase(All2AllManagerBase):

354 ) -> tuple[torch.Tensor, torch.Tensor]:

355 raise NotImplementedError

356

357 + def dispatch(

358 + self,

359 + hidden_states: torch.Tensor,

360 + topk_weights: torch.Tensor,

361 + topk_ids: torch.Tensor,

362 + is_sequence_parallel: bool = False,

363 + extra_tensors: list[torch.Tensor] | None = None,

364 + ) -> (

365 + tuple[torch.Tensor, torch.Tensor, torch.Tensor]

366 + | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]

367 + ):

368 + raise NotImplementedError

369 +

370 def combine(

371 self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False

372 ) -> torch.Tensor:

vllm/distributed/device_communicators/base_device_communicator.py

+43 -4

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import threading

4 -from typing import Any

4 from weakref import WeakValueDictionary

6 import torch

@@ -64,13 +63,32 @@ class All2AllManagerBase:

63 # and reuse it for the same config.

64 raise NotImplementedError

67 - def dispatch(

66 + def dispatch_router_logits(

67 self,

68 hidden_states: torch.Tensor,

69 router_logits: torch.Tensor,

70 is_sequence_parallel: bool = False,

71 extra_tensors: list[torch.Tensor] | None = None,

73 - ) -> Any:

72 + ) -> (

73 + tuple[torch.Tensor, torch.Tensor]

74 + | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]

75 + ):

76 + # Subclasses should either:

77 + # - implement handling for extra_tensors, or

78 + # - raise a clear error if extra_tensors is not supported.

79 + raise NotImplementedError

80 +

81 + def dispatch(

82 + self,

83 + hidden_states: torch.Tensor,

84 + topk_weights: torch.Tensor,

85 + topk_ids: torch.Tensor,

86 + is_sequence_parallel: bool = False,

87 + extra_tensors: list[torch.Tensor] | None = None,

88 + ) -> (

89 + tuple[torch.Tensor, torch.Tensor, torch.Tensor]

90 + | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]

91 + ):

92 # Subclasses should either:

93 # - implement handling for extra_tensors, or

94 # - raise a clear error if extra_tensors is not supported.

@@ -280,7 +298,7 @@ class DeviceCommunicatorBase:

298 for module in moe_modules:

299 module.maybe_init_modular_kernel()

300

283 - def dispatch(

301 + def dispatch_router_logits(

302 self,

303 hidden_states: torch.Tensor,

304 router_logits: torch.Tensor,

@@ -294,8 +312,29 @@ class DeviceCommunicatorBase:

312 Dispatch the hidden states and router logits to the appropriate device.

313 This is a no-op in the base class.

314 """

315 + if extra_tensors is not None:

316 + return hidden_states, router_logits, extra_tensors

317 return hidden_states, router_logits

318

319 + def dispatch(

320 + self,

321 + hidden_states: torch.Tensor,

322 + topk_weights: torch.Tensor,

323 + topk_ids: torch.Tensor,

324 + is_sequence_parallel: bool = False,

325 + extra_tensors: list[torch.Tensor] | None = None,

326 + ) -> (

327 + tuple[torch.Tensor, torch.Tensor, torch.Tensor]

328 + | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]

329 + ):

330 + """

331 + Dispatch the hidden states and topk weights/ids to the appropriate device.

332 + This is a no-op in the base class.

333 + """

334 + if extra_tensors is not None:

335 + return hidden_states, topk_weights, topk_ids, extra_tensors

336 + return hidden_states, topk_weights, topk_ids

337 +

338 def combine(

339 self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False

340 ) -> torch.Tensor:

vllm/distributed/device_communicators/cpu_communicator.py

+43 -7

@@ -130,29 +130,65 @@ class CpuCommunicator(DeviceCommunicatorBase):

130 ) -> dict[str, torch.Tensor | Any]:

131 return self.dist_module.recv_tensor_dict(src)

132

133 - def dispatch( # type: ignore[override]

133 + def dispatch_router_logits(

134 self,

135 hidden_states: torch.Tensor,

136 router_logits: torch.Tensor,

137 is_sequence_parallel: bool = False,

138 extra_tensors: list[torch.Tensor] | None = None,

139 - ) -> tuple[torch.Tensor, torch.Tensor]:

139 + ) -> (

140 + tuple[torch.Tensor, torch.Tensor]

141 + | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]

142 + ):

143 + """

144 + Dispatch the hidden states and router logits to the appropriate device.

145 + This is a no-op in the base class.

146 + """

147 +

148 assert self.all2all_manager is not None

141 - return self.all2all_manager.dispatch(

149 + return self.all2all_manager.dispatch_router_logits(

150 hidden_states,

151 router_logits,

152 is_sequence_parallel,

145 - extra_tensors, # type: ignore[call-arg]

153 + extra_tensors,

154 + )

155 +

156 + def dispatch(

157 + self,

158 + hidden_states: torch.Tensor,

159 + topk_weights: torch.Tensor,

160 + topk_ids: torch.Tensor,

161 + is_sequence_parallel: bool = False,

162 + extra_tensors: list[torch.Tensor] | None = None,

163 + ) -> (

164 + tuple[torch.Tensor, torch.Tensor, torch.Tensor]

165 + | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]

166 + ):

167 + """

168 + Dispatch the hidden states and topk weights/ids to the appropriate device.

169 + This is a no-op in the base class.

170 + """

171 + assert self.all2all_manager is not None

172 + return self.all2all_manager.dispatch(

173 + hidden_states,

174 + topk_weights,

175 + topk_ids,

176 + is_sequence_parallel,

177 + extra_tensors=extra_tensors,

178 )

179

180 def combine(

181 self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False

182 ) -> torch.Tensor:

183 + """

184 + Combine the hidden states and router logits from the appropriate device.

185 + This is a no-op in the base class.

186 + """

187 assert self.all2all_manager is not None

152 - hidden_states = self.all2all_manager.combine(

153 - hidden_states, is_sequence_parallel

188 + return self.all2all_manager.combine(

189 + hidden_states,

190 + is_sequence_parallel,

191 )

155 - return hidden_states

192

193

194 class _CPUSHMDistributed:

vllm/distributed/device_communicators/cuda_communicator.py

+39 -6

@@ -322,7 +322,7 @@ class CudaCommunicator(DeviceCommunicatorBase):

322

323 return output_list

324

325 - def dispatch( # type: ignore[override]

325 + def dispatch_router_logits(

326 self,

327 hidden_states: torch.Tensor,

328 router_logits: torch.Tensor,

@@ -332,19 +332,52 @@ class CudaCommunicator(DeviceCommunicatorBase):

332 tuple[torch.Tensor, torch.Tensor]

333 | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]

334 ):

335 + """

336 + Dispatch the hidden states and router logits to the appropriate device.

337 + This is a no-op in the base class.

338 + """

339 +

340 assert self.all2all_manager is not None

336 - return self.all2all_manager.dispatch(

341 + return self.all2all_manager.dispatch_router_logits(

342 hidden_states,

343 router_logits,

344 is_sequence_parallel,

340 - extra_tensors, # type: ignore[call-arg]

345 + extra_tensors,

346 + )

347 +

348 + def dispatch(

349 + self,

350 + hidden_states: torch.Tensor,

351 + topk_weights: torch.Tensor,

352 + topk_ids: torch.Tensor,

353 + is_sequence_parallel: bool = False,

354 + extra_tensors: list[torch.Tensor] | None = None,

355 + ) -> (

356 + tuple[torch.Tensor, torch.Tensor, torch.Tensor]

357 + | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]

358 + ):

359 + """

360 + Dispatch the hidden states and topk weights/ids to the appropriate device.

361 + This is a no-op in the base class.

362 + """

363 + assert self.all2all_manager is not None

364 + return self.all2all_manager.dispatch(

365 + hidden_states,

366 + topk_weights,

367 + topk_ids,

368 + is_sequence_parallel,

369 + extra_tensors=extra_tensors,

370 )

371

372 def combine(

373 self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False

374 ) -> torch.Tensor:

375 + """

376 + Combine the hidden states and router logits from the appropriate device.

377 + This is a no-op in the base class.

378 + """

379 assert self.all2all_manager is not None

347 - hidden_states = self.all2all_manager.combine(

348 - hidden_states, is_sequence_parallel

380 + return self.all2all_manager.combine(

381 + hidden_states,

382 + is_sequence_parallel,

383 )

350 - return hidden_states

vllm/distributed/device_communicators/mnnvl_compat.py

+11

@@ -1,5 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from typing import Any

4 +

5 import torch.distributed as dist

6 from flashinfer.comm.mnnvl import CommBackend as CommBackend

@@ -23,5 +25,14 @@ class CustomCommunicator(CommBackend):

25 dist.all_gather_object(gathered, data, group=self._group)

26 return gathered

28 + # NOTE(rob): CommBackend is an abstract class, and bcast/barrier

29 + # are unimplemented on vLLM side. If we need to utilize these

30 + # methods in the future, can create a concrete implementation.

31 + def bcast(self, data: Any, root: int) -> Any:

32 + raise NotImplementedError

33 +

34 + def barrier(self) -> None:

35 + raise NotImplementedError

36 +

37 def Split(self, color: int, key: int) -> "CustomCommunicator":

38 return self

vllm/distributed/device_communicators/xpu_communicator.py

+43 -7

@@ -196,26 +196,62 @@ class XpuCommunicator(DeviceCommunicatorBase):

196 def broadcast(self, input_: torch.Tensor, src: int = 0) -> None:

197 dist.broadcast(input_, src=src, group=self.device_group)

198

199 - def dispatch(

199 + def dispatch_router_logits(

200 self,

201 hidden_states: torch.Tensor,

202 router_logits: torch.Tensor,

203 is_sequence_parallel: bool = False,

204 extra_tensors: list[torch.Tensor] | None = None,

205 - ) -> tuple[torch.Tensor, torch.Tensor]:

205 + ) -> (

206 + tuple[torch.Tensor, torch.Tensor]

207 + | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]

208 + ):

209 + """

210 + Dispatch the hidden states and router logits to the appropriate device.

211 + This is a no-op in the base class.

212 + """

213 +

214 assert self.all2all_manager is not None

207 - return self.all2all_manager.dispatch(

215 + return self.all2all_manager.dispatch_router_logits(

216 hidden_states,

217 router_logits,

218 is_sequence_parallel,

211 - extra_tensors, # type: ignore[call-arg]

219 + extra_tensors,

220 + )

221 +

222 + def dispatch(

223 + self,

224 + hidden_states: torch.Tensor,

225 + topk_weights: torch.Tensor,

226 + topk_ids: torch.Tensor,

227 + is_sequence_parallel: bool = False,

228 + extra_tensors: list[torch.Tensor] | None = None,

229 + ) -> (

230 + tuple[torch.Tensor, torch.Tensor, torch.Tensor]

231 + | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]

232 + ):

233 + """

234 + Dispatch the hidden states and topk weights/ids to the appropriate device.

235 + This is a no-op in the base class.

236 + """

237 + assert self.all2all_manager is not None

238 + return self.all2all_manager.dispatch(

239 + hidden_states,

240 + topk_weights,

241 + topk_ids,

242 + is_sequence_parallel,

243 + extra_tensors=extra_tensors,

244 )

245

246 def combine(

247 self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False

248 ) -> torch.Tensor:

249 + """

250 + Combine the hidden states and router logits from the appropriate device.

251 + This is a no-op in the base class.

252 + """

253 assert self.all2all_manager is not None

218 - hidden_states = self.all2all_manager.combine(

219 - hidden_states, is_sequence_parallel

254 + return self.all2all_manager.combine(

255 + hidden_states,

256 + is_sequence_parallel,

257 )

221 - return hidden_states

vllm/distributed/parallel_state.py

+24 -2

@@ -1000,7 +1000,7 @@ class GroupCoordinator:

1000 if self.device_communicator is not None:

1001 self.device_communicator.prepare_communication_buffer_for_model(model)

1002

1003 - def dispatch(

1003 + def dispatch_router_logits(

1004 self,

1005 hidden_states: torch.Tensor,

1006 router_logits: torch.Tensor,

@@ -1011,7 +1011,7 @@ class GroupCoordinator:

1011 | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]

1012 ):

1013 if self.device_communicator is not None:

1014 - return self.device_communicator.dispatch( # type: ignore[call-arg]

1014 + return self.device_communicator.dispatch_router_logits(

1015 hidden_states,

1016 router_logits,

1017 is_sequence_parallel,

@@ -1020,6 +1020,28 @@ class GroupCoordinator:

1020 else:

1021 return hidden_states, router_logits

1022

1023 + def dispatch(

1024 + self,

1025 + hidden_states: torch.Tensor,

1026 + topk_weights: torch.Tensor,

1027 + topk_ids: torch.Tensor,

1028 + is_sequence_parallel: bool = False,

1029 + extra_tensors: list[torch.Tensor] | None = None,

1030 + ) -> (

1031 + tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]

1032 + | tuple[torch.Tensor, torch.Tensor, torch.Tensor]

1033 + ):

1034 + if self.device_communicator is not None:

1035 + return self.device_communicator.dispatch(

1036 + hidden_states,

1037 + topk_weights,

1038 + topk_ids,

1039 + is_sequence_parallel,

1040 + extra_tensors,

1041 + )

1042 + else:

1043 + return hidden_states, topk_weights, topk_ids

1044 +

1045 def combine(

1046 self, hidden_states, is_sequence_parallel: bool = False

1047 ) -> torch.Tensor:

vllm/model_executor/layers/fused_moe/all2all_utils.py

+54 -6

@@ -7,17 +7,27 @@ import torch

7 from vllm.distributed import (

8 get_ep_group,

9 )

10 +from vllm.logger import init_logger

11 from vllm.model_executor.layers.fused_moe.config import (

12 FusedMoEConfig,

13 FusedMoEParallelConfig,

14 FusedMoEQuantConfig,

15 )

16 +from vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize import (

17 + FlashInferA2APrepareAndFinalize,

18 +)

19 from vllm.model_executor.layers.fused_moe.modular_kernel import (

20 FusedMoEPrepareAndFinalize,

21 )

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

23 + MoEPrepareAndFinalizeNaiveEP,

24 + MoEPrepareAndFinalizeNoEP,

25 +)

26 from vllm.platforms import current_platform

27 from vllm.utils.import_utils import has_deep_ep, has_mori, has_pplx

29 +logger = init_logger(__name__)

30 +

31 if current_platform.is_cuda_alike():

32 if has_pplx():

33 from .pplx_prepare_finalize import (

@@ -70,20 +80,46 @@ def maybe_make_prepare_finalize(

80 moe: FusedMoEConfig,

81 quant_config: FusedMoEQuantConfig | None,

82 routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

83 + allow_new_interface: bool = False,

84 ) -> FusedMoEPrepareAndFinalize | None:

85 + # NOTE(rob): we are migrating each quant_method to hold the MK

86 + # in all cases. The allow_new_interface=False flag allow us to fall

87 + # back to the old method for methods that have not yet been migrated.

88 + #

89 + # In old method:

90 + # * maybe_init_modular_kernel() calls this function. If we are

91 + # using no Dp/Ep or naive all2all, we return None this function

92 + # returns None and no ModularKernelMethod is created. If non-naive

93 + # all2all is used, this returns a PrepareAndFinalize object and

94 + # a ModularKernelMethod is created.

95 + # In new method:

96 + # * maybe_make_prepare_finalize() is called from the oracle. We

97 + # always return a PrepareAndFinalize object and the quant method

98 + # holds the ModularKernel.

99 if not moe.moe_parallel_config.use_all2all_kernels:

75 - return None

100 + if not allow_new_interface:

101 + return None

102 +

103 + # For DP/TP case, fall back to naive P/F.

104 + if moe.moe_parallel_config.dp_size > 1:

105 + logger.info_once(

106 + "Detected DP deployment with no --enable-expert-parallel. "

107 + "Falling back to AllGather+ReduceScatter dispatch/combine."

108 + )

109 + return MoEPrepareAndFinalizeNaiveEP(

110 + is_sequence_parallel=moe.moe_parallel_config.is_sequence_parallel,

111 + num_dispatchers=(

112 + get_ep_group().device_communicator.all2all_manager.world_size

113 + ),

114 + )

115 + else:

116 + return MoEPrepareAndFinalizeNoEP()

117

118 all2all_manager = get_ep_group().device_communicator.all2all_manager

119 assert all2all_manager is not None

120

121 prepare_finalize: FusedMoEPrepareAndFinalize | None = None

122

82 - # TODO(rob): update this as part of the MoE refactor.

83 - assert not moe.use_flashinfer_cutlass_kernels, (

84 - "Must be created in modelopt.py or fp8.py"

85 - )

86 -

123 if moe.use_pplx_kernels:

124 assert quant_config is not None

125

@@ -203,4 +239,16 @@ def maybe_make_prepare_finalize(

239 use_fp8_dispatch=use_fp8_dispatch,

240 )

241

242 + elif moe.use_fi_all2allv_kernels:

243 + assert quant_config is not None

244 + prepare_finalize = FlashInferA2APrepareAndFinalize(

245 + num_dispatchers=all2all_manager.world_size,

246 + )

247 +

248 + elif moe.use_naive_all2all_kernels and allow_new_interface:

249 + prepare_finalize = MoEPrepareAndFinalizeNaiveEP(

250 + is_sequence_parallel=(moe.moe_parallel_config.is_sequence_parallel),

251 + num_dispatchers=all2all_manager.world_size,

252 + )

253 +

254 return prepare_finalize

vllm/model_executor/layers/fused_moe/config.py

+16 -10

@@ -20,7 +20,6 @@ from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (

20 )

21 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape

22 from vllm.platforms import current_platform

23 -from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe

23 from vllm.utils.import_utils import has_triton_kernels

24 from vllm.utils.math_utils import cdiv

@@ -862,6 +861,7 @@ class FusedMoEParallelConfig:

861

862 use_ep: bool # whether to use EP or not

863 all2all_backend: str # all2all backend for MoE communication

864 + is_sequence_parallel: bool # whether sequence parallelism is used

865 enable_eplb: bool # whether to enable expert load balancing

866

867 @property

@@ -883,6 +883,12 @@ class FusedMoEParallelConfig:

883 def use_deepep_ll_kernels(self):

884 return self.use_all2all_kernels and self.all2all_backend == "deepep_low_latency"

885

886 + @property

887 + def use_fi_all2allv_kernels(self):

888 + return (

889 + self.use_all2all_kernels and self.all2all_backend == "flashinfer_all2allv"

890 + )

891 +

892 @property

893 def use_batched_activation_format(self):

894 return self.use_deepep_ll_kernels or self.use_pplx_kernels

@@ -1014,6 +1020,7 @@ class FusedMoEParallelConfig:

1020 ep_rank=0,

1021 use_ep=False,

1022 all2all_backend=vllm_parallel_config.all2all_backend,

1023 + is_sequence_parallel=vllm_parallel_config.use_sequence_parallel_moe,

1024 enable_eplb=vllm_parallel_config.enable_eplb,

1025 )

1026 # DP + EP / TP + EP / DP + TP + EP

@@ -1033,6 +1040,7 @@ class FusedMoEParallelConfig:

1040 ep_rank=ep_rank,

1041 use_ep=True,

1042 all2all_backend=vllm_parallel_config.all2all_backend,

1043 + is_sequence_parallel=vllm_parallel_config.use_sequence_parallel_moe,

1044 enable_eplb=vllm_parallel_config.enable_eplb,

1045 )

1046

@@ -1051,6 +1059,7 @@ class FusedMoEParallelConfig:

1059 use_ep=False,

1060 all2all_backend="naive",

1061 enable_eplb=False,

1062 + is_sequence_parallel=False,

1063 )

1064

1065

@@ -1145,12 +1154,9 @@ class FusedMoEConfig:

1154 return self.moe_parallel_config.use_mori_kernels

1155

1156 @property

1148 - def use_flashinfer_cutlass_kernels(self):

1149 - """

1150 - Whether to use FlashInfer cutlass kernels for NVFP4 MoE.

1151 - """

1152 - return (

1153 - envs.VLLM_USE_FLASHINFER_MOE_FP4

1154 - and has_flashinfer_cutlass_fused_moe()

1155 - and envs.VLLM_FLASHINFER_MOE_BACKEND == "throughput"

1156 - )

1157 + def use_fi_all2allv_kernels(self):

1158 + return self.moe_parallel_config.use_fi_all2allv_kernels

1159 +

1160 + @property

1161 + def use_naive_all2all_kernels(self):

1162 + return self.moe_parallel_config.use_naive_all2all_kernels

vllm/model_executor/layers/fused_moe/cutlass_moe.py

+14 -6

@@ -103,7 +103,14 @@ def run_cutlass_moe_fp8(

103 or a2_scale.size(0) == a1q.shape[0]

104 ), "Intermediate scale shape mismatch"

105 assert out_dtype in [torch.half, torch.bfloat16], "Invalid output dtype"

106 - if expert_map is not None:

106 +

107 + # NOTE(rob): the expert_map is used for the STANDARD case and

108 + # the batched format is used by the BATCHED case.

109 + # TODO(rob): update the MK interface to only pass the expert_map

110 + # during the STANDARD case to make this clearer across all kernels.

111 + if use_batched_format:

112 + assert expert_num_tokens is not None

113 + else:

114 assert expert_num_tokens is None

115

116 # We have two modes: batched experts and non-batched experts.

@@ -379,7 +386,10 @@ class CutlassExpertsFp8(CutlassExpertsFp8Base):

386 # needed for STANDARD activation format kernels in DP/EP mode.

387 # Note that the BATCHED activation format does not use

388 # the expert map for identifying experts.

382 - return not moe_parallel_config.use_all2all_kernels

389 + return not (

390 + moe_parallel_config.use_fi_all2allv_kernels

391 + or moe_parallel_config.use_deepep_ht_kernels

392 + )

393

394 def supports_chunking(self) -> bool:

395 return True

@@ -641,10 +651,8 @@ def run_cutlass_moe_fp4(

651

652

653 class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):

644 - @staticmethod

645 - def expects_unquantized_inputs(

646 - moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig

647 - ) -> bool:

654 + @property

655 + def expects_unquantized_inputs(self) -> bool:

656 return True

657

658 @staticmethod

vllm/model_executor/layers/fused_moe/deep_gemm_moe.py

+2 -1

@@ -148,7 +148,8 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

148

149 @staticmethod

150 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:

151 - return True

151 + # NOTE(rob): discovered an IMA with this combination. Needs investigation.

152 + return not moe_parallel_config.use_fi_all2allv_kernels

153

154 def supports_chunking(self) -> bool:

155 return True

vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py

+24 -8

@@ -103,6 +103,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

103 num_experts: int,

104 a1_scale: torch.Tensor | None,

105 quant_config: FusedMoEQuantConfig,

106 + defer_input_quant: bool,

107 ) -> Callable:

108 has_scales = token_scales is not None

109

@@ -174,6 +175,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

175 expert_topk_weights,

176 a1_scale,

177 quant_config,

178 + defer_input_quant=defer_input_quant,

179 )

180

181 def _receiver(

@@ -187,6 +189,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

189 expert_topk_weights: torch.Tensor | None,

190 a1_scale: torch.Tensor | None,

191 quant_config: FusedMoEQuantConfig,

192 + defer_input_quant: bool,

193 ) -> mk.PrepareResultType:

194 if event.event is not None:

195 event.current_stream_wait()

@@ -221,14 +224,15 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

224 expert_num_tokens_per_expert_list, device=expert_x.device

225 )

226

224 - # Dispatch and Quant

225 - # DeepEP kernels only support dispatching block-quantized

226 - # activation scales.

227 - # Dispatch in bfloat16 and quantize afterwards

228 - if not quant_config.is_block_quantized:

227 + # * For non-block quant, dispatch in b16 and quantize now as

228 + # DeepEP kernels only support dispatching block scales.

229 + # * For expert kernels that require unquantized inputs,

230 + # defer quantization to FusedMoEExpertsPermuteUnpermute.

231 + if not quant_config.is_block_quantized and not defer_input_quant:

232 # Quantize after dispatch.

233 expert_x_scale = None

234 if expert_x.numel() != 0:

235 + # TODO: support per_act_token_quant,

236 expert_x, expert_x_scale = moe_kernel_quantize_input(

237 expert_x,

238 a1_scale,

@@ -257,6 +261,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

261 expert_map: torch.Tensor | None,

262 apply_router_weight_on_input: bool,

263 quant_config: FusedMoEQuantConfig,

264 + defer_input_quant: bool = False,

265 ) -> mk.ReceiverType:

266 if apply_router_weight_on_input:

267 topk = topk_ids.size(1)

@@ -266,8 +271,12 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

271 )

272 a1 = a1 * topk_weights.to(a1.dtype)

273

269 - if quant_config.is_block_quantized:

270 - # Quant and Dispatch

274 + # * DeepEP only supports fp8 block scales so quantize

275 + # before the dispatch for these models.

276 + # * For all other quantization, dispatch after.

277 + # * For expert kernels that require unquantized inputs,

278 + # defer quantization to FusedMoEExpertsPermuteUnpermute.

279 + if quant_config.is_block_quantized and not defer_input_quant:

280 a1q, a1q_scale = moe_kernel_quantize_input(

281 a1,

282 quant_config.a1_scale,

@@ -281,7 +290,11 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

290 else:

291 a1q = a1

292 a1q_scale = None

284 - a1_post_scale = quant_config.a1_scale

293 + a1_post_scale = (

294 + quant_config.a1_gscale

295 + if quant_config.quant_dtype == "nvfp4"

296 + else quant_config.a1_scale

297 + )

298

299 return self._do_dispatch(

300 tokens=a1q,

@@ -291,6 +304,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

304 num_experts=num_experts,

305 a1_scale=a1_post_scale,

306 quant_config=quant_config,

307 + defer_input_quant=defer_input_quant,

308 )

309

310 def prepare(

@@ -302,6 +316,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

316 expert_map: torch.Tensor | None,

317 apply_router_weight_on_input: bool,

318 quant_config: FusedMoEQuantConfig,

319 + defer_input_quant: bool = False,

320 ) -> mk.PrepareResultType:

321 receiver = self.prepare_async(

322 a1,

@@ -311,6 +326,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

326 expert_map,

327 apply_router_weight_on_input,

328 quant_config,

329 + defer_input_quant,

330 )

331 return receiver()

332

vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py

+13

@@ -242,7 +242,14 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

242 expert_map: torch.Tensor | None,

243 apply_router_weight_on_input: bool,

244 quant_config: FusedMoEQuantConfig,

245 + defer_input_quant: bool = False,

246 ) -> tuple[Callable, mk.ReceiverType]:

247 + if defer_input_quant:

248 + raise NotImplementedError(

249 + f"{self.__class__.__name__} does not support defer_input_quant=True. "

250 + "Please select an MoE kernel that accepts quantized inputs."

251 + )

252 +

253 hidden_size = a1.size(1)

254 assert hidden_size in self.SUPPORTED_HIDDEN_SIZES, (

255 f"Hidden Size {hidden_size} not in supported list of hidden sizes"

@@ -344,7 +351,13 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

351 expert_map: torch.Tensor | None,

352 apply_router_weight_on_input: bool,

353 quant_config: FusedMoEQuantConfig,

354 + defer_input_quant: bool = False,

355 ) -> mk.PrepareResultType:

356 + if defer_input_quant:

357 + raise NotImplementedError(

358 + f"{self.__class__.__name__} does not support defer_input_quant=True. "

359 + "Please select an MoE kernel that accepts quantized inputs."

360 + )

361 hook, receiver = self.prepare_async(

362 a1,

363 topk_weights,

vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py new

+226

@@ -0,0 +1,226 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +import torch

5 +

6 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk

7 +from vllm.distributed import get_ep_group

8 +from vllm.distributed.device_communicators.base_device_communicator import (

9 + All2AllManagerBase,

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.utils import moe_kernel_quantize_input

14 +from vllm.utils.flashinfer import nvfp4_block_scale_interleave

15 +

16 +

17 +def get_local_sizes():

18 + return get_forward_context().dp_metadata.get_chunk_sizes_across_dp_rank()

19 +

20 +

21 +class FlashInferA2APrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

22 + """Base class for FlashInfer MoE prepare and finalize operations."""

23 +

24 + def __init__(

25 + self,

26 + num_dispatchers: int = 1,

27 + ):

28 + super().__init__()

29 + self.num_dispatchers_ = num_dispatchers

30 + self.all2all_manager = get_ep_group().device_communicator.all2all_manager

31 +

32 + @property

33 + def activation_format(self) -> mk.FusedMoEActivationFormat:

34 + return mk.FusedMoEActivationFormat.Standard

35 +

36 + def max_num_tokens_per_rank(self) -> int | None:

37 + return None

38 +

39 + def topk_indices_dtype(self) -> torch.dtype | None:

40 + return None

41 +

42 + def num_dispatchers(self) -> int:

43 + return self.num_dispatchers_

44 +

45 + def output_is_reduced(self) -> bool:

46 + return False

47 +

48 + def _apply_router_weight_on_input(

49 + self,

50 + a1: torch.Tensor,

51 + topk_weights: torch.Tensor,

52 + topk_ids: torch.Tensor,

53 + apply_router_weight_on_input: bool,

54 + ) -> None:

55 + """Apply router weight on input if needed."""

56 + if apply_router_weight_on_input:

57 + topk = topk_ids.size(1)

58 + assert topk == 1, (

59 + "apply_router_weight_on_input is only implemented for topk=1"

60 + )

61 + a1.mul_(topk_weights.to(a1.dtype))

62 +

63 + def prepare(

64 + self,

65 + a1: torch.Tensor,

66 + topk_weights: torch.Tensor,

67 + topk_ids: torch.Tensor,

68 + num_experts: int,

69 + expert_map: torch.Tensor | None,

70 + apply_router_weight_on_input: bool,

71 + quant_config: FusedMoEQuantConfig,

72 + defer_input_quant: bool = False,

73 + ) -> mk.PrepareResultType:

74 + self._apply_router_weight_on_input(

75 + a1, topk_weights, topk_ids, apply_router_weight_on_input

76 + )

77 + global_num_tokens_cpu = get_local_sizes()

78 + top_k = topk_ids.size(1)

79 +

80 + (self.alltoall_info, topk_ids, topk_weights, a1q, a1q_scale) = (

81 + flashinfer_alltoall_dispatch(

82 + self.all2all_manager,

83 + global_num_tokens_cpu,

84 + a1,

85 + quant_config.a1_gscale,

86 + topk_ids,

87 + topk_weights,

88 + top_k,

89 + num_experts,

90 + quant_config,

91 + defer_input_quant=defer_input_quant,

92 + )

93 + )

94 +

95 + return a1q, a1q_scale, None, topk_ids, topk_weights

96 +

97 + def finalize(

98 + self,

99 + output: torch.Tensor,

100 + fused_expert_output: torch.Tensor,

101 + topk_weights: torch.Tensor,

102 + topk_ids: torch.Tensor,

103 + apply_router_weight_on_input: bool,

104 + weight_and_reduce_impl: mk.TopKWeightAndReduce,

105 + ) -> None:

106 + top_k = topk_ids.size(1)

107 + token_count = output.shape[0]

108 + fused_expert_output = flashinfer_alltoall_combine(

109 + self.all2all_manager,

110 + fused_expert_output,

111 + top_k=top_k,

112 + token_count=token_count,

113 + alltoall_info=self.alltoall_info,

114 + )

115 + output.copy_(fused_expert_output)

116 +

117 +

118 +def flashinfer_alltoall_dispatch(

119 + all2all_manager: All2AllManagerBase,

120 + global_num_tokens_cpu: list[int],

121 + x: torch.Tensor,

122 + gs: torch.Tensor,

123 + topk_ids: torch.Tensor,

124 + topk_weights: torch.Tensor,

125 + top_k: int,

126 + num_experts: int,

127 + quant_config: FusedMoEQuantConfig,

128 + defer_input_quant: bool = False,

129 +):

130 + from flashinfer.comm.trtllm_alltoall import MnnvlMoe

131 +

132 + assert all2all_manager.ensure_alltoall_workspace_initialized(), (

133 + "FlashInfer AllToAll workspace not available"

134 + )

135 +

136 + ep_rank = all2all_manager.rank

137 + ep_size = all2all_manager.world_size

138 + max_num_token = (

139 + max(global_num_tokens_cpu) if global_num_tokens_cpu is not None else x.shape[0]

140 + )

141 + orig_topk_weights_dtype = topk_weights.dtype

142 + alltoall_info, topk_ids, topk_weights, _ = (

143 + MnnvlMoe.mnnvl_moe_alltoallv_prepare_without_allgather(

144 + topk_ids,

145 + topk_weights,

146 + None,

147 + all2all_manager.prepare_workspace_tensor,

148 + max_num_token,

149 + ep_rank,

150 + ep_size,

151 + num_experts,

152 + num_experts,

153 + top_k,

154 + )

155 + )

156 + topk_weights = topk_weights.view(dtype=orig_topk_weights_dtype)

157 +

158 + if not defer_input_quant:

159 + x, x_sf = moe_kernel_quantize_input(

160 + x,

161 + gs,

162 + quant_config.quant_dtype,

163 + quant_config.per_act_token_quant,

164 + quant_config.block_shape,

165 + # NOTE: swizzling pads the scales to multiple of 128

166 + # which makes the scales tensor different shape than

167 + # the hidden states, breaking the A2A kernel. So, we

168 + # delay the swizzling until after the A2A.

169 + is_fp4_scale_swizzled=False,

170 + )

171 +

172 + x = MnnvlMoe.mnnvl_moe_alltoallv(

173 + x,

174 + alltoall_info,

175 + all2all_manager.workspace_tensor,

176 + ep_rank,

177 + ep_size,

178 + )

179 +

180 + x_sf = MnnvlMoe.mnnvl_moe_alltoallv(

181 + x_sf,

182 + alltoall_info,

183 + all2all_manager.workspace_tensor,

184 + ep_rank,

185 + ep_size,

186 + )

187 +

188 + # Swizzle after the A2A if nvfp4.

189 + if quant_config.quant_dtype == "nvfp4":

190 + if x_sf.element_size() == 1:

191 + x_sf = x_sf.view(torch.uint8)

192 + x_sf = nvfp4_block_scale_interleave(x_sf)

193 + else:

194 + # Block-scale path: pass activations through without quantization

195 + x_sf = None

196 + x = MnnvlMoe.mnnvl_moe_alltoallv(

197 + x,

198 + alltoall_info,

199 + all2all_manager.workspace_tensor,

200 + ep_rank,

201 + ep_size,

202 + )

203 + return alltoall_info, topk_ids, topk_weights, x, x_sf

204 +

205 +

206 +def flashinfer_alltoall_combine(

207 + all2all_manager: All2AllManagerBase,

208 + output: torch.Tensor,

209 + top_k: int,

210 + token_count: int,

211 + alltoall_info,

212 +):

213 + from flashinfer.comm.trtllm_alltoall import MnnvlMoe

214 +

215 + assert all2all_manager.ensure_alltoall_workspace_initialized(), (

216 + "FlashInfer AllToAll workspace not available"

217 + )

218 + return MnnvlMoe.mnnvl_moe_alltoallv_combine(

219 + output,

220 + alltoall_info,

221 + all2all_manager.workspace_tensor,

222 + ep_rank=all2all_manager.rank,

223 + ep_size=all2all_manager.world_size,

224 + top_k=top_k,

225 + token_count=token_count,

226 + )

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py

+8 -16

@@ -78,16 +78,9 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):

78 # - skip input activation quantization (kernel applies scaling)

79 self.use_deepseek_fp8_block_scale = quant_config.is_block_quantized

81 - @staticmethod

82 - def expects_unquantized_inputs(

83 - moe_config: mk.FusedMoEConfig, quant_config: FusedMoEQuantConfig

84 - ) -> bool:

85 - # NVFP4 TP kernels and FP8 block-quantized kernels apply

86 - # input quantization inside FusedMoEPermuteExpertsUnpermute.

87 - return (

88 - quant_config.use_nvfp4_w4a4

89 - and not moe_config.moe_parallel_config.use_all2all_kernels

90 - ) or (quant_config.use_fp8_w8a8 and quant_config.is_block_quantized)

81 + @property

82 + def expects_unquantized_inputs(self) -> bool:

83 + return self.quant_config.use_fp8_w8a8 and self.quant_config.is_block_quantized

85 @staticmethod

86 def _supports_current_device() -> bool:

@@ -144,10 +137,8 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):

137 # FLASHINFER_CUTLASS currently uses its down P/F, which does not

138 # work with SP. This will be removed in follow up after we get

139 # rid of the FlashInfer specific P/F function.

147 - return (

148 - moe_parallel_config.dp_size == 1

149 - or moe_parallel_config.dp_size == moe_parallel_config.ep_size

150 - )

140 + # TODO: the per-tensor fp8 kernels don't work with MNNVL FI A2As.

141 + return not moe_parallel_config.is_sequence_parallel

142

143 @staticmethod

144 def activation_format() -> mk.FusedMoEActivationFormat:

@@ -194,8 +185,9 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):

185 """

186 workspace1 = (M, K)

187 workspace2 = (0,)

197 - # For TP, the quantization is fused with fused_moe call.

198 - output_shape = (M, K * 2 if self.quant_dtype == "nvfp4" and self.use_dp else K)

188 + # For NVFP4, the output is stored in a packed int8 format,

189 + # so the actual hidden dim is 2x the size of K here.

190 + output_shape = (M, K * 2 if self.quant_dtype == "nvfp4" else K)

191 # The workspace is determined by `aq`, since it comes after any

192 # potential communication op and is involved in the expert computation.

193 return (workspace1, workspace2, output_shape)

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py deleted

-373

@@ -1,373 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -

4 -import torch

5 -

6 -import vllm.model_executor.layers.fused_moe.modular_kernel as mk

7 -from vllm.distributed import get_dp_group, get_ep_group

8 -from vllm.distributed.device_communicators.base_device_communicator import (

9 - All2AllManagerBase,

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

19 -from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input

20 -from vllm.utils.flashinfer import nvfp4_block_scale_interleave

21 -

22 -

23 -def get_local_sizes():

24 - return get_forward_context().dp_metadata.get_chunk_sizes_across_dp_rank()

25 -

26 -

27 -class FlashInferCutlassMoEPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

28 - """Base class for FlashInfer MoE prepare and finalize operations."""

29 -

30 - def __init__(

31 - self,

32 - use_dp: bool,

33 - num_dispatchers: int = 1,

34 - use_deepseek_fp8_block_scale: bool = False,

35 - ):

36 - super().__init__()

37 - self.num_dispatchers_ = num_dispatchers

38 - self.use_dp = use_dp

39 - self.local_tokens = None

40 - # Toggle for DeepSeek-style FP8 block-scale path where activations are

41 - # not quantized here and weight block scales are consumed by the kernel.

42 - self.use_deepseek_fp8_block_scale = use_deepseek_fp8_block_scale

43 -

44 - @property

45 - def activation_format(self) -> mk.FusedMoEActivationFormat:

46 - return mk.FusedMoEActivationFormat.Standard

47 -

48 - def max_num_tokens_per_rank(self) -> int | None:

49 - return None

50 -

51 - def topk_indices_dtype(self) -> torch.dtype | None:

52 - return None

53 -

54 - def num_dispatchers(self) -> int:

55 - return self.num_dispatchers_

56 -

57 - def output_is_reduced(self) -> bool:

58 - return False

59 -

60 - def _apply_router_weight_on_input(

61 - self,

62 - a1: torch.Tensor,

63 - topk_weights: torch.Tensor,

64 - topk_ids: torch.Tensor,

65 - apply_router_weight_on_input: bool,

66 - ) -> None:

67 - """Apply router weight on input if needed."""

68 - if apply_router_weight_on_input:

69 - topk = topk_ids.size(1)

70 - assert topk == 1, (

71 - "apply_router_weight_on_input is only implemented for topk=1"

72 - )

73 - a1.mul_(topk_weights.to(a1.dtype))

74 -

75 -

76 -class FlashInferAllToAllMoEPrepareAndFinalize(FlashInferCutlassMoEPrepareAndFinalize):

77 - """FlashInfer implementation using AllToAll communication."""

78 -

79 - def __init__(

80 - self,

81 - use_dp: bool,

82 - num_dispatchers: int = 1,

83 - use_deepseek_fp8_block_scale: bool = False,

84 - ):

85 - super().__init__(use_dp, num_dispatchers, use_deepseek_fp8_block_scale)

86 - self.alltoall_info = None

87 -

88 - # Initialize all2all_manager only for DP case

89 - self.all2all_manager = None

90 - if self.use_dp:

91 - self.all2all_manager = get_ep_group().device_communicator.all2all_manager

92 -

93 - def prepare(

94 - self,

95 - a1: torch.Tensor,

96 - topk_weights: torch.Tensor,

97 - topk_ids: torch.Tensor,

98 - num_experts: int,

99 - expert_map: torch.Tensor | None,

100 - apply_router_weight_on_input: bool,

101 - quant_config: FusedMoEQuantConfig,

102 - ) -> mk.PrepareResultType:

103 - self._apply_router_weight_on_input(

104 - a1, topk_weights, topk_ids, apply_router_weight_on_input

105 - )

106 -

107 - if not self.use_dp:

108 - # Non-DP case: quantize activations unless using block-scale path

109 - if not self.use_deepseek_fp8_block_scale:

110 - a1q, a1q_scale = moe_kernel_quantize_input(

111 - a1,

112 - quant_config.a1_gscale,

113 - quant_config.quant_dtype,

114 - quant_config.per_act_token_quant,

115 - quant_config.block_shape,

116 - is_fp4_scale_swizzled=not self.use_dp,

117 - )

118 - else:

119 - a1q = a1

120 - a1q_scale = None

121 - else:

122 - # DP case: use FlashInfer AllToAll

123 - global_num_tokens_cpu = get_local_sizes()

124 - top_k = topk_ids.size(1)

125 -

126 - (self.alltoall_info, topk_ids, topk_weights, a1q, a1q_scale) = (

127 - flashinfer_alltoall_dispatch(

128 - self.all2all_manager,

129 - global_num_tokens_cpu,

130 - a1,

131 - quant_config.a1_gscale,

132 - topk_ids,

133 - topk_weights,

134 - top_k,

135 - num_experts,

136 - quant_config,

137 - use_deepseek_fp8_block_scale=self.use_deepseek_fp8_block_scale,

138 - )

139 - )

140 -

141 - return a1q, a1q_scale, None, topk_ids, topk_weights

142 -

143 - def finalize(

144 - self,

145 - output: torch.Tensor,

146 - fused_expert_output: torch.Tensor,

147 - topk_weights: torch.Tensor,

148 - topk_ids: torch.Tensor,

149 - apply_router_weight_on_input: bool,

150 - weight_and_reduce_impl: mk.TopKWeightAndReduce,

151 - ) -> None:

152 - if self.use_dp:

153 - top_k = topk_ids.size(1)

154 - token_count = output.shape[0]

155 - fused_expert_output = flashinfer_alltoall_combine(

156 - self.all2all_manager,

157 - fused_expert_output,

158 - top_k=top_k,

159 - token_count=token_count,

160 - alltoall_info=self.alltoall_info,

161 - )

162 - output.copy_(fused_expert_output)

163 -

164 -

165 -class FlashInferAllGatherMoEPrepareAndFinalize(FlashInferCutlassMoEPrepareAndFinalize):

166 - def __init__(

167 - self,

168 - use_dp: bool,

169 - num_dispatchers: int = 1,

170 - use_deepseek_fp8_block_scale: bool = False,

171 - ):

172 - super().__init__(use_dp, num_dispatchers, use_deepseek_fp8_block_scale)

173 -

174 - def prepare(

175 - self,

176 - a1: torch.Tensor,

177 - topk_weights: torch.Tensor,

178 - topk_ids: torch.Tensor,

179 - num_experts: int,

180 - expert_map: torch.Tensor | None,

181 - apply_router_weight_on_input: bool,

182 - quant_config: FusedMoEQuantConfig,

183 - ) -> mk.PrepareResultType:

184 - self._apply_router_weight_on_input(

185 - a1, topk_weights, topk_ids, apply_router_weight_on_input

186 - )

187 - is_nvfp4 = quant_config.quant_dtype == "nvfp4"

188 - if not self.use_dp and is_nvfp4:

189 - return a1, None, None, topk_ids, topk_weights

190 -

191 - if not self.use_deepseek_fp8_block_scale:

192 - a1q, a1q_scale = moe_kernel_quantize_input(

193 - a1,

194 - quant_config.a1_gscale if is_nvfp4 else quant_config.a1_scale,

195 - quant_config.quant_dtype,

196 - quant_config.per_act_token_quant,

197 - quant_config.block_shape,

198 - is_fp4_scale_swizzled=not self.use_dp,

199 - )

200 - else:

201 - # Block-scale path: pass activations through, omit per-token scales

202 - a1q = a1

203 - a1q_scale = None

204 -

205 - if self.use_dp:

206 - # Build gather list conditionally - omit a1q_scale if None

207 - # (block-scale path)

208 - gather_list = [topk_weights, topk_ids, a1q]

209 - if a1q_scale is not None:

210 - gather_list.append(a1q_scale)

211 - gathered = get_dp_group().all_gatherv(

212 - gather_list,

213 - dim=0,

214 - sizes=get_local_sizes(),

215 - )

216 - topk_weights, topk_ids, a1q, a1q_scale = gathered

217 - else:

218 - gathered = get_dp_group().all_gatherv(

219 - gather_list,

220 - dim=0,

221 - sizes=get_local_sizes(),

222 - )

223 - topk_weights, topk_ids, a1q = gathered

224 - a1q_scale = None

225 -

226 - if is_nvfp4 and a1q_scale is not None:

227 - if a1q_scale.element_size() == 1:

228 - a1q_scale = a1q_scale.view(torch.uint8)

229 - a1q_scale = nvfp4_block_scale_interleave(a1q_scale)

230 -

231 - return a1q, a1q_scale, None, topk_ids, topk_weights

232 -

233 - def finalize(

234 - self,

235 - output: torch.Tensor,

236 - fused_expert_output: torch.Tensor,

237 - topk_weights: torch.Tensor,

238 - topk_ids: torch.Tensor,

239 - apply_router_weight_on_input: bool,

240 - weight_and_reduce_impl: mk.TopKWeightAndReduce,

241 - ) -> None:

242 - assert isinstance(weight_and_reduce_impl, TopKWeightAndReduceNoOP)

243 -

244 - if self.use_dp:

245 - fused_expert_output = get_dp_group().reduce_scatterv(

246 - fused_expert_output, dim=0, sizes=get_local_sizes()

247 - )

248 - output.copy_(fused_expert_output)

249 -

250 -

251 -def flashinfer_alltoall_dispatch(

252 - all2all_manager: All2AllManagerBase,

253 - global_num_tokens_cpu: list[int],

254 - x: torch.Tensor,

255 - gs: torch.Tensor,

256 - topk_ids: torch.Tensor,

257 - topk_weights: torch.Tensor,

258 - top_k: int,

259 - num_experts: int,

260 - quant_config: FusedMoEQuantConfig,

261 - use_deepseek_fp8_block_scale: bool = False,

262 -):

263 - from flashinfer.comm.trtllm_alltoall import MnnvlMoe

264 -

265 - assert all2all_manager.ensure_alltoall_workspace_initialized(), (

266 - "FlashInfer AllToAll workspace not available"

267 - )

268 -

269 - ep_rank = all2all_manager.rank

270 - ep_size = all2all_manager.world_size

271 - max_num_token = (

272 - max(global_num_tokens_cpu) if global_num_tokens_cpu is not None else x.shape[0]

273 - )

274 - orig_topk_weights_dtype = topk_weights.dtype

275 - alltoall_info, topk_ids, topk_weights, _ = (

276 - MnnvlMoe.mnnvl_moe_alltoallv_prepare_without_allgather(

277 - topk_ids,

278 - topk_weights,

279 - None,

280 - all2all_manager.prepare_workspace_tensor,

281 - max_num_token,

282 - ep_rank,

283 - ep_size,

284 - num_experts,

285 - num_experts,

286 - top_k,

287 - )

288 - )

289 - topk_weights = topk_weights.view(dtype=orig_topk_weights_dtype)

290 -

291 - if not use_deepseek_fp8_block_scale:

292 - x, x_sf = moe_kernel_quantize_input(

293 - x,

294 - gs,

295 - quant_config.quant_dtype,

296 - quant_config.per_act_token_quant,

297 - quant_config.block_shape,

298 - is_fp4_scale_swizzled=False, # delay swizzle to after comm

299 - )

300 - x = MnnvlMoe.mnnvl_moe_alltoallv(

301 - x,

302 - alltoall_info,

303 - all2all_manager.workspace_tensor,

304 - ep_rank,

305 - ep_size,

306 - )

307 -

308 - x_sf = MnnvlMoe.mnnvl_moe_alltoallv(

309 - x_sf,

310 - alltoall_info,

311 - all2all_manager.workspace_tensor,

312 - ep_rank,

313 - ep_size,

314 - )

315 - if quant_config.quant_dtype == "nvfp4":

316 - x_sf = nvfp4_block_scale_interleave(x_sf)

317 - else:

318 - # Block-scale path: pass activations through without quantization

319 - x_sf = None

320 - x = MnnvlMoe.mnnvl_moe_alltoallv(

321 - x,

322 - alltoall_info,

323 - all2all_manager.workspace_tensor,

324 - ep_rank,

325 - ep_size,

326 - )

327 - return alltoall_info, topk_ids, topk_weights, x, x_sf

328 -

329 -

330 -def flashinfer_alltoall_combine(

331 - all2all_manager: All2AllManagerBase,

332 - output: torch.Tensor,

333 - top_k: int,

334 - token_count: int,

335 - alltoall_info,

336 -):

337 - from flashinfer.comm.trtllm_alltoall import MnnvlMoe

338 -

339 - assert all2all_manager.ensure_alltoall_workspace_initialized(), (

340 - "FlashInfer AllToAll workspace not available"

341 - )

342 - return MnnvlMoe.mnnvl_moe_alltoallv_combine(

343 - output,

344 - alltoall_info,

345 - all2all_manager.workspace_tensor,

346 - ep_rank=all2all_manager.rank,

347 - ep_size=all2all_manager.world_size,

348 - top_k=top_k,

349 - token_count=token_count,

350 - )

351 -

352 -

353 -def create_flashinfer_prepare_finalize(

354 - use_dp: bool,

355 - use_nvfp4: bool = False,

356 - enable_alltoallv: bool = False,

357 - use_deepseek_fp8_block_scale: bool = False,

358 -) -> FlashInferCutlassMoEPrepareAndFinalize | MoEPrepareAndFinalizeNoEP:

359 - """Factory function to create the appropriate FlashInfer implementation."""

360 -

361 - if use_dp:

362 - if enable_alltoallv:

363 - assert use_nvfp4

364 - return FlashInferAllToAllMoEPrepareAndFinalize(use_dp)

365 - return FlashInferAllGatherMoEPrepareAndFinalize(

366 - use_dp=True,

367 - use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale,

368 - )

369 - else:

370 - # CUTLASS FP8 BLOCK and CUTLASS NVFP4 apply input quantization

371 - # in a single call with the MoE experts kernel.

372 - defer_input_quant = use_deepseek_fp8_block_scale or use_nvfp4

373 - return MoEPrepareAndFinalizeNoEP(defer_input_quant=defer_input_quant)

vllm/model_executor/layers/fused_moe/fused_batched_moe.py

@@ -533,7 +533,13 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

533 expert_map: torch.Tensor | None,

534 apply_router_weight_on_input: bool,

535 quant_config: FusedMoEQuantConfig,

536 + defer_input_quant: bool = False,

537 ) -> mk.PrepareResultType:

538 + if defer_input_quant:

539 + raise NotImplementedError(

540 + f"{self.__class__.__name__} does not support defer_input_quant=True. "

541 + "Please select an MoE kernel that accepts quantized inputs."

542 + )

543 assert a1.dim() == 2

544 assert topk_ids.dim() == 2

545 assert topk_ids.size(0) == a1.size(0)

vllm/model_executor/layers/fused_moe/fused_marlin_moe.py

+1 -1

@@ -593,7 +593,7 @@ class MarlinExpertsBase(mk.FusedMoEPermuteExpertsUnpermute):

593

594 @staticmethod

595 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:

596 - return True

596 + return not moe_parallel_config.use_fi_all2allv_kernels

597

598 @property

599 def quant_type_id(self) -> int:

vllm/model_executor/layers/fused_moe/fused_moe.py

+1 -1

@@ -1951,7 +1951,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

1951

1952 @staticmethod

1953 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:

1954 - return True

1954 + return not moe_parallel_config.use_fi_all2allv_kernels

1955

1956 def supports_chunking(self) -> bool:

1957 return True

vllm/model_executor/layers/fused_moe/fused_moe_method_base.py

+16

@@ -5,6 +5,7 @@ from abc import abstractmethod

6 import torch

8 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk

9 from vllm.logger import init_logger

10 from vllm.model_executor.layers.fused_moe.config import (

11 FusedMoEConfig,

@@ -26,6 +27,19 @@ class FusedMoEMethodBase(QuantizeMethodBase):

27 super().__init__()

28 self.moe: FusedMoEConfig = moe

29 self.moe_quant_config: FusedMoEQuantConfig | None = None

30 + self.moe_mk: mk.FusedMoEModularKernel | None = None

31 +

32 + @property

33 + def supports_internal_mk(self) -> bool:

34 + # NOTE(rob): temporary attribute to indicate support for

35 + # completed migration to the new internal MK interface.

36 + return self.moe_mk is not None

37 +

38 + @property

39 + def mk_owns_shared_expert(self) -> bool:

40 + # NOTE(rob): temporary attribute to indicate support for

41 + # completed migration to the new internal MK interface.

42 + return self.moe_mk is not None and self.moe_mk.shared_experts is not None

44 @abstractmethod

45 def create_weights(

@@ -91,6 +105,8 @@ class FusedMoEMethodBase(QuantizeMethodBase):

105

106 @property

107 def topk_indices_dtype(self) -> torch.dtype | None:

108 + if self.moe_mk is not None:

109 + return self.moe_mk.prepare_finalize.topk_indices_dtype()

110 return None

111

112 @property

vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py

+4 -7

@@ -30,11 +30,11 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):

30 ):

31 super().__init__(old_quant_method.moe)

32 self.moe_quant_config = old_quant_method.moe_quant_config

33 - self.fused_experts = experts

33 + self.moe_mk = experts

34 self.disable_expert_map = getattr(

35 old_quant_method,

36 "disable_expert_map",

37 - not self.fused_experts.supports_expert_map(),

37 + not self.moe_mk.supports_expert_map(),

38 )

39 self.old_quant_method = old_quant_method

40 assert not self.old_quant_method.is_monolithic

@@ -57,10 +57,6 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):

57 ),

58 )

60 - @property

61 - def topk_indices_dtype(self) -> torch.dtype | None:

62 - return self.fused_experts.prepare_finalize.topk_indices_dtype()

63 -

60 @property

61 def supports_eplb(self) -> bool:

62 return self.old_quant_method.supports_eplb

@@ -96,7 +92,8 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):

92 topk_weights: torch.Tensor,

93 topk_ids: torch.Tensor,

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

99 - return self.fused_experts(

95 + assert self.moe_mk is not None

96 + return self.moe_mk(

97 hidden_states=x,

98 w1=layer.w13_weight,

99 w2=layer.w2_weight,

vllm/model_executor/layers/fused_moe/layer.py

+13 -17

@@ -571,9 +571,6 @@ class FusedMoE(CustomOp):

571 device=vllm_config.device_config.device,

572 routing_method=self.routing_method_type,

573 )

574 - self.moe_config_use_flashinfer_cutlass_kernels = (

575 - self.moe_config.use_flashinfer_cutlass_kernels

576 - )

574 if self.use_mori_kernels:

575 assert self.rocm_aiter_fmoe_enabled, (

576 "Mori needs to be used with aiter fused_moe for now."

@@ -646,6 +643,11 @@ class FusedMoE(CustomOp):

643 # This is called after all weight loading and post-processing, so it

644 # should be safe to swap out the quant_method.

645 def maybe_init_modular_kernel(self) -> None:

646 + # NOTE(rob): WIP refactor. For quant methods that own the MK

647 + # we create the MK during process_weights_after_loading.

648 + if self.quant_method.supports_internal_mk or self.quant_method.is_monolithic:

649 + return None

650 +

651 self.ensure_moe_quant_config_init()

652 # routing_tables only needed for round-robin expert placement with

653 # DeepEP all2all backend.

@@ -728,14 +730,6 @@ class FusedMoE(CustomOp):

730 def use_mori_kernels(self):

731 return self.moe_parallel_config.use_mori_kernels

732

731 - @property

732 - def use_flashinfer_cutlass_kernels(self):

733 - return (

734 - self.moe_quant_config is not None

735 - and self.moe_quant_config.quant_dtype == "nvfp4"

736 - and self.moe_config_use_flashinfer_cutlass_kernels

737 - )

738 -

733 @property

734 def use_marlin_kernels(self):

735 return getattr(self.quant_method, "use_marlin", False)

@@ -746,7 +740,7 @@ class FusedMoE(CustomOp):

740 self.moe_parallel_config.use_pplx_kernels

741 or self.moe_parallel_config.use_deepep_ll_kernels

742 or self.moe_parallel_config.use_mori_kernels

749 - or (self.dp_size > 1 and self.use_flashinfer_cutlass_kernels)

743 + or self.moe_parallel_config.use_fi_all2allv_kernels

744 ) and envs.VLLM_ENABLE_MOE_DP_CHUNK

745

746 @property

@@ -1532,7 +1526,7 @@ class FusedMoE(CustomOp):

1526 assert self.quant_method is not None

1527 return (

1528 isinstance(self.quant_method, FusedMoEModularMethod)

1535 - and self.quant_method.fused_experts.output_is_reduced()

1529 + and self.quant_method.moe_mk.output_is_reduced() # type: ignore[union-attr]

1530 )

1531

1532 def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor):

@@ -1765,7 +1759,7 @@ class FusedMoE(CustomOp):

1759 self.ensure_dp_chunking_init()

1760

1761 has_separate_shared_experts = (

1768 - not isinstance(self.quant_method, FusedMoEModularMethod)

1762 + not self.quant_method.mk_owns_shared_expert

1763 and self.shared_experts is not None

1764 )

1765

@@ -1789,8 +1783,10 @@ class FusedMoE(CustomOp):

1783 hidden_states, router_logits, has_separate_shared_experts

1784 )

1785

1792 - do_naive_dispatch_combine: bool = self.dp_size > 1 and not isinstance(

1793 - self.quant_method, FusedMoEModularMethod

1786 + # NOTE(rob): once we finish migrating all the quant methods to use

1787 + # MKs, we can remove the naive dispatch/combine path from here.

1788 + do_naive_dispatch_combine = (

1789 + self.dp_size > 1 and not self.quant_method.supports_internal_mk

1790 )

1791

1792 ctx = get_forward_context()

@@ -1818,7 +1814,7 @@ class FusedMoE(CustomOp):

1814 else:

1815 hidden_states_to_dispatch = hidden_states

1816

1821 - dispatch_res = get_ep_group().dispatch(

1817 + dispatch_res = get_ep_group().dispatch_router_logits(

1818 hidden_states_to_dispatch,

1819 router_logits,

1820 self.is_sequence_parallel,

vllm/model_executor/layers/fused_moe/modular_kernel.py

+12 -4

@@ -180,6 +180,7 @@ class FusedMoEPrepareAndFinalize(ABC):

180 expert_map: torch.Tensor | None,

181 apply_router_weight_on_input: bool,

182 quant_config: FusedMoEQuantConfig,

183 + defer_input_quant: bool,

184 ) -> PrepareResultType:

185 """

186 Perform any quantization (and/or) dispatching needed for this kernel.

@@ -192,6 +193,9 @@ class FusedMoEPrepareAndFinalize(ABC):

193 - apply_router_weight_on_input: When True, apply the weights to the

194 activations, before quantization + dispatching.

195 - quant_config: Quantization info provided by the fused experts.

196 + - defer_input_quant: Runtime parameter indicating whether or not to

197 + defer input quantization to the FusedMoEPermuteExpertsUnpermute

198 + in cases where the compute kernel expects unquantized inputs

199

200 Returns a tuple of:

201 - quantized + dispatched a.

@@ -220,6 +224,7 @@ class FusedMoEPrepareAndFinalize(ABC):

224 expert_map: torch.Tensor | None,

225 apply_router_weight_on_input: bool,

226 quant_config: FusedMoEQuantConfig,

227 + defer_input_quant: bool,

228 ) -> tuple[Callable, ReceiverType] | ReceiverType:

229 """

230 Perform any quantization (and/or) dispatching needed for this kernel

@@ -235,6 +240,9 @@ class FusedMoEPrepareAndFinalize(ABC):

240 space to the local expert space of the expert parallel shard.

241 - apply_router_weight_on_input: When True, apply the weights to the

242 activations, before quantization + dispatching.

243 + - defer_input_quant: Runtime parameter indicating whether or not to

244 + defer input quantization to the FusedMoEPermuteExpertsUnpermute

245 + in cases where the compute kernel expects unquantized inputs

246

247 Returns a callback or a hook callback pair that when invoked waits for

248 results from other workers and has the same return signature as

@@ -407,10 +415,8 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

415 self.max_num_tokens = max_num_tokens

416 self.num_dispatchers = num_dispatchers

417

410 - @staticmethod

411 - def expects_unquantized_inputs(

412 - moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig

413 - ) -> bool:

418 + @property

419 + def expects_unquantized_inputs(self) -> bool:

420 """

421 Whether or not the PrepareFinalize should defer input quantization

422 in the prepare step. If True, then the Experts kernel will

@@ -1069,6 +1075,7 @@ class FusedMoEModularKernel(torch.nn.Module):

1075 expert_map,

1076 apply_router_weight_on_input,

1077 self.fused_experts.quant_config,

1078 + defer_input_quant=self.fused_experts.expects_unquantized_inputs,

1079 )

1080 else:

1081 # Overlap shared expert compute with all2all dispatch.

@@ -1081,6 +1088,7 @@ class FusedMoEModularKernel(torch.nn.Module):

1088 expert_map,

1089 apply_router_weight_on_input,

1090 self.fused_experts.quant_config,

1091 + defer_input_quant=self.fused_experts.expects_unquantized_inputs,

1092 )

1093

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

vllm/model_executor/layers/fused_moe/mori_prepare_finalize.py

@@ -58,6 +58,7 @@ class MoriPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

58 expert_map: torch.Tensor | None,

59 apply_router_weight_on_input: bool,

60 quant_config: FusedMoEQuantConfig,

61 + defer_input_quant: bool = False,

62 ) -> mk.PrepareResultType:

63 """

64 Returns a tuple of:

@@ -69,6 +70,11 @@ class MoriPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

70 - Optional dispatched expert topk IDs

71 - Optional dispatched expert topk weight

72 """

73 + if defer_input_quant:

74 + raise NotImplementedError(

75 + f"{self.__class__.__name__} does not support defer_input_quant=True. "

76 + "Please select an MoE kernel that accepts quantized inputs."

77 + )

78 assert not apply_router_weight_on_input, (

79 "mori does not support apply_router_weight_on_input=True now."

80 )

vllm/model_executor/layers/fused_moe/oracle/fp8.py

+31 -47

@@ -8,6 +8,9 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk

8 from vllm import envs

9 from vllm._aiter_ops import rocm_aiter_ops

10 from vllm.logger import init_logger

11 +from vllm.model_executor.layers.fused_moe.all2all_utils import (

12 + maybe_make_prepare_finalize,

13 +)

14 from vllm.model_executor.layers.fused_moe.config import (

15 FusedMoEConfig,

16 FusedMoEQuantConfig,

@@ -17,9 +20,6 @@ from vllm.model_executor.layers.fused_moe.config import (

20 from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe import (

21 is_supported_config_trtllm,

22 )

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

21 - MoEPrepareAndFinalizeNoEP,

22 -)

23 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

24 FlashinferMoeBackend,

25 get_flashinfer_moe_backend,

@@ -465,68 +465,52 @@ def make_fp8_moe_quant_config(

465 )

466

467

468 -def make_fp8_moe_kernel_for_mkm(

468 +def make_fp8_moe_kernel(

469 + moe_quant_config: FusedMoEQuantConfig,

470 moe_config: FusedMoEConfig,

470 - quant_config: FusedMoEQuantConfig,

471 experts_cls: type[mk.FusedMoEPermuteExpertsUnpermute],

472 - prepare_finalize: mk.FusedMoEPrepareAndFinalize,

473 -) -> mk.FusedMoEPermuteExpertsUnpermute:

472 + fp8_backend: Fp8MoeBackend,

473 + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

474 + shared_experts: torch.nn.Module | None = None,

475 +) -> tuple[mk.FusedMoEModularKernel, bool]:

476 + # Create Prepare/Finalize.

477 + prepare_finalize = maybe_make_prepare_finalize(

478 + moe=moe_config,

479 + quant_config=moe_quant_config,

480 + routing_tables=routing_tables,

481 + allow_new_interface=True,

482 + )

483 + assert prepare_finalize is not None

484 +

485 + logger.info_once("Using %s", prepare_finalize.__class__.__name__)

486 +

487 + # Create Experts.

488 if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts:

475 - max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank()

476 - assert max_num_tokens_per_rank is not None

489 + max_num_tokens = prepare_finalize.max_num_tokens_per_rank()

490 + assert max_num_tokens is not None

491 experts = experts_cls(

492 moe_config=moe_config,

479 - quant_config=quant_config,

480 - max_num_tokens=max_num_tokens_per_rank,

493 + quant_config=moe_quant_config,

494 + max_num_tokens=max_num_tokens,

495 num_dispatchers=prepare_finalize.num_dispatchers(),

496 )

497 else:

498 experts = experts_cls(

499 moe_config=moe_config,

486 - quant_config=quant_config,

500 + quant_config=moe_quant_config,

501 )

502

489 - logger.debug_once("Using %s", experts.__class__.__name__)

490 - return experts

491 -

492 -

493 -def make_fp8_moe_kernel(

494 - moe_quant_config: FusedMoEQuantConfig,

495 - moe_config: FusedMoEConfig,

496 - fp8_backend: Fp8MoeBackend,

497 - experts_cls: type[mk.FusedMoEPermuteExpertsUnpermute],

498 -) -> tuple[mk.FusedMoEModularKernel, bool]:

499 - # TODO(rob): unify after we merge tp and dp/ep.

500 - if (

501 - moe_config.moe_parallel_config.use_all2all_kernels

502 - and moe_config.moe_parallel_config.all2all_backend

503 - not in ["allgather_reducescatter", "naive"]

504 - ):

505 - raise ValueError(

506 - "Fp8 Oracle should not create non-naive A2A P/F. "

507 - "This should happen via the ModularKernelMethod."

508 - )

509 -

510 - # Create Prepare/Finalize.

511 - prepare_finalize = MoEPrepareAndFinalizeNoEP(

512 - defer_input_quant=experts_cls.expects_unquantized_inputs(

513 - moe_config, moe_quant_config

514 - ),

515 - )

516 -

517 - # Create Experts.

518 - experts = experts_cls(

519 - moe_config=moe_config,

520 - quant_config=moe_quant_config,

521 - )

522 -

503 # NOTE(rob): we only want the mk to control the shared_expert

504 # if using all2all (for SBO). bnell is making this explict in

505 # the new MoE runner class.

506 kernel = mk.FusedMoEModularKernel(

507 prepare_finalize,

508 experts,

529 - shared_experts=None,

509 + shared_experts=(

510 + shared_experts

511 + if moe_config.moe_parallel_config.use_all2all_kernels

512 + else None

513 + ),

514 moe_parallel_config=moe_config.moe_parallel_config,

515 )

516

vllm/model_executor/layers/fused_moe/oracle/nvfp4.py

+30 -46

@@ -7,6 +7,9 @@ import torch

7 import vllm.envs as envs

8 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

9 from vllm.logger import init_logger

10 +from vllm.model_executor.layers.fused_moe.all2all_utils import (

11 + maybe_make_prepare_finalize,

12 +)

13 from vllm.model_executor.layers.fused_moe.config import (

14 FusedMoEConfig,

15 FusedMoEQuantConfig,

@@ -14,9 +17,6 @@ from vllm.model_executor.layers.fused_moe.config import (

17 nvfp4_moe_quant_config,

18 nvfp4_w4a16_moe_quant_config,

19 )

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

18 - MoEPrepareAndFinalizeNoEP,

19 -)

20 from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (

21 is_supported_config_trtllm,

22 prepare_nvfp4_moe_layer_for_fi_or_cutlass,

@@ -391,67 +391,51 @@ def make_nvfp4_moe_quant_config(

391 )

392

393

394 -def make_nvfp4_moe_kernel_for_mkm(

394 +def make_nvfp4_moe_kernel(

395 + moe_quant_config: FusedMoEQuantConfig,

396 moe_config: FusedMoEConfig,

396 - quant_config: FusedMoEQuantConfig,

397 experts_cls: type[mk.FusedMoEPermuteExpertsUnpermute],

398 - prepare_finalize: mk.FusedMoEPrepareAndFinalize,

399 -) -> mk.FusedMoEPermuteExpertsUnpermute:

398 + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

399 + shared_experts: torch.nn.Module | None = None,

400 +) -> mk.FusedMoEModularKernel:

401 + # Create Prepare/Finalize.

402 + prepare_finalize = maybe_make_prepare_finalize(

403 + moe=moe_config,

404 + quant_config=moe_quant_config,

405 + routing_tables=routing_tables,

406 + allow_new_interface=True,

407 + )

408 + assert prepare_finalize is not None

409 +

410 + logger.info_once("Using %s", prepare_finalize.__class__.__name__)

411 +

412 + # Create Experts.

413 if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts:

401 - max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank()

402 - assert max_num_tokens_per_rank is not None

414 + max_num_tokens = prepare_finalize.max_num_tokens_per_rank()

415 + assert max_num_tokens is not None

416 experts = experts_cls(

417 moe_config=moe_config,

405 - quant_config=quant_config,

406 - max_num_tokens=max_num_tokens_per_rank,

418 + quant_config=moe_quant_config,

419 + max_num_tokens=max_num_tokens,

420 num_dispatchers=prepare_finalize.num_dispatchers(),

421 )

422 else:

423 experts = experts_cls(

424 moe_config=moe_config,

412 - quant_config=quant_config,

425 + quant_config=moe_quant_config,

426 )

427

415 - logger.debug_once("Using %s", experts.__class__.__name__)

416 - return experts

417 -

418 -

419 -def make_nvfp4_moe_kernel(

420 - moe_quant_config: FusedMoEQuantConfig,

421 - moe_config: FusedMoEConfig,

422 - experts_cls: type[mk.FusedMoEPermuteExpertsUnpermute],

423 -) -> mk.FusedMoEModularKernel:

424 - # TODO(rob): unify after we merge tp and dp/ep.

425 - if (

426 - moe_config.moe_parallel_config.use_all2all_kernels

427 - and moe_config.moe_parallel_config.all2all_backend

428 - not in ["allgather_reducescatter", "naive"]

429 - ):

430 - raise ValueError(

431 - "NvFP4 Oracle should not create non-naive A2A P/F. "

432 - "This should happen via the ModularKernelMethod."

433 - )

434 -

435 - # Create Prepare/Finalize.

436 - prepare_finalize = MoEPrepareAndFinalizeNoEP(

437 - defer_input_quant=experts_cls.expects_unquantized_inputs(

438 - moe_config, moe_quant_config

439 - ),

440 - )

441 -

442 - # Create Experts.

443 - experts = experts_cls(

444 - moe_config=moe_config,

445 - quant_config=moe_quant_config,

446 - )

447 -

428 # NOTE(rob): we only want the mk to control the shared_expert

429 # if using all2all (for SBO). bnell is making this explict in

430 # the new MoE runner class.

431 kernel = mk.FusedMoEModularKernel(

432 prepare_finalize,

433 experts,

454 - shared_experts=None,

434 + shared_experts=(

435 + shared_experts

436 + if moe_config.moe_parallel_config.use_all2all_kernels

437 + else None

438 + ),

439 moe_parallel_config=moe_config.moe_parallel_config,

440 )

441

vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py

@@ -106,7 +106,14 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

106 expert_map: torch.Tensor | None,

107 apply_router_weight_on_input: bool,

108 quant_config: FusedMoEQuantConfig,

109 + defer_input_quant: bool = False,

110 ) -> tuple[Callable, mk.ReceiverType]:

111 + if defer_input_quant:

112 + raise NotImplementedError(

113 + f"{self.__class__.__name__} does not support defer_input_quant=True. "

114 + "Please select an MoE kernel that accepts quantized inputs."

115 + )

116 +

117 num_tokens = a1.size(0) # M

118 hidden_dim = a1.size(-1) # K

119

@@ -274,6 +281,7 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

281 expert_map: torch.Tensor | None,

282 apply_router_weight_on_input: bool,

283 quant_config: FusedMoEQuantConfig,

284 + defer_input_quant: bool = False,

285 ) -> mk.PrepareResultType:

286 hook, receiver = self.prepare_async(

287 a1,

@@ -283,6 +291,7 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

291 expert_map,

292 apply_router_weight_on_input,

293 quant_config,

294 + defer_input_quant=defer_input_quant,

295 )

296 hook()

297 return receiver()

vllm/model_executor/layers/fused_moe/prepare_finalize.py

+125 -5

@@ -4,18 +4,25 @@

4 import torch

6 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

7 +from vllm.distributed import get_ep_group

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

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

10 TopKWeightAndReduceContiguous,

11 TopKWeightAndReduceDelegate,

12 )

13 from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input

14 +from vllm.utils.flashinfer import nvfp4_block_scale_interleave

15 -class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):

16 - def __init__(self, defer_input_quant: bool = False) -> None:

17 +class MoEPrepareAndFinalizeNaiveEP(mk.FusedMoEPrepareAndFinalize):

18 + def __init__(

19 + self,

20 + is_sequence_parallel: bool = False,

21 + num_dispatchers: int = 1,

22 + ) -> None:

23 super().__init__()

18 - self.defer_input_quant = defer_input_quant

24 + self.is_sequence_parallel = is_sequence_parallel

25 + self._num_dispatchers = num_dispatchers

27 @property

28 def activation_format(self) -> mk.FusedMoEActivationFormat:

@@ -27,6 +34,113 @@ class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):

34 def topk_indices_dtype(self) -> torch.dtype | None:

35 return None

37 + def num_dispatchers(self) -> int:

38 + return self._num_dispatchers

39 +

40 + def output_is_reduced(self) -> bool:

41 + return False

42 +

43 + def prepare(

44 + self,

45 + a1: torch.Tensor,

46 + topk_weights: torch.Tensor,

47 + topk_ids: torch.Tensor,

48 + num_experts: int,

49 + expert_map: torch.Tensor | None,

50 + apply_router_weight_on_input: bool,

51 + quant_config: FusedMoEQuantConfig,

52 + defer_input_quant: bool = False,

53 + ) -> mk.PrepareResultType:

54 + if apply_router_weight_on_input:

55 + topk = topk_ids.size(1)

56 + assert topk == 1, (

57 + "apply_router_weight_on_input is only implemented for topk=1"

58 + )

59 + # Note: do not use inplace for shared experts overlap

60 + a1 = a1 * topk_weights.to(a1.dtype)

61 +

62 + # Defer input quantization to the MoE kernel.

63 + use_nvfp4 = quant_config.use_nvfp4_w4a4

64 + if defer_input_quant:

65 + a1q = a1

66 + a1q_scale = None

67 + else:

68 + a1q, a1q_scale = moe_kernel_quantize_input(

69 + a1,

70 + quant_config.a1_gscale if use_nvfp4 else quant_config.a1_scale,

71 + quant_config.quant_dtype,

72 + quant_config.per_act_token_quant,

73 + quant_config.block_shape,

74 + # NOTE: swizzling pads the scales to multiple of 128

75 + # which makes the scales tensor different shape than

76 + # the hidden states, breaking the A2A kernel. So, we

77 + # delay the swizzling until after the A2A.

78 + is_fp4_scale_swizzled=False,

79 + )

80 +

81 + # Skip gathering scales if we have static quantization

82 + # (the scale is a scalar, replicated on all ranks) or

83 + # if quantization is deferred.

84 + skip_gather_scales = a1q_scale is None or a1q_scale.ndim == 0

85 + scales = None if skip_gather_scales else [a1q_scale]

86 +

87 + res = get_ep_group().dispatch(

88 + a1q,

89 + topk_weights,

90 + topk_ids,

91 + is_sequence_parallel=self.is_sequence_parallel,

92 + extra_tensors=scales,

93 + )

94 + if skip_gather_scales:

95 + a1q, topk_weights, topk_ids = res

96 + else:

97 + a1q, topk_weights, topk_ids, scales = res

98 + assert scales is not None and len(scales) == 1

99 + a1q_scale = scales[0]

100 + if quant_config.quant_dtype == "nvfp4":

101 + assert a1q_scale is not None

102 + if a1q_scale.element_size() == 1:

103 + a1q_scale = a1q_scale.view(torch.uint8)

104 + a1q_scale = nvfp4_block_scale_interleave(a1q_scale)

105 +

106 + return a1q, a1q_scale, None, topk_ids, topk_weights

107 +

108 + def finalize(

109 + self,

110 + output: torch.Tensor,

111 + fused_expert_output: torch.Tensor,

112 + topk_weights: torch.Tensor,

113 + topk_ids: torch.Tensor,

114 + apply_router_weight_on_input: bool,

115 + weight_and_reduce_impl: mk.TopKWeightAndReduce,

116 + ) -> None:

117 + if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):

118 + weight_and_reduce_impl = TopKWeightAndReduceContiguous()

119 +

120 + out = weight_and_reduce_impl.apply(

121 + output=None,

122 + fused_expert_output=fused_expert_output,

123 + topk_weights=topk_weights,

124 + topk_ids=topk_ids,

125 + apply_router_weight_on_input=apply_router_weight_on_input,

126 + )

127 +

128 + output.copy_(

129 + get_ep_group().combine(out, is_sequence_parallel=self.is_sequence_parallel)

130 + )

131 +

132 +

133 +class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):

134 + @property

135 + def activation_format(self) -> mk.FusedMoEActivationFormat:

136 + return mk.FusedMoEActivationFormat.Standard

137 +

138 + def max_num_tokens_per_rank(self) -> int | None:

139 + return None

140 +

141 + def topk_indices_dtype(self) -> torch.dtype | None:

142 + return None

143 +

144 def num_dispatchers(self) -> int:

145 return 1

146

@@ -42,6 +156,7 @@ class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):

156 expert_map: torch.Tensor | None,

157 apply_router_weight_on_input: bool,

158 quant_config: FusedMoEQuantConfig,

159 + defer_input_quant: bool = False,

160 ) -> mk.PrepareResultType:

161 if apply_router_weight_on_input:

162 topk = topk_ids.size(1)

@@ -54,12 +169,17 @@ class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):

169

170 # Defer input quant to moe kernel for backends (e.g. AITER, FI)

171 # which use a single kernel call for quant + experts.

57 - if self.defer_input_quant:

172 + if defer_input_quant:

173 return a1, None, None, None, None

174

175 + input_sf = (

176 + quant_config.a1_gscale

177 + if quant_config.use_nvfp4_w4a4

178 + else quant_config.a1_scale

179 + )

180 a1q, a1q_scale = moe_kernel_quantize_input(

181 a1,

62 - quant_config.a1_scale,

182 + input_sf,

183 quant_config.quant_dtype,

184 quant_config.per_act_token_quant,

185 quant_config.block_shape,

vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py

+5 -8

@@ -287,17 +287,14 @@ def rocm_aiter_fused_experts(

287

288

289 class AiterExperts(mk.FusedMoEPermuteExpertsUnpermute):

290 + @property

291 + def expects_unquantized_inputs(self) -> bool:

292 + return True

293 +

294 @staticmethod

295 def activation_format() -> mk.FusedMoEActivationFormat:

296 return mk.FusedMoEActivationFormat.Standard

297

294 - @staticmethod

295 - def expects_unquantized_inputs(

296 - fused_moe_config: mk.FusedMoEConfig, quant_config: FusedMoEQuantConfig

297 - ) -> bool:

298 - # AITER fused MoE kernels handle input quantization internally.

299 - return True

300 -

298 @staticmethod

299 def _supports_current_device() -> bool:

300 return rocm_aiter_ops.is_fused_moe_enabled()

@@ -329,7 +326,7 @@ class AiterExperts(mk.FusedMoEPermuteExpertsUnpermute):

326

327 @staticmethod

328 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:

332 - return True

329 + return not moe_parallel_config.use_fi_all2allv_kernels

330

331 def supports_expert_map(self):

332 return True

vllm/model_executor/layers/fused_moe/shared_fused_moe.py

+1 -1

@@ -37,7 +37,7 @@ class SharedFusedMoE(FusedMoE):

37 use_overlapped

38 and not (

39 (self.enable_eplb and backend != "allgather_reducescatter")

40 - or (self.moe_config.use_flashinfer_cutlass_kernels and self.dp_size > 1)

40 + or self.moe_parallel_config.use_fi_all2allv_kernels

41 )

42 and self._shared_experts is not None

43 )

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py

+35 -85

@@ -43,7 +43,6 @@ from vllm.model_executor.layers.fused_moe.oracle.fp8 import (

43 Fp8MoeBackend,

44 convert_to_fp8_moe_kernel_format,

45 make_fp8_moe_kernel,

46 - make_fp8_moe_kernel_for_mkm,

46 make_fp8_moe_quant_config,

47 select_fp8_moe_backend,

48 )

@@ -53,7 +52,6 @@ from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (

52 is_global_sf_supported_for_nvfp4_backend,

53 make_mxfp4_moe_quant_config,

54 make_nvfp4_moe_kernel,

56 - make_nvfp4_moe_kernel_for_mkm,

55 make_nvfp4_moe_quant_config,

56 select_nvfp4_moe_backend,

57 )

@@ -67,7 +65,6 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (

65 )

66 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

67 apply_fi_trtllm_fp8_per_tensor_moe,

70 - build_flashinfer_fp8_cutlass_moe_prepare_finalize,

68 )

69 from vllm.model_executor.layers.quantization.utils.fp8_utils import (

70 process_fp8_input_tensor_strategy_moe,

@@ -243,7 +240,6 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):

240 self.group_size = 32

241 self.mxfp4_backend = NvFp4MoeBackend.MARLIN

242 self.experts_cls = MarlinExperts

246 - self.kernel: mk.FusedMoEModularKernel | None = None

243

244 def create_weights(

245 self,

@@ -320,7 +316,7 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):

316 w13_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale

317 )

318

323 - def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

319 + def process_weights_after_loading(self, layer: FusedMoE) -> None:

320 layer.w13_weight = torch.nn.Parameter(

321 layer.w13_weight_packed.data, requires_grad=False

322 )

@@ -335,10 +331,12 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):

331

332 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

333 if self.moe_quant_config is not None:

338 - self.kernel = make_nvfp4_moe_kernel(

334 + self.moe_mk = make_nvfp4_moe_kernel(

335 moe_quant_config=self.moe_quant_config,

336 moe_config=self.moe,

337 experts_cls=self.experts_cls,

338 + shared_experts=layer.shared_experts,

339 + routing_tables=layer._maybe_init_expert_routing_tables(),

340 )

341

342 def apply(

@@ -348,8 +346,8 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):

346 topk_weights: torch.Tensor,

347 topk_ids: torch.Tensor,

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

351 - assert self.kernel is not None

352 - return self.kernel(

349 + assert self.moe_mk is not None

350 + return self.moe_mk(

351 x,

352 layer.w13_weight,

353 layer.w2_weight,

@@ -380,19 +378,10 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

378 activation_key=None if use_a16 else kNvfp4Dynamic,

379 )

380

383 - # Delay creation of the kernel until after process-weights.

384 - self.kernel: mk.FusedMoEModularKernel | None = None

385 -

381 self.use_global_sf = is_global_sf_supported_for_nvfp4_backend(

382 self.nvfp4_backend

383 )

384

390 - @property

391 - def topk_indices_dtype(self) -> torch.dtype | None:

392 - if self.kernel is not None:

393 - return self.kernel.prepare_finalize.topk_indices_dtype()

394 - return None

395 -

385 def create_weights(

386 self,

387 layer: torch.nn.Module,

@@ -506,7 +495,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

495 )

496 set_weight_attrs(w2_input_scale, extra_weight_attrs)

497

509 - def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

498 + def process_weights_after_loading(self, layer: FusedMoE) -> None:

499 """

500 Convert NVFP4 MoE weights into kernel format and setup the kernel.

501 """

@@ -572,48 +561,33 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

561 # TODO(rob): unify these so FP8MoEMethod owns the ModularKernel

562 # in both cases.

563 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

575 - if self.moe_quant_config and (

576 - (not self.moe.moe_parallel_config.use_all2all_kernels)

577 - or self.moe.moe_parallel_config.use_naive_all2all_kernels

578 - ):

564 + if self.moe_quant_config:

565 assert self.experts_cls is not None

580 - self.kernel = make_nvfp4_moe_kernel(

566 + self.moe_mk = make_nvfp4_moe_kernel(

567 moe_quant_config=self.moe_quant_config,

568 moe_config=self.moe,

569 experts_cls=self.experts_cls,

570 + shared_experts=layer.shared_experts,

571 + routing_tables=layer._maybe_init_expert_routing_tables(),

572 )

573

574 def maybe_make_prepare_finalize(

575 self,

576 routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

577 ) -> mk.FusedMoEPrepareAndFinalize | None:

590 - if self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM:

591 - return None

592 - elif self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_CUTLASS:

593 - # For no-EP case, don't use the MKM framework.

594 - if not self.moe.moe_parallel_config.use_all2all_kernels:

595 - return None

596 -

597 - prepare_finalize = build_flashinfer_fp8_cutlass_moe_prepare_finalize(

598 - self.moe,

599 - use_deepseek_fp8_block_scale=False,

600 - )

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

602 - return prepare_finalize

603 - return super().maybe_make_prepare_finalize(routing_tables)

578 + raise ValueError(

579 + f"{self.__class__.__name__} uses the new modular kernel initialization "

580 + "logic. This function should not be called."

581 + )

582

583 def select_gemm_impl(

584 self,

585 prepare_finalize: mk.FusedMoEPrepareAndFinalize,

586 layer: torch.nn.Module,

587 ) -> mk.FusedMoEPermuteExpertsUnpermute:

610 - assert self.moe_quant_config is not None

611 - assert self.experts_cls is not None

612 - return make_nvfp4_moe_kernel_for_mkm(

613 - moe_config=self.moe,

614 - quant_config=self.moe_quant_config,

615 - experts_cls=self.experts_cls,

616 - prepare_finalize=prepare_finalize,

588 + raise ValueError(

589 + f"{self.__class__.__name__} uses the new modular kernel initialization "

590 + "logic. This function should not be called."

591 )

592

593 def get_fused_moe_quant_config(

@@ -684,8 +658,8 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

658 global_num_experts=layer.global_num_experts,

659 )

660 else:

687 - assert self.kernel is not None

688 - return self.kernel(

661 + assert self.moe_mk is not None

662 + return self.moe_mk(

663 x,

664 layer.w13_weight,

665 layer.w2_weight,

@@ -759,15 +733,6 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

733 allow_vllm_cutlass=True,

734 )

735

762 - # Delay creation of the kernel until after process-weights.

763 - self.kernel: mk.FusedMoEModularKernel | None = None

764 -

765 - @property

766 - def topk_indices_dtype(self) -> torch.dtype | None:

767 - if self.kernel is not None:

768 - return self.kernel.prepare_finalize.topk_indices_dtype()

769 - return None

770 -

736 def create_weights(

737 self,

738 layer: torch.nn.Module,

@@ -927,7 +892,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

892 layer.w13_input_scale = None

893 layer.w2_input_scale = None

894

930 - def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

895 + def process_weights_after_loading(self, layer: FusedMoE) -> None:

896 # Allow for accessing weights and scales in standard way.

897 w13 = layer.w13_weight

898 w2 = layer.w2_weight

@@ -989,49 +954,34 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

954 # TODO(rob): unify these so FP8MoEMethod owns the ModularKernel

955 # in both cases.

956 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

992 - if self.moe_quant_config and (

993 - (not self.moe.moe_parallel_config.use_all2all_kernels)

994 - or self.moe.moe_parallel_config.use_naive_all2all_kernels

995 - ):

957 + if self.moe_quant_config:

958 assert self.experts_cls is not None

997 - self.kernel, self.use_inplace = make_fp8_moe_kernel(

959 + self.moe_mk, self.use_inplace = make_fp8_moe_kernel(

960 moe_quant_config=self.moe_quant_config,

961 moe_config=self.moe,

962 fp8_backend=self.fp8_backend,

963 experts_cls=self.experts_cls,

964 + routing_tables=layer._maybe_init_expert_routing_tables(),

965 + shared_experts=layer.shared_experts,

966 )

967

968 def maybe_make_prepare_finalize(

969 self,

970 routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

971 ) -> mk.FusedMoEPrepareAndFinalize | None:

1008 - if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:

1009 - return None

1010 - elif self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:

1011 - # For no-EP case, don't use the MKM framework.

1012 - if not self.moe.moe_parallel_config.use_all2all_kernels:

1013 - return None

1014 -

1015 - prepare_finalize = build_flashinfer_fp8_cutlass_moe_prepare_finalize(

1016 - self.moe,

1017 - use_deepseek_fp8_block_scale=self.block_quant,

1018 - )

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

1020 - return prepare_finalize

1021 - return super().maybe_make_prepare_finalize(routing_tables)

972 + raise ValueError(

973 + f"{self.__class__.__name__} uses the new modular kernel initialization "

974 + "logic. This function should not be called."

975 + )

976

977 def select_gemm_impl(

978 self,

979 prepare_finalize: mk.FusedMoEPrepareAndFinalize,

980 layer: torch.nn.Module,

1027 - ) -> FusedMoEPermuteExpertsUnpermute:

1028 - assert self.moe_quant_config is not None

1029 - assert self.experts_cls is not None

1030 - return make_fp8_moe_kernel_for_mkm(

1031 - moe_config=self.moe,

1032 - quant_config=self.moe_quant_config,

1033 - experts_cls=self.experts_cls,

1034 - prepare_finalize=prepare_finalize,

981 + ) -> mk.FusedMoEPermuteExpertsUnpermute:

982 + raise ValueError(

983 + f"{self.__class__.__name__} uses the new modular kernel initialization "

984 + "logic. This function should not be called."

985 )

986

987 def get_fused_moe_quant_config(

@@ -1120,8 +1070,8 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

1070 topk_ids: torch.Tensor,

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

1072 assert not self.is_monolithic

1123 - assert self.kernel is not None

1124 - return self.kernel(

1073 + assert self.moe_mk is not None

1074 + return self.moe_mk(

1075 x,

1076 layer.w13_weight,

1077 layer.w2_weight,

vllm/model_executor/layers/quantization/fp8.py

+14 -40

@@ -33,7 +33,6 @@ from vllm.model_executor.layers.fused_moe.oracle.fp8 import (

33 Fp8MoeBackend,

34 convert_to_fp8_moe_kernel_format,

35 make_fp8_moe_kernel,

36 - make_fp8_moe_kernel_for_mkm,

36 make_fp8_moe_quant_config,

37 select_fp8_moe_backend,

38 )

@@ -53,7 +52,6 @@ from vllm.model_executor.layers.quantization.kernels.scaled_mm import (

52 from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod

53 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

54 apply_fi_trtllm_fp8_per_tensor_moe,

56 - build_flashinfer_fp8_cutlass_moe_prepare_finalize,

55 )

56 from vllm.model_executor.layers.quantization.utils.fp8_utils import (

57 W8A8BlockFp8LinearOp,

@@ -679,15 +677,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):

677 allow_vllm_cutlass=False,

678 )

679

682 - # Delay creation of the kernel until after process-weights.

683 - self.kernel: mk.FusedMoEModularKernel | None = None

684 -

685 - @property

686 - def topk_indices_dtype(self) -> torch.dtype | None:

687 - if self.kernel is not None:

688 - return self.kernel.prepare_finalize.topk_indices_dtype()

689 - return None

690 -

680 def create_weights(

681 self,

682 layer: Module,

@@ -813,7 +802,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

802

803 def _setup_kernel(

804 self,

816 - layer: Module,

805 + layer: FusedMoE,

806 w13: torch.Tensor,

807 w2: torch.Tensor,

808 w13_scale: torch.Tensor,

@@ -845,16 +834,15 @@ class Fp8MoEMethod(FusedMoEMethodBase):

834 # TODO(rob): unify these so FP8MoEMethod owns the ModularKernel

835 # in both cases.

836 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

848 - if self.moe_quant_config and (

849 - (not self.moe.moe_parallel_config.use_all2all_kernels)

850 - or self.moe.moe_parallel_config.use_naive_all2all_kernels

851 - ):

837 + if self.moe_quant_config:

838 assert self.experts_cls is not None

853 - self.kernel, self.use_inplace = make_fp8_moe_kernel(

839 + self.moe_mk, self.use_inplace = make_fp8_moe_kernel(

840 moe_quant_config=self.moe_quant_config,

841 moe_config=self.moe,

842 fp8_backend=self.fp8_backend,

843 experts_cls=self.experts_cls,

844 + routing_tables=layer._maybe_init_expert_routing_tables(),

845 + shared_experts=layer.shared_experts,

846 )

847

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

@@ -909,33 +897,19 @@ class Fp8MoEMethod(FusedMoEMethodBase):

897 self,

898 routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

899 ) -> mk.FusedMoEPrepareAndFinalize | None:

912 - if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:

913 - return None

914 - elif self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:

915 - # For no-EP case, don't use the MKM framework.

916 - if not self.moe.moe_parallel_config.use_all2all_kernels:

917 - return None

918 -

919 - prepare_finalize = build_flashinfer_fp8_cutlass_moe_prepare_finalize(

920 - self.moe,

921 - use_deepseek_fp8_block_scale=self.block_quant,

922 - )

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

924 - return prepare_finalize

925 - return super().maybe_make_prepare_finalize(routing_tables)

900 + raise ValueError(

901 + f"{self.__class__.__name__} uses the new modular kernel initialization "

902 + "logic. This function should not be called."

903 + )

904

905 def select_gemm_impl(

906 self,

907 prepare_finalize: FusedMoEPrepareAndFinalize,

908 layer: torch.nn.Module,

909 ) -> FusedMoEPermuteExpertsUnpermute:

932 - assert self.moe_quant_config is not None

933 - assert self.experts_cls is not None

934 - return make_fp8_moe_kernel_for_mkm(

935 - moe_config=self.moe,

936 - quant_config=self.moe_quant_config,

937 - experts_cls=self.experts_cls,

938 - prepare_finalize=prepare_finalize,

910 + raise ValueError(

911 + f"{self.__class__.__name__} uses the new modular kernel initialization "

912 + "logic. This function should not be called."

913 )

914

915 def get_fused_moe_quant_config(

@@ -1037,9 +1011,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1011 topk_weights: torch.Tensor,

1012 topk_ids: torch.Tensor,

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

1040 - assert self.kernel is not None

1014 + assert self.moe_mk is not None

1015 assert not self.is_monolithic

1042 - return self.kernel(

1016 + return self.moe_mk(

1017 x,

1018 layer.w13_weight,

1019 layer.w2_weight,

vllm/model_executor/layers/quantization/modelopt.py

+27 -77

@@ -26,7 +26,6 @@ from vllm.model_executor.layers.fused_moe.oracle.fp8 import (

26 Fp8MoeBackend,

27 convert_to_fp8_moe_kernel_format,

28 make_fp8_moe_kernel,

29 - make_fp8_moe_kernel_for_mkm,

29 make_fp8_moe_quant_config,

30 select_fp8_moe_backend,

31 )

@@ -35,7 +34,6 @@ from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (

34 convert_to_nvfp4_moe_kernel_format,

35 is_global_sf_supported_for_nvfp4_backend,

36 make_nvfp4_moe_kernel,

38 - make_nvfp4_moe_kernel_for_mkm,

37 make_nvfp4_moe_quant_config,

38 select_nvfp4_moe_backend,

39 )

@@ -54,13 +52,11 @@ from vllm.model_executor.layers.quantization.kernels.scaled_mm import (

52 )

53 from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod

54 from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (

57 - build_flashinfer_fp4_cutlass_moe_prepare_finalize,

55 flashinfer_trtllm_fp4_moe,

56 flashinfer_trtllm_fp4_routed_moe,

57 )

58 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

59 apply_fi_trtllm_fp8_per_tensor_moe,

63 - build_flashinfer_fp8_cutlass_moe_prepare_finalize,

60 )

61 from vllm.model_executor.layers.quantization.utils.fp8_utils import (

62 W8A8BlockFp8LinearOp,

@@ -739,47 +735,23 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

735 activation_key=kFp8StaticTensorSym,

736 )

737

742 - # Delay creation of the kernel until after process-weights.

743 - self.kernel: mk.FusedMoEModularKernel | None = None

744 -

745 - @property

746 - def topk_indices_dtype(self) -> torch.dtype | None:

747 - if self.kernel is not None:

748 - return self.kernel.prepare_finalize.topk_indices_dtype()

749 - return None

750 -

738 def maybe_make_prepare_finalize(

739 self,

740 routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

741 ) -> mk.FusedMoEPrepareAndFinalize | None:

755 - # TRT LLM not supported with all2all yet.

756 - if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:

757 - return None

758 - elif self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:

759 - # For no-EP case, don't use the MKM framework.

760 - if not self.moe.moe_parallel_config.use_all2all_kernels:

761 - return None

762 -

763 - prepare_finalize = build_flashinfer_fp8_cutlass_moe_prepare_finalize(

764 - self.moe,

765 - use_deepseek_fp8_block_scale=False,

766 - )

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

768 - return prepare_finalize

769 - return super().maybe_make_prepare_finalize(routing_tables)

742 + raise ValueError(

743 + f"{self.__class__.__name__} uses the new modular kernel initialization "

744 + "logic. This function should not be called."

745 + )

746

747 def select_gemm_impl(

748 self,

749 prepare_finalize: mk.FusedMoEPrepareAndFinalize,

750 layer: torch.nn.Module,

751 ) -> mk.FusedMoEPermuteExpertsUnpermute:

776 - assert self.moe_quant_config is not None

777 - assert self.experts_cls is not None

778 - return make_fp8_moe_kernel_for_mkm(

779 - moe_config=self.moe,

780 - quant_config=self.moe_quant_config,

781 - experts_cls=self.experts_cls,

782 - prepare_finalize=prepare_finalize,

752 + raise ValueError(

753 + f"{self.__class__.__name__} uses the new modular kernel initialization "

754 + "logic. This function should not be called."

755 )

756

757 def create_weights(

@@ -863,7 +835,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

835

836 def _setup_kernel(

837 self,

866 - layer: torch.nn.Module,

838 + layer: FusedMoE,

839 w13: torch.Tensor,

840 w2: torch.Tensor,

841 w13_scale: torch.Tensor,

@@ -893,11 +865,13 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

865 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

866 if self.moe_quant_config:

867 assert self.experts_cls is not None

896 - self.kernel, self.use_inplace = make_fp8_moe_kernel(

868 + self.moe_mk, self.use_inplace = make_fp8_moe_kernel(

869 moe_quant_config=self.moe_quant_config,

870 moe_config=self.moe,

871 fp8_backend=self.fp8_backend,

872 experts_cls=self.experts_cls,

873 + routing_tables=layer._maybe_init_expert_routing_tables(),

874 + shared_experts=layer.shared_experts,

875 )

876

877 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

@@ -998,8 +972,8 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

972 f"but got {layer.activation}"

973 )

974

1001 - assert self.kernel is not None

1002 - return self.kernel(

975 + assert self.moe_mk is not None

976 + return self.moe_mk(

977 x,

978 layer.w13_weight,

979 layer.w2_weight,

@@ -1379,50 +1353,27 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1353 activation_key=kNvfp4Dynamic,

1354 )

1355

1382 - # Delay creation of the kernel until after process-weights.

1383 - self.kernel: mk.FusedMoEModularKernel | None = None

1384 -

1356 self.use_global_sf = is_global_sf_supported_for_nvfp4_backend(

1357 self.nvfp4_backend

1358 )

1359

1389 - @property

1390 - def topk_indices_dtype(self) -> torch.dtype | None:

1391 - if self.kernel is not None:

1392 - return self.kernel.prepare_finalize.topk_indices_dtype()

1393 - return None

1394 -

1360 def maybe_make_prepare_finalize(

1361 self,

1362 routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

1363 ) -> mk.FusedMoEPrepareAndFinalize | None:

1399 - if self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM:

1400 - return None

1401 - elif self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_CUTLASS:

1402 - # For no-EP case, don't use the MKM framework.

1403 - if not self.moe.moe_parallel_config.use_all2all_kernels:

1404 - return None

1405 - # For now, fp4 moe only works with the flashinfer dispatcher.

1406 - prepare_finalize = build_flashinfer_fp4_cutlass_moe_prepare_finalize(

1407 - self.moe

1408 - )

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

1410 - return prepare_finalize

1411 - else:

1412 - return super().maybe_make_prepare_finalize(routing_tables)

1364 + raise ValueError(

1365 + f"{self.__class__.__name__} uses the new modular kernel initialization "

1366 + "logic. This function should not be called."

1367 + )

1368

1369 def select_gemm_impl(

1370 self,

1371 prepare_finalize: mk.FusedMoEPrepareAndFinalize,

1372 layer: torch.nn.Module,

1373 ) -> mk.FusedMoEPermuteExpertsUnpermute:

1419 - assert self.moe_quant_config is not None

1420 - assert self.experts_cls is not None

1421 - return make_nvfp4_moe_kernel_for_mkm(

1422 - moe_config=self.moe,

1423 - quant_config=self.moe_quant_config,

1424 - experts_cls=self.experts_cls,

1425 - prepare_finalize=prepare_finalize,

1374 + raise ValueError(

1375 + f"{self.__class__.__name__} uses the new modular kernel initialization "

1376 + "logic. This function should not be called."

1377 )

1378

1379 def uses_weight_scale_2_pattern(self) -> bool:

@@ -1547,7 +1498,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1498 )

1499 layer.register_parameter("w2_input_scale", w2_input_scale)

1500

1550 - def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

1501 + def process_weights_after_loading(self, layer: FusedMoE) -> None:

1502 """

1503 Convert NVFP4 MoE weights into kernel format and setup the kernel.

1504 """

@@ -1599,15 +1550,14 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1550 # TODO(rob): unify these so FP8MoEMethod owns the ModularKernel

1551 # in both cases.

1552 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

1602 - if self.moe_quant_config and (

1603 - (not self.moe.moe_parallel_config.use_all2all_kernels)

1604 - or self.moe.moe_parallel_config.use_naive_all2all_kernels

1605 - ):

1553 + if self.moe_quant_config:

1554 assert self.experts_cls is not None

1607 - self.kernel = make_nvfp4_moe_kernel(

1555 + self.moe_mk = make_nvfp4_moe_kernel(

1556 moe_quant_config=self.moe_quant_config,

1557 moe_config=self.moe,

1558 experts_cls=self.experts_cls,

1559 + shared_experts=layer.shared_experts,

1560 + routing_tables=layer._maybe_init_expert_routing_tables(),

1561 )

1562

1563 @property

@@ -1708,8 +1658,8 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1658 global_num_experts=layer.global_num_experts,

1659 )

1660 else:

1711 - assert self.kernel is not None

1712 - return self.kernel(

1661 + assert self.moe_mk is not None

1662 + return self.moe_mk(

1663 x,

1664 layer.w13_weight,

1665 layer.w2_weight,

vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py

-15

@@ -15,9 +15,6 @@ from vllm.model_executor.layers.fused_moe.config import (

15 FusedMoEParallelConfig,

16 RoutingMethodType,

17 )

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

19 - create_flashinfer_prepare_finalize,

20 -)

18 from vllm.model_executor.layers.quantization.utils.quant_utils import (

19 QuantKey,

20 kNvfp4Dynamic,

@@ -42,7 +39,6 @@ __all__ = [

39 "is_flashinfer_fp4_cutlass_moe_available",

40 "is_flashinfer_fp4_cutedsl_moe_available",

41 "reorder_w1w3_to_w3w1",

45 - "build_flashinfer_fp4_cutlass_moe_prepare_finalize",

42 ]

44 #

@@ -163,17 +159,6 @@ def reorder_w1w3_to_w3w1(

159 )

160

161

166 -def build_flashinfer_fp4_cutlass_moe_prepare_finalize(

167 - moe: FusedMoEConfig,

168 -) -> mk.FusedMoEPrepareAndFinalize:

169 - """Create a FlashInfer CUTLASS fused-MoE prepare finalize kernel"""

170 - use_dp = moe.moe_parallel_config.dp_size > 1

171 - enable_alltoallv = moe.moe_parallel_config.all2all_backend == "flashinfer_all2allv"

172 - return create_flashinfer_prepare_finalize(

173 - use_dp=use_dp, use_nvfp4=True, enable_alltoallv=enable_alltoallv

174 - )

175 -

176 -

162 def prepare_static_weights_for_trtllm_fp4_moe(

163 # args_dequant,

164 # args,

vllm/model_executor/layers/quantization/utils/flashinfer_utils.py

-19

@@ -4,15 +4,8 @@ from enum import Enum

5 import torch

7 -import vllm.model_executor.layers.fused_moe.modular_kernel as mk

7 from vllm import envs

8 from vllm.logger import init_logger

10 -from vllm.model_executor.layers.fused_moe.config import (

11 - FusedMoEConfig,

12 -)

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

14 - create_flashinfer_prepare_finalize,

15 -)

9 from vllm.platforms import current_platform

10 from vllm.utils.math_utils import round_up

@@ -163,18 +156,6 @@ def make_fp8_moe_alpha_scales_for_fi(

156 return g1_alphas, g2_alphas

157

158

166 -def build_flashinfer_fp8_cutlass_moe_prepare_finalize(

167 - moe: FusedMoEConfig | None, use_deepseek_fp8_block_scale: bool = False

168 -) -> mk.FusedMoEPrepareAndFinalize:

169 - """Create a FlashInfer CUTLASS fused-MoE prepare finalize kernel"""

170 - use_dp = moe.moe_parallel_config.dp_size > 1 if moe is not None else False

171 - # Propagate block-scale flag so prepare/finalize can skip act quantization

172 - # and inform the kernel to consume per-block weight scales.

173 - return create_flashinfer_prepare_finalize(

174 - use_dp, use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale

175 - )

176 -

177 -

159 def get_flashinfer_moe_backend() -> FlashinferMoeBackend:

160 backend_map = {

161 "throughput": FlashinferMoeBackend.CUTLASS,

vllm/model_executor/warmup/deep_gemm_warmup.py

+3 -3

@@ -14,7 +14,6 @@ from vllm.distributed.parallel_state import get_dp_group, is_global_first_rank

14 from vllm.model_executor.layers.fused_moe.deep_gemm_moe import DeepGemmExperts

15 from vllm.model_executor.layers.fused_moe.deep_gemm_utils import compute_aligned_M

16 from vllm.model_executor.layers.fused_moe.layer import FusedMoE, FusedMoEModularMethod

17 -from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel

17 from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (

18 TritonOrDeepGemmExperts,

19 )

@@ -169,9 +168,10 @@ def _fused_moe_grouped_gemm_may_use_deep_gemm(module: torch.nn.Module) -> bool:

168 # modular kernels could invoke deep_gemm_moe_fp8

169 return True

170

172 - mk: FusedMoEModularKernel = module.quant_method.fused_experts

171 # Further check if the ModularKernel implementation uses the DeepGemmExperts

174 - return isinstance(mk.fused_experts, (DeepGemmExperts, TritonOrDeepGemmExperts))

172 + return isinstance(

173 + module.quant_method.moe_mk, (DeepGemmExperts, TritonOrDeepGemmExperts)

174 + )

175

176

177 FP8_GEMM_NT_WARMUP_CACHE: set[torch.Size] = set()