@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 97995f6376

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

[MoE Refactor] Create MK for TRTLLM Kernels (#32564)

Signed-off-by: Robert Shaw <robshaw@redhat.com> Signed-off-by: Robert Shaw <rshaw@neuralmagic.com> Signed-off-by: Robert Shaw <robertgshaw2@gmail.com> Co-authored-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <rshaw@neuralmagic.com>

Robert Shaw committed Mar 3, 2026 at 13:39 UTC 97995f6376fd3dae7a67624055ddf038233e181e

77 files changed +2574 -2086

.buildkite/test_areas/kernels.yaml

+2 -1

@@ -44,7 +44,8 @@ steps:

44 - vllm/envs.py

45 - vllm/config

46 commands:

47 - - pytest -v -s kernels/moe --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT

47 + - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT

48 + - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT

49 parallelism: 2

51 - label: Kernels Mamba Test

benchmarks/kernels/benchmark_cutlass_moe_fp8.py

+17 -11

@@ -12,12 +12,12 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk

12 from tests.kernels.moe.utils import make_dummy_moe_config

13 from vllm import _custom_ops as ops

14 from vllm.model_executor.layers.fused_moe.activation import MoEActivation

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

16 + maybe_make_prepare_finalize,

17 +)

18 from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config

19 from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8

20 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk

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

19 - MoEPrepareAndFinalizeNoEP,

20 -)

21 from vllm.platforms import current_platform

22 from vllm.utils.argparse_utils import FlexibleArgumentParser

23 from vllm.v1.worker.workspace import init_workspace_manager

@@ -137,15 +137,21 @@ def bench_run(

137 per_out_ch_quant=per_out_ch,

138 )

139

140 - fn = mk.FusedMoEModularKernel(

141 - MoEPrepareAndFinalizeNoEP(),

140 + moe_config = make_dummy_moe_config(

141 + num_experts=num_experts,

142 + hidden_dim=k,

143 + intermediate_size_per_partition=n,

144 + in_dtype=a.dtype,

145 + )

146 + fn = mk.FusedMoEKernel(

147 + maybe_make_prepare_finalize(

148 + moe=moe_config,

149 + quant_config=quant_config,

150 + allow_new_interface=True,

151 + use_monolithic=False,

152 + ),

153 CutlassExpertsFp8(

143 - moe_config=make_dummy_moe_config(

144 - num_experts=num_experts,

145 - hidden_dim=k,

146 - intermediate_size_per_partition=n,

147 - in_dtype=a.dtype,

148 - ),

154 + moe_config=moe_config,

155 quant_config=quant_config,

156 ),

157 )

benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py

+26 -9

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

15 from tests.kernels.moe.utils import make_dummy_moe_config

16 from vllm import _custom_ops as ops

17 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config

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

19 + maybe_make_prepare_finalize,

20 +)

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

22 fp8_w8a8_moe_quant_config,

23 nvfp4_moe_quant_config,

@@ -23,9 +26,6 @@ from vllm.model_executor.layers.fused_moe.cutlass_moe import (

26 CutlassExpertsFp4,

27 )

28 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk

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

27 - MoEPrepareAndFinalizeNoEP,

28 -)

29 from vllm.scalar_type import scalar_types

30 from vllm.utils.argparse_utils import FlexibleArgumentParser

31 from vllm.v1.worker.workspace import init_workspace_manager

@@ -196,10 +196,21 @@ def bench_run(

196 g2_alphas=w2_gs,

197 )

198

199 - kernel = mk.FusedMoEModularKernel(

200 - MoEPrepareAndFinalizeNoEP(),

199 + moe_config = make_dummy_moe_config(

200 + num_experts=num_experts,

201 + hidden_dim=k,

202 + intermediate_size_per_partition=n,

203 + in_dtype=a.dtype,

204 + )

205 + kernel = mk.FusedMoEKernel(

206 + maybe_make_prepare_finalize(

207 + moe=moe_config,

208 + quant_config=quant_config,

209 + allow_new_interface=True,

210 + use_monolithic=False,

211 + ),

212 CutlassExpertsFp4(

202 - make_dummy_moe_config(),

213 + moe_config=moe_config,

214 quant_config=quant_config,

215 ),

216 )

@@ -240,11 +251,17 @@ def bench_run(

251 g1_alphas=w1_gs,

252 g2_alphas=w2_gs,

253 )

254 + moe_config = make_dummy_moe_config()

255

244 - kernel = mk.FusedMoEModularKernel(

245 - MoEPrepareAndFinalizeNoEP(),

256 + kernel = mk.FusedMoEKernel(

257 + maybe_make_prepare_finalize(

258 + moe=moe_config,

259 + quant_config=quant_config,

260 + allow_new_interface=True,

261 + use_monolithic=False,

262 + ),

263 CutlassExpertsFp4(

247 - make_dummy_moe_config(),

264 + moe_config=moe_config,

265 quant_config=quant_config,

266 ),

267 )

benchmarks/kernels/benchmark_grouped_gemm_cutlass.py

+31 -19

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

9 from tests.kernels.moe.utils import make_dummy_moe_config

10 from vllm import _custom_ops as ops

11 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config

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

13 + maybe_make_prepare_finalize,

14 +)

15 from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config

16 from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8

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

18 fused_experts,

19 fused_topk,

20 )

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

19 - MoEPrepareAndFinalizeNoEP,

20 -)

21 from vllm.utils.argparse_utils import FlexibleArgumentParser

22 from vllm.v1.worker.workspace import init_workspace_manager

@@ -131,16 +131,22 @@ def bench_run(

131 w2_scale=w2_scale,

132 per_act_token_quant=per_act_token,

133 )

134 + moe_config = make_dummy_moe_config(

135 + num_experts=w2.shape[0],

136 + hidden_dim=w2.shape[1],

137 + intermediate_size_per_partition=w2.shape[2],

138 + in_dtype=a.dtype,

139 + )

140

135 - fn = mk.FusedMoEModularKernel(

136 - MoEPrepareAndFinalizeNoEP(),

141 + fn = mk.FusedMoEKernel(

142 + maybe_make_prepare_finalize(

143 + moe=moe_config,

144 + quant_config=quant_config,

145 + allow_new_interface=True,

146 + use_monolithic=False,

147 + ),

148 CutlassExpertsFp8(

138 - moe_config=make_dummy_moe_config(

139 - num_experts=w2.shape[0],

140 - hidden_dim=w2.shape[1],

141 - intermediate_size_per_partition=w2.shape[2],

142 - in_dtype=a.dtype,

143 - ),

149 + moe_config=moe_config,

150 quant_config=quant_config,

151 ),

152 )

@@ -163,16 +169,22 @@ def bench_run(

169 w2_scale=w2_scale,

170 per_act_token_quant=per_act_token,

171 )

172 + moe_config = make_dummy_moe_config(

173 + num_experts=w2.shape[0],

174 + hidden_dim=w2.shape[1],

175 + intermediate_size_per_partition=w2.shape[2],

176 + in_dtype=a.dtype,

177 + )

178

167 - fn = mk.FusedMoEModularKernel(

168 - MoEPrepareAndFinalizeNoEP(),

179 + fn = mk.FusedMoEKernel(

180 + maybe_make_prepare_finalize(

181 + moe=moe_config,

182 + quant_config=quant_config,

183 + allow_new_interface=True,

184 + use_monolithic=False,

185 + ),

186 CutlassExpertsFp8(

170 - moe_config=make_dummy_moe_config(

171 - num_experts=w2.shape[0],

172 - hidden_dim=w2.shape[1],

173 - intermediate_size_per_partition=w2.shape[2],

174 - in_dtype=a.dtype,

175 - ),

187 + moe_config=moe_config,

188 quant_config=quant_config,

189 ),

190 )

benchmarks/kernels/benchmark_moe.py

+37 -17

@@ -17,6 +17,9 @@ from ray.experimental.tqdm_ray import tqdm

18 from vllm.model_executor.layers.fused_moe import fused_topk

19 from vllm.model_executor.layers.fused_moe.activation import MoEActivation

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

21 + maybe_make_prepare_finalize,

22 +)

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

24 FusedMoEConfig,

25 FusedMoEParallelConfig,

@@ -242,24 +245,33 @@ def benchmark_config(

245

246 deep_gemm_experts = None

247 if use_deep_gemm:

245 - deep_gemm_experts = mk.FusedMoEModularKernel(

246 - prepare_finalize=MoEPrepareAndFinalizeNoEP(),

248 + moe_config = (

249 + FusedMoEConfig(

250 + num_experts=num_experts,

251 + experts_per_token=topk,

252 + hidden_dim=hidden_size,

253 + intermediate_size_per_partition=shard_intermediate_size,

254 + num_local_experts=num_experts,

255 + num_logical_experts=num_experts,

256 + activation=MoEActivation.SILU,

257 + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),

258 + in_dtype=init_dtype,

259 + routing_method=RoutingMethodType.TopK,

260 + device="cuda",

261 + ),

262 + )

263 + deep_gemm_experts = mk.FusedMoEKernel(

264 + prepare_finalize=maybe_make_prepare_finalize(

265 + moe=moe_config,

266 + quant_config=quant_config,

267 + allow_new_interface=True,

268 + use_monolithic=False,

269 + ),

270 fused_experts=TritonOrDeepGemmExperts(

248 - moe_config=FusedMoEConfig(

249 - num_experts=num_experts,

250 - experts_per_token=topk,

251 - hidden_dim=hidden_size,

252 - intermediate_size_per_partition=shard_intermediate_size,

253 - num_local_experts=num_experts,

254 - num_logical_experts=num_experts,

255 - activation=MoEActivation.SILU,

256 - moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),

257 - in_dtype=init_dtype,

258 - routing_method=RoutingMethodType.TopK,

259 - device="cuda",

260 - ),

271 + moe_config=moe_config,

272 quant_config=quant_config,

273 ),

274 + inplace=not disable_inplace(),

275 )

276

277 with override_config(config):

@@ -269,8 +281,16 @@ def benchmark_config(

281

282 inplace = not disable_inplace()

283 if use_deep_gemm:

272 - return deep_gemm_experts(

273 - x, w1, w2, topk_weights, topk_ids, inplace=inplace

284 + return deep_gemm_experts.apply(

285 + x,

286 + w1,

287 + w2,

288 + topk_weights,

289 + topk_ids,

290 + activation=MoEActivation.SILU,

291 + global_num_experts=num_experts,

292 + apply_router_weight_on_input=False,

293 + expert_map=False,

294 )

295 return fused_experts(

296 x,

docs/design/dbo.md

+1 -1

@@ -81,7 +81,7 @@ The current implementation has all `dbo_yield` and `dbo_maybe_run_recv_hook` cal

82 The `make_ubatch_context` function initializes two `UBatchContexts`, one for each UBatch thread. It takes two CUDA streams, the preexisting `ForwardContexts` and a CPU thread barrier. This function should be used exclusively to instantiate `UBatchContexts`. It will handle all of the event initialization.

84 -The `dbo_register_recv_hook` method registers a callback that can be returned by the `FusedMoEPrepareAndFinalize` class in the other UBatch thread’s `UBatchContext`. The callback will be run when the other thread calls `dbo_maybe_run_recv_hook`. This is typically used to wait on an all-to-all kernel.

84 +The `dbo_register_recv_hook` method registers a callback that can be returned by the `FusedMoEPrepareAndFinalizeModular` class in the other UBatch thread’s `UBatchContext`. The callback will be run when the other thread calls `dbo_maybe_run_recv_hook`. This is typically used to wait on an all-to-all kernel.

86 The `dbo_maybe_run_recv_hook` method runs a callback that’s set by the `dbo_register_recv_hook` function if that callback exists.

docs/design/fused_moe_modular_kernel.md

+52 -52

@@ -37,31 +37,31 @@ The rest of the document will focus on the Contiguous / Non-Batched case. Extrap

37 FusedMoEModularKernel splits the FusedMoE operation into 3 parts,

39 1. TopKWeightAndReduce

40 -2. FusedMoEPrepareAndFinalize

41 -3. FusedMoEPermuteExpertsUnpermute

40 +2. FusedMoEPrepareAndFinalizeModular

41 +3. FusedMoEExpertsModular

43 ### TopKWeightAndReduce

45 -The TopK Weight Application and Reduction components happen right after the Unpermute operation and before the All2All Combine. Note that the `FusedMoEPermuteExpertsUnpermute` is responsible for the Unpermute and `FusedMoEPrepareAndFinalize` is responsible for the All2All Combine. There is value in doing the TopK Weight Application and Reduction in the `FusedMoEPermuteExpertsUnpermute`. But some implementations choose to do it `FusedMoEPrepareAndFinalize`. In order to enable this flexibility, we have a TopKWeightAndReduce abstract class.

45 +The TopK Weight Application and Reduction components happen right after the Unpermute operation and before the All2All Combine. Note that the `FusedMoEExpertsModular` is responsible for the Unpermute and `FusedMoEPrepareAndFinalizeModular` is responsible for the All2All Combine. There is value in doing the TopK Weight Application and Reduction in the `FusedMoEExpertsModular`. But some implementations choose to do it `FusedMoEPrepareAndFinalizeModular`. In order to enable this flexibility, we have a TopKWeightAndReduce abstract class.

47 Please find the implementations of TopKWeightAndReduce [here](../../vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py).

49 -`FusedMoEPrepareAndFinalize::finalize()` method accepts a `TopKWeightAndReduce` argument that is invoked inside the method.

50 -The `FusedMoEModularKernel` acts as a bridge between the `FusedMoEPermuteExpertsUnpermute` and `FusedMoEPerpareAndFinalize` implementations to determine where the TopK Weight Application and Reduction happens.

49 +`FusedMoEPrepareAndFinalizeModular::finalize()` method accepts a `TopKWeightAndReduce` argument that is invoked inside the method.

50 +The `FusedMoEModularKernel` acts as a bridge between the `FusedMoEExpertsModular` and `FusedMoEPerpareAndFinalize` implementations to determine where the TopK Weight Application and Reduction happens.

52 -* `FusedMoEPermuteExpertsUnpermute::finalize_weight_and_reduce_impl` method returns `TopKWeightAndReduceNoOp` if the `FusedMoEPermuteExpertsUnpermute` implementation does the weight application and reduction itself.

53 -* `FusedMoEPermuteExpertsUnpermute::finalize_weight_and_reduce_impl` method returns `TopKWeightAndReduceContiguous` / `TopKWeightAndReduceNaiveBatched` / `TopKWeightAndReduceDelegate` if the `FusedMoEPermuteExpertsUnpermute` implementation needs the `FusedMoEPrepareAndFinalize::finalize()` to do the weight application and reduction.

52 +* `FusedMoEExpertsModular::finalize_weight_and_reduce_impl` method returns `TopKWeightAndReduceNoOp` if the `FusedMoEExpertsModular` implementation does the weight application and reduction itself.

53 +* `FusedMoEExpertsModular::finalize_weight_and_reduce_impl` method returns `TopKWeightAndReduceContiguous` / `TopKWeightAndReduceNaiveBatched` / `TopKWeightAndReduceDelegate` if the `FusedMoEExpertsModular` implementation needs the `FusedMoEPrepareAndFinalizeModular::finalize()` to do the weight application and reduction.

55 -### FusedMoEPrepareAndFinalize

55 +### FusedMoEPrepareAndFinalizeModular

57 -The `FusedMoEPrepareAndFinalize` abstract class exposes `prepare`, `prepare_no_receive` and `finalize` functions.

58 -The `prepare` function is responsible for input activation Quantization and All2All Dispatch. If implemented, The `prepare_no_receive` is like `prepare` except it does not wait to receive results from other workers. Instead it returns a "receiver" callback that must be invoked to wait for the final results of worker. It is not required that this method is supported by all `FusedMoEPrepareAndFinalize` classes, but if it is available, it can be used to interleave work with the initial all to all communication, e.g. interleaving shared experts with fused experts. The `finalize` function is responsible for invoking the All2All Combine. Additionally the `finalize` function may or may not do the TopK weight application and reduction (Please refer to the TopKWeightAndReduce section)

57 +The `FusedMoEPrepareAndFinalizeModular` abstract class exposes `prepare`, `prepare_no_receive` and `finalize` functions.

58 +The `prepare` function is responsible for input activation Quantization and All2All Dispatch. If implemented, The `prepare_no_receive` is like `prepare` except it does not wait to receive results from other workers. Instead it returns a "receiver" callback that must be invoked to wait for the final results of worker. It is not required that this method is supported by all `FusedMoEPrepareAndFinalizeModular` classes, but if it is available, it can be used to interleave work with the initial all to all communication, e.g. interleaving shared experts with fused experts. The `finalize` function is responsible for invoking the All2All Combine. Additionally the `finalize` function may or may not do the TopK weight application and reduction (Please refer to the TopKWeightAndReduce section)

60 -![FusedMoEPrepareAndFinalize Blocks](../assets/design/fused_moe_modular_kernel/prepare_and_finalize_blocks.png)

60 +![FusedMoEPrepareAndFinalizeModular Blocks](../assets/design/fused_moe_modular_kernel/prepare_and_finalize_blocks.png)

62 -### FusedMoEPermuteExpertsUnpermute

62 +### FusedMoEExpertsModular

64 -The `FusedMoEPermuteExpertsUnpermute` class is where the crux of the MoE operations happen. The `FusedMoEPermuteExpertsUnpermute` abstract class exposes a few important functions,

64 +The `FusedMoEExpertsModular` class is where the crux of the MoE operations happen. The `FusedMoEExpertsModular` abstract class exposes a few important functions,

66 * apply()

67 * workspace_shapes()

@@ -81,25 +81,25 @@ The `apply` method is where the implementations perform

82 #### workspace_shapes()

84 -The core FusedMoE implementation performs a series of operations. It would be inefficient to create output memory for each of these operations separately. To that effect, implementations are required to declare 2 workspace shapes, the workspace datatype and the FusedMoE output shape as outputs of the workspace_shapes() method. This information is used to allocate the workspace tensors and the output tensor in `FusedMoEModularKernel::forward()` and passed on to the `FusedMoEPermuteExpertsUnpermute::apply()` method. The workspaces could then be used as intermediate buffers in the FusedMoE implementation.

84 +The core FusedMoE implementation performs a series of operations. It would be inefficient to create output memory for each of these operations separately. To that effect, implementations are required to declare 2 workspace shapes, the workspace datatype and the FusedMoE output shape as outputs of the workspace_shapes() method. This information is used to allocate the workspace tensors and the output tensor in `FusedMoEModularKernel::forward()` and passed on to the `FusedMoEExpertsModular::apply()` method. The workspaces could then be used as intermediate buffers in the FusedMoE implementation.

86 #### finalize_weight_and_reduce_impl()

88 -It is sometimes efficient to perform TopK weight application and Reduction inside the `FusedMoEPermuteExpertsUnpermute::apply()`. Find an example [here](https://github.com/vllm-project/vllm/pull/20228). We have a `TopKWeightAndReduce` abstract class to facilitate such implementations. Please refer to the TopKWeightAndReduce section.

89 -`FusedMoEPermuteExpertsUnpermute::finalize_weight_and_reduce_impl()` returns the `TopKWeightAndReduce` object that the implementation wants the `FusedMoEPrepareAndFinalize::finalize()` to use.

88 +It is sometimes efficient to perform TopK weight application and Reduction inside the `FusedMoEExpertsModular::apply()`. Find an example [here](https://github.com/vllm-project/vllm/pull/20228). We have a `TopKWeightAndReduce` abstract class to facilitate such implementations. Please refer to the TopKWeightAndReduce section.

89 +`FusedMoEExpertsModular::finalize_weight_and_reduce_impl()` returns the `TopKWeightAndReduce` object that the implementation wants the `FusedMoEPrepareAndFinalizeModular::finalize()` to use.

91 -![FusedMoEPermuteExpertsUnpermute Blocks](../assets/design/fused_moe_modular_kernel/fused_experts_blocks.png)

91 +![FusedMoEExpertsModular Blocks](../assets/design/fused_moe_modular_kernel/fused_experts_blocks.png)

93 ### FusedMoEModularKernel

95 -`FusedMoEModularKernel` is composed of the `FusedMoEPrepareAndFinalize` and `FusedMoEPermuteExpertsUnpermute` objects.

95 +`FusedMoEModularKernel` is composed of the `FusedMoEPrepareAndFinalizeModular` and `FusedMoEExpertsModular` objects.

96 `FusedMoEModularKernel` pseudocode/sketch,

98 ```py

99 class FusedMoEModularKernel:

100 def __init__(self,

101 - prepare_finalize: FusedMoEPrepareAndFinalize,

102 - fused_experts: FusedMoEPermuteExpertsUnpermute):

101 + prepare_finalize: FusedMoEPrepareAndFinalizeModular,

102 + fused_experts: FusedMoEExpertsModular):

103

104 self.prepare_finalize = prepare_finalize

105 self.fused_experts = fused_experts

@@ -128,53 +128,53 @@ class FusedMoEModularKernel:

128

129 ## How-To

130

131 -### How To Add a FusedMoEPrepareAndFinalize Type

131 +### How To Add a FusedMoEPrepareAndFinalizeModular Type

132

133 -Typically a FusedMoEPrepareAndFinalize type is backed by an All2All Dispatch & Combine implementation / kernel. For example,

133 +Typically a FusedMoEPrepareAndFinalizeModular type is backed by an All2All Dispatch & Combine implementation / kernel. For example,

134

135 * DeepEPHTPrepareAndFinalize type is backed by DeepEP High-Throughput All2All kernels, and

136 * DeepEPLLPrepareAndFinalize type is backed by DeepEP Low-Latency All2All kernels.

137

138 #### Step 1: Add an All2All manager

139

140 -The purpose of the All2All Manager is to set up the All2All kernel implementations. The `FusedMoEPrepareAndFinalize` implementations typically fetch a kernel-implementation "handle" from the All2All Manager to invoke the Dispatch and Combine functions. Please look at the All2All Manager implementations [here](../../vllm/distributed/device_communicators/all2all.py).

140 +The purpose of the All2All Manager is to set up the All2All kernel implementations. The `FusedMoEPrepareAndFinalizeModular` implementations typically fetch a kernel-implementation "handle" from the All2All Manager to invoke the Dispatch and Combine functions. Please look at the All2All Manager implementations [here](../../vllm/distributed/device_communicators/all2all.py).

141

142 -#### Step 2: Add a FusedMoEPrepareAndFinalize Type

142 +#### Step 2: Add a FusedMoEPrepareAndFinalizeModular Type

143

144 -This section describes the significance of the various functions exposed by the `FusedMoEPrepareAndFinalize` abstract class.

144 +This section describes the significance of the various functions exposed by the `FusedMoEPrepareAndFinalizeModular` abstract class.

145

146 -`FusedMoEPrepareAndFinalize::prepare()`: The prepare method implements the Quantization and All2All Dispatch. Typically the Dispatch function from the relevant All2All Manager is invoked.

146 +`FusedMoEPrepareAndFinalizeModular::prepare()`: The prepare method implements the Quantization and All2All Dispatch. Typically the Dispatch function from the relevant All2All Manager is invoked.

147

148 -`FusedMoEPrepareAndFinalize::has_prepare_no_receive()`: Indicates whether or not this subclass implements `prepare_no_receive`. Defaults to False.

148 +`FusedMoEPrepareAndFinalizeModular::has_prepare_no_receive()`: Indicates whether or not this subclass implements `prepare_no_receive`. Defaults to False.

149

150 -`FusedMoEPrepareAndFinalize::prepare_no_receive()`: The prepare_no_receive method implements the Quantization and All2All Dispatch. It does not wait for the result of the dispatch operation but instead returns a thunk that can be invoked to wait for the final results. Typically the Dispatch function from the relevant All2All Manager is invoked.

150 +`FusedMoEPrepareAndFinalizeModular::prepare_no_receive()`: The prepare_no_receive method implements the Quantization and All2All Dispatch. It does not wait for the result of the dispatch operation but instead returns a thunk that can be invoked to wait for the final results. Typically the Dispatch function from the relevant All2All Manager is invoked.

151

152 -`FusedMoEPrepareAndFinalize::finalize()`: Maybe perform TopK Weight Application and Reduction and All2All Combine. Typically the Combine function from the relevant All2AllManager is invoked.

152 +`FusedMoEPrepareAndFinalizeModular::finalize()`: Maybe perform TopK Weight Application and Reduction and All2All Combine. Typically the Combine function from the relevant All2AllManager is invoked.

153

154 -`FusedMoEPrepareAndFinalize::activation_format()`: Return `FusedMoEActivationFormat.BatchedExperts` if the output of the prepare method (i.e. the All2All dispatch) is Batched. Return `FusedMoEActivationFormat.Standard` otherwise.

154 +`FusedMoEPrepareAndFinalizeModular::activation_format()`: Return `FusedMoEActivationFormat.BatchedExperts` if the output of the prepare method (i.e. the All2All dispatch) is Batched. Return `FusedMoEActivationFormat.Standard` otherwise.

155

156 -`FusedMoEPrepareAndFinalize::topk_indices_dtype()`: Data type of the TopK ids. Some All2All kernels have strict requirements pertaining to the data type of the TopK ids. This requirement is passed on to the `FusedMoe::select_experts` function so it could be respected. If there are no strict requirements return None.

156 +`FusedMoEPrepareAndFinalizeModular::topk_indices_dtype()`: Data type of the TopK ids. Some All2All kernels have strict requirements pertaining to the data type of the TopK ids. This requirement is passed on to the `FusedMoe::select_experts` function so it could be respected. If there are no strict requirements return None.

157

158 -`FusedMoEPrepareAndFinalize::max_num_tokens_per_rank()`: This is the maximum number of tokens that would be submitted to the All2All Dispatch at once.

158 +`FusedMoEPrepareAndFinalizeModular::max_num_tokens_per_rank()`: This is the maximum number of tokens that would be submitted to the All2All Dispatch at once.

159

160 -`FusedMoEPrepareAndFinalize::num_dispatchers()`: Total number of dispatching units. This value determines the size of the Dispatch output. The Dispatch output is of shape (num_local_experts, max_num_tokens, K). Here max_num_tokens = num_dispatchers() * max_num_tokens_per_rank().

160 +`FusedMoEPrepareAndFinalizeModular::num_dispatchers()`: Total number of dispatching units. This value determines the size of the Dispatch output. The Dispatch output is of shape (num_local_experts, max_num_tokens, K). Here max_num_tokens = num_dispatchers() * max_num_tokens_per_rank().

161

162 -We suggest picking an already existing `FusedMoEPrepareAndFinalize` implementation that matches your All2All implementation closely and using it as a reference.

162 +We suggest picking an already existing `FusedMoEPrepareAndFinalizeModular` implementation that matches your All2All implementation closely and using it as a reference.

163

164 -### How To Add a FusedMoEPermuteExpertsUnpermute Type

164 +### How To Add a FusedMoEExpertsModular Type

165

166 -FusedMoEPermuteExpertsUnpermute performs the core of the FusedMoE operations. The various functions exposed by the abstract class and their significance is as follows,

166 +FusedMoEExpertsModular performs the core of the FusedMoE operations. The various functions exposed by the abstract class and their significance is as follows,

167

168 -`FusedMoEPermuteExpertsUnpermute::activation_formats()`: Return the supported Input and Output activation formats. i.e. Contiguous / Batched format.

168 +`FusedMoEExpertsModular::activation_formats()`: Return the supported Input and Output activation formats. i.e. Contiguous / Batched format.

169

170 -`FusedMoEPermuteExpertsUnpermute::supports_chunking()`: Return True if the implementation supports chunking. Typically

170 +`FusedMoEExpertsModular::supports_chunking()`: Return True if the implementation supports chunking. Typically

171 implementations that input `FusedMoEActivationFormat.Standard` support chunking and `FusedMoEActivationFormat.BatchedExperts` do not.

172

173 -`FusedMoEPermuteExpertsUnpermute::supports_expert_map()`: Return True if the implementation supports expert map.

173 +`FusedMoEExpertsModular::supports_expert_map()`: Return True if the implementation supports expert map.

174

175 -`FusedMoEPermuteExpertsUnpermute::workspace_shapes()` /

176 -`FusedMoEPermuteExpertsUnpermute::finalize_weight_and_reduce_impl` /

177 -`FusedMoEPermuteExpertsUnpermute::apply`: Refer to `FusedMoEPermuteExpertsUnpermute` section above.

175 +`FusedMoEExpertsModular::workspace_shapes()` /

176 +`FusedMoEExpertsModular::finalize_weight_and_reduce_impl` /

177 +`FusedMoEExpertsModular::apply`: Refer to `FusedMoEExpertsModular` section above.

178

179 ### FusedMoEModularKernel Initialization

180

@@ -186,14 +186,14 @@ implementations that input `FusedMoEActivationFormat.Standard` support chunking

186

187 #### maybe_make_prepare_finalize

188

189 -The `maybe_make_prepare_finalize` method is responsible for constructing an instance of `FusedMoEPrepareAndFinalize` when appropriate based on the current all2all backend, e.g. when EP + DP is enabled. The base class method currently constructs all the `FusedMoEPrepareAndFinalize` objects for the EP+DP case. Derived classes can override this method to construct prepare/finalize objects for different scenarios, e.g. `ModelOptNvFp4FusedMoE` can construct a `FlashInferCutlassMoEPrepareAndFinalize` for the EP+TP case.

189 +The `maybe_make_prepare_finalize` method is responsible for constructing an instance of `FusedMoEPrepareAndFinalizeModular` when appropriate based on the current all2all backend, e.g. when EP + DP is enabled. The base class method currently constructs all the `FusedMoEPrepareAndFinalizeModular` objects for the EP+DP case. Derived classes can override this method to construct prepare/finalize objects for different scenarios, e.g. `ModelOptNvFp4FusedMoE` can construct a `FlashInferCutlassMoEPrepareAndFinalize` for the EP+TP case.

190 Please refer to the implementations in,

191

192 * `ModelOptNvFp4FusedMoE`

193

194 #### select_gemm_impl

195

196 -The `select_gemm_impl` method is undefined in the base class. It is the responsibility of the derived class to implement a method that constructs a valid/appropriate `FusedMoEPermuteExpertsUnpermute` object.

196 +The `select_gemm_impl` method is undefined in the base class. It is the responsibility of the derived class to implement a method that constructs a valid/appropriate `FusedMoEExpertsModular` object.

197 Please refer to the implementations in,

198

199 * `UnquantizedFusedMoEMethod`

@@ -205,7 +205,7 @@ derived classes.

205

206 #### init_prepare_finalize

207

208 -Based on the input and env settings, the `init_prepare_finalize` method creates the appropriate `FusedMoEPrepareAndFinalize` object. The method then queries `select_gemm_impl` for the appropriate `FusedMoEPermuteExpertsUnpermute` object and builds the `FusedMoEModularKernel` object

208 +Based on the input and env settings, the `init_prepare_finalize` method creates the appropriate `FusedMoEPrepareAndFinalizeModular` object. The method then queries `select_gemm_impl` for the appropriate `FusedMoEExpertsModular` object and builds the `FusedMoEModularKernel` object

209

210 Please take a look at [init_prepare_finalize](https://github.com/vllm-project/vllm/blob/1cbf951ba272c230823b947631065b826409fa62/vllm/model_executor/layers/fused_moe/layer.py#L188).

211 **Important**: The `FusedMoEMethodBase` derived classes use the `FusedMoEMethodBase::fused_experts` object in their `apply` methods. When settings permit the construction of a valid `FusedMoEModularKernel` object, we override `FusedMoEMethodBase::fused_experts` with it. This essentially makes the derived classes agnostic to what FusedMoE implementation is used.

@@ -214,9 +214,9 @@ Please take a look at [init_prepare_finalize](https://github.com/vllm-project/vl

214

215 We have `FusedMoEModularKernel` unit tests at [test_modular_kernel_combinations.py](../../tests/kernels/moe/test_modular_kernel_combinations.py).

216

217 -The unit test iterates through all combinations of `FusedMoEPrepareAndFinalize` and `FusedMoEPremuteExpertsUnpermute` types and if they are

217 +The unit test iterates through all combinations of `FusedMoEPrepareAndFinalizeModular` and `FusedMoEPremuteExpertsUnpermute` types and if they are

218 compatible, runs some correctness tests.

219 -If you are adding some `FusedMoEPrepareAndFinalize` / `FusedMoEPermuteExpertsUnpermute` implementations,

219 +If you are adding some `FusedMoEPrepareAndFinalizeModular` / `FusedMoEExpertsModular` implementations,

220

221 1. Add the implementation type to `MK_ALL_PREPARE_FINALIZE_TYPES` and `MK_FUSED_EXPERT_TYPES` in [mk_objects.py](../../tests/kernels/moe/modular_kernel_tools/mk_objects.py) respectively.

222 2. Update `Config::is_batched_prepare_finalize()`, `Config::is_batched_fused_experts()`, `Config::is_standard_fused_experts()`,

@@ -225,24 +225,24 @@ If you are adding some `FusedMoEPrepareAndFinalize` / `FusedMoEPermuteExpertsUnp

225

226 Doing this will add the new implementation to the test suite.

227

228 -### How To Check `FusedMoEPrepareAndFinalize` & `FusedMoEPermuteExpertsUnpermute` Compatibility

228 +### How To Check `FusedMoEPrepareAndFinalizeModular` & `FusedMoEExpertsModular` Compatibility

229

230 The unit test file [test_modular_kernel_combinations.py](../../tests/kernels/moe/test_modular_kernel_combinations.py) can also be executed as a standalone script.

231 Example: `python3 -m tests.kernels.moe.test_modular_kernel_combinations --pf-type DeepEPLLPrepareAndFinalize --experts-type BatchedTritonExperts`

232 -As a side effect, this script can be used to test `FusedMoEPrepareAndFinalize` & `FusedMoEPermuteExpertsUnpermute` compatibility. When invoked

232 +As a side effect, this script can be used to test `FusedMoEPrepareAndFinalizeModular` & `FusedMoEExpertsModular` compatibility. When invoked

233 with incompatible types, the script will error.

234

235 ### How To Profile

236

237 Please take a look at [profile_modular_kernel.py](../../tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py)

238 The script can be used to generate Torch traces for a single `FusedMoEModularKernel::forward()` call for any compatible

239 -`FusedMoEPrepareAndFinalize` and `FusedMoEPermuteExpertsUnpermute` types.

239 +`FusedMoEPrepareAndFinalizeModular` and `FusedMoEExpertsModular` types.

240 Example: `python3 -m tests.kernels.moe.modular_kernel_tools.profile_modular_kernel --pf-type DeepEPLLPrepareAndFinalize --experts-type BatchedTritonExperts`

241

242 -## FusedMoEPrepareAndFinalize Implementations

242 +## FusedMoEPrepareAndFinalizeModular Implementations

243

244 See [Fused MoE Kernel features](./moe_kernel_features.md#fused-moe-modular-all2all-backends) for a list of all the available modular prepare and finalize subclasses.

245

246 -## FusedMoEPermuteExpertsUnpermute

246 +## FusedMoEExpertsModular

247

248 See [Fused MoE Kernel features](./moe_kernel_features.md#fused-moe-experts-kernels) for a list of all the available modular experts.

docs/design/moe_kernel_features.md

+7 -9

@@ -4,17 +4,17 @@ The purpose of this document is to provide an overview of the various MoE kernel

5 ## Fused MoE Modular All2All backends

7 -There are a number of all2all communication backends that are used to implement expert parallelism (EP) for the `FusedMoE` layer. The different `FusedMoEPrepareAndFinalize` subclasses provide an interface for each all2all backend.

7 +There are a number of all2all communication backends that are used to implement expert parallelism (EP) for the `FusedMoE` layer. The different `FusedMoEPrepareAndFinalizeModular` subclasses provide an interface for each all2all backend.

9 The following table describes the relevant features of each backend, i.e. activation format, supported quantization schemes and async support.

11 -The output activation format (standard or batched) corresponds to the output of the prepare step of the `FusedMoEPrepareAndFinalize` subclass, and the finalize step requires the same format. All the backend `prepare` methods expect activations in the standard format and all the `finalize` methods return activations in standard format. More details on the formats can be found in the [Fused MoE Modular Kernel](./fused_moe_modular_kernel.md) document.

11 +The output activation format (standard or batched) corresponds to the output of the prepare step of the `FusedMoEPrepareAndFinalizeModular` subclass, and the finalize step requires the same format. All the backend `prepare` methods expect activations in the standard format and all the `finalize` methods return activations in standard format. More details on the formats can be found in the [Fused MoE Modular Kernel](./fused_moe_modular_kernel.md) document.

13 -The quantization types and formats enumerate which quantization schemes are supported by each `FusedMoEPrepareAndFinalize` class. The quantization can happen before or after the dispatch based on the format the all2all backend supports, e.g. deepep_high_throughput supports only block-quantized fp8 format. Any other format will result in dispatching in higher precision and quantizing afterwards. The output of the prepare step for each backend is the quantized type. The finalize step generally requires the same input type as the original activations, e.g. if the original input is bfloat16 and the quantization scheme is fp8 with per-tensor scales, `prepare` will return fp8/per-tensor scale activations and `finalize` will take bfloat16 activations. See the diagrams in [Fused MoE Modular Kernel](./fused_moe_modular_kernel.md) for more details on the types and formats of activations at each step of the MoE process. If no quantization type is specified, the kernel operates on float16 and/or bfloat16.

13 +The quantization types and formats enumerate which quantization schemes are supported by each `FusedMoEPrepareAndFinalizeModular` class. The quantization can happen before or after the dispatch based on the format the all2all backend supports, e.g. deepep_high_throughput supports only block-quantized fp8 format. Any other format will result in dispatching in higher precision and quantizing afterwards. The output of the prepare step for each backend is the quantized type. The finalize step generally requires the same input type as the original activations, e.g. if the original input is bfloat16 and the quantization scheme is fp8 with per-tensor scales, `prepare` will return fp8/per-tensor scale activations and `finalize` will take bfloat16 activations. See the diagrams in [Fused MoE Modular Kernel](./fused_moe_modular_kernel.md) for more details on the types and formats of activations at each step of the MoE process. If no quantization type is specified, the kernel operates on float16 and/or bfloat16.

15 Async backends support the use of DBO (Dual Batch Overlap) and shared expert overlap (where shared experts are computed during the combine step).

17 -Certain models require the topk weights to be applied to the input activations rather than the output activations when topk==1, e.g. Llama. For modular kernels, this feature is supported by the `FusedMoEPrepareAndFinalize` subclass. For non-modular kernels, it is up to the experts function to deal with this flag.

17 +Certain models require the topk weights to be applied to the input activations rather than the output activations when topk==1, e.g. Llama. For modular kernels, this feature is supported by the `FusedMoEPrepareAndFinalizeModular` subclass. For non-modular kernels, it is up to the experts function to deal with this flag.

19 Unless otherwise specified, backends are controlled via the `--all2all-backend` command-line argument (or the `all2all_backend` parameter in `ParallelConfig`). All backends except `flashinfer` only work with EP+DP or EP+TP. `Flashinfer` can work with EP or DP without EP.

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

36 | deepep_high_throughput | standard | fp8 | G(128),A,T2 | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] |

37 | deepep_low_latency | batched | fp8 | G(128),A,T3 | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize.DeepEPLLPrepareAndFinalize] |

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

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

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

40 !!! info "Table key"

41 1. All types: mxfp4, nvfp4, int4, int8, fp8

@@ -75,9 +73,9 @@ Each experts kernel supports one or more activation functions, e.g. silu or gelu

74 As with the backends, some experts support applying topk weights on the input activations. The entries in the column in this table only apply to the non-modular experts.

78 -Most experts flavors include an equivalent modular interface which will be a subclass of `FusedMoEPermuteExpertsUnpermute`.

76 +Most experts flavors include an equivalent modular interface which will be a subclass of `FusedMoEExpertsModular`.

80 -To be used with a particular `FusedMoEPrepareAndFinalize` subclass, MoE kernels must have compatible activation formats, quantization types and quantization formats.

78 +To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE kernels must have compatible activation formats, quantization types and quantization formats.

81 |--------|-------------------|--------------|---------------|---------------------|-----------------------|---------|--------|

@@ -106,7 +104,7 @@ To be used with a particular `FusedMoEPrepareAndFinalize` subclass, MoE kernels

104

105 The following table shows "families" of modular kernels that are intended to work together. There are some combinations which may work but have not yet been tested, e.g. flashinfer with other fp8 experts. Note that the "naive" backend will work with any non-modular experts.

106

109 -| backend | `FusedMoEPrepareAndFinalize` subclasses | `FusedMoEPermuteExpertsUnpermute` subclasses |

107 +| backend | `FusedMoEPrepareAndFinalizeModular` subclasses | `FusedMoEExpertsModular` subclasses |

108 |---------|-----------------------------------------|----------------------------------------------|

109 | deepep_high_throughput | `DeepEPHTPrepareAndFinalize` | `DeepGemmExperts`,`TritonExperts`,`TritonOrDeepGemmExperts`,`CutlassExpertsFp8`, `MarlinExperts` |

110 | deepep_low_latency | `DeepEPLLPrepareAndFinalize` | `BatchedDeepGemmExperts`,`BatchedTritonExperts`,`CutlassBatchedExpertsFp8`,`BatchedMarlinExperts` |

tests/kernels/moe/modular_kernel_tools/cli_args.py

+2 -2

@@ -17,13 +17,13 @@ from .mk_objects import (

19 def make_config_arg_parser(description: str):

20 - def to_pf_class_type(s: str) -> mk.FusedMoEPrepareAndFinalize:

20 + def to_pf_class_type(s: str) -> mk.FusedMoEPrepareAndFinalizeModular:

21 for pf in MK_ALL_PREPARE_FINALIZE_TYPES:

22 if pf.__name__ == s:

23 return pf

24 raise ValueError(f"Cannot find a PrepareFinalize type that matches {s}")

26 - def to_experts_class_type(s: str) -> mk.FusedMoEPermuteExpertsUnpermute:

26 + def to_experts_class_type(s: str) -> mk.FusedMoEExpertsModular:

27 for fe in MK_FUSED_EXPERT_TYPES:

28 if fe.__name__ == s:

29 return fe

tests/kernels/moe/modular_kernel_tools/common.py

+5 -4

@@ -66,7 +66,7 @@ class Config:

66 quant_config: TestMoEQuantConfig | None

68 prepare_finalize_type: mk.FusedMoEPrepareAndFinalize

69 - fused_experts_type: mk.FusedMoEPermuteExpertsUnpermute

69 + fused_experts_type: mk.FusedMoEExperts

71 fused_moe_chunk_size: int | None

72 world_size: int

@@ -566,7 +566,7 @@ def make_modular_kernel(

566 config: Config,

567 vllm_config: VllmConfig,

568 quant_config: FusedMoEQuantConfig,

569 -) -> mk.FusedMoEModularKernel:

569 +) -> mk.FusedMoEKernel:

570 def next_power_of_2(x):

571 import math

572

@@ -613,7 +613,7 @@ def make_modular_kernel(

613 config.N,

614 )

615

616 - modular_kernel = mk.FusedMoEModularKernel(

616 + modular_kernel = mk.FusedMoEKernel(

617 prepare_finalize=prepare_finalize,

618 fused_experts=fused_experts,

619 inplace=False,

@@ -667,6 +667,7 @@ def run_modular_kernel(

667 "w2": rank_weights.w2,

668 "topk_weights": rank_tensors.topk_weights,

669 "topk_ids": topk_ids,

670 + "activation": MoEActivation.SILU,

671 "expert_map": rank_tensors.expert_map,

672 "global_num_experts": config.E,

673 "apply_router_weight_on_input": config.topk == 1

@@ -684,6 +685,6 @@ def run_modular_kernel(

685 num_tokens=num_tokens,

686 num_tokens_across_dp=num_tokens_across_dp,

687 ):

687 - out = mk.forward(**mk_kwargs)

688 + out = mk.apply(**mk_kwargs)

689

690 return out

tests/kernels/moe/modular_kernel_tools/mk_objects.py

+14 -12

@@ -20,7 +20,7 @@ from vllm.model_executor.layers.fused_moe.fused_batched_moe import (

20 NaiveBatchedExperts,

21 )

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

23 - MoEPrepareAndFinalizeNoEP,

23 + MoEPrepareAndFinalizeNoDPEPModular,

24 )

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

26 TritonOrDeepGemmExperts,

@@ -71,12 +71,14 @@ class ExpertInfo:

71 needs_aiter: bool = False

74 -PREPARE_FINALIZE_INFO: dict[mk.FusedMoEPrepareAndFinalize, PrepareFinalizeInfo] = {}

75 -EXPERT_INFO: dict[mk.FusedMoEPermuteExpertsUnpermute, ExpertInfo] = {}

76 -MK_ALL_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalize] = []

77 -MK_MULTI_GPU_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalize] = []

78 -MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalize] = []

79 -MK_FUSED_EXPERT_TYPES: list[mk.FusedMoEPermuteExpertsUnpermute] = []

74 +PREPARE_FINALIZE_INFO: dict[

75 + mk.FusedMoEPrepareAndFinalizeModular, PrepareFinalizeInfo

76 +] = {}

77 +EXPERT_INFO: dict[mk.FusedMoEExpertsModular, ExpertInfo] = {}

78 +MK_ALL_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalizeModular] = []

79 +MK_MULTI_GPU_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalizeModular] = []

80 +MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalizeModular] = []

81 +MK_FUSED_EXPERT_TYPES: list[mk.FusedMoEExpertsModular] = []

83 standard_format = mk.FusedMoEActivationFormat.Standard

84 batched_format = mk.FusedMoEActivationFormat.BatchedExperts

@@ -162,7 +164,7 @@ def expert_info(kind) -> ExpertInfo:

164

165

166 register_prepare_and_finalize(

165 - MoEPrepareAndFinalizeNoEP,

167 + MoEPrepareAndFinalizeNoDPEPModular,

168 standard_format,

169 common_float_types,

170 blocked_quantization_support=True,

@@ -239,14 +241,14 @@ if has_mori():

241

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

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

242 - FlashInferCutlassMoEPrepareAndFinalize,

244 + FlashInferA2APrepareAndFinalize,

245 )

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

247 FlashInferExperts,

248 )

249

250 register_prepare_and_finalize(

249 - FlashInferCutlassMoEPrepareAndFinalize,

251 + FlashInferA2APrepareAndFinalize,

252 standard_format,

253 nvfp4_types + fp8_types,

254 blocked_quantization_support=True,

@@ -430,12 +432,12 @@ def make_cutlass_strides(

432

433

434 def make_fused_experts(

433 - fused_experts_type: mk.FusedMoEPermuteExpertsUnpermute,

435 + fused_experts_type: mk.FusedMoEExpertsModular,

436 moe: FusedMoEConfig,

437 quant_config: FusedMoEQuantConfig,

438 num_dispatchers: int,

439 N: int,

438 -) -> mk.FusedMoEPermuteExpertsUnpermute:

440 +) -> mk.FusedMoEExpertsModular:

441 if (

442 fused_experts_type.activation_format()

443 == mk.FusedMoEActivationFormat.BatchedExperts

tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py

+1 -1

@@ -72,7 +72,7 @@ def profile_modular_kernel(

72 "apply_router_weight_on_input": config.topk == 1,

73 }

75 - do_profile(mk.forward, mk_kwargs, pgi, config)

75 + do_profile(mk.apply, mk_kwargs, pgi, config)

78 def rank_worker(

tests/kernels/moe/test_batched_deepgemm.py

+12 -5

@@ -4,6 +4,7 @@

4 import pytest

5 import torch

7 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

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

9 BatchedDeepGemmExperts,

10 )

@@ -12,7 +13,7 @@ from vllm.model_executor.layers.fused_moe.fused_batched_moe import (

13 BatchedPrepareAndFinalize,

14 BatchedTritonExperts,

15 )

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

16 +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel

17 from vllm.utils.deep_gemm import calc_diff, is_deep_gemm_supported

19 from .test_deepgemm import make_block_quant_fp8_weights

@@ -74,19 +75,22 @@ def test_batched_deepgemm_vs_triton(

75 quant_config=quant_config,

76 moe_config=make_dummy_moe_config(),

77 )

77 - mk_triton = FusedMoEModularKernel(

78 + mk_triton = FusedMoEKernel(

79 prep_finalize,

80 triton_experts,

81 inplace=False,

82 )

83 - out_triton = mk_triton(

84 + out_triton = mk_triton.apply(

85 hidden_states=a,

86 w1=w1,

87 w2=w2,

88 topk_weights=topk_weights,

89 topk_ids=topk_ids,

90 + activation=MoEActivation.SILU,

91 global_num_experts=E,

92 + expert_map=None,

93 + apply_router_weight_on_input=False,

94 )

96 # deepgemm

@@ -96,19 +100,22 @@ def test_batched_deepgemm_vs_triton(

100 quant_config=quant_config,

101 moe_config=make_dummy_moe_config(),

102 )

99 - mk_deepgemm = FusedMoEModularKernel(

103 + mk_deepgemm = FusedMoEKernel(

104 prep_finalize,

105 deepgemm_experts,

106 inplace=False,

107 )

108

105 - out_deepgemm = mk_deepgemm(

109 + out_deepgemm = mk_deepgemm.apply(

110 hidden_states=a,

111 w1=w1,

112 w2=w2,

113 topk_weights=topk_weights,

114 topk_ids=topk_ids,

115 + activation=MoEActivation.SILU,

116 global_num_experts=E,

117 + expert_map=None,

118 + apply_router_weight_on_input=False,

119 )

120

121 diff = calc_diff(out_deepgemm, out_triton)

tests/kernels/moe/test_block_fp8.py

+29 -8

@@ -21,15 +21,16 @@ from vllm.model_executor.layers.fused_moe import (

21 fused_experts,

22 fused_topk,

23 )

24 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

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 fp8_w8a8_moe_quant_config,

30 )

31 from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (

32 _valid_deep_gemm_shape,

33 )

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

31 - MoEPrepareAndFinalizeNoEP,

32 -)

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

35 TritonOrDeepGemmExperts,

36 )

@@ -193,7 +194,17 @@ def test_w8a8_block_fp8_fused_moe(

194 a, w1, w2, topk_weights, topk_ids, quant_config=quant_config

195 )

196

196 - m_out = m_fused_moe(a, w1, w2, topk_weights, topk_ids)

197 + m_out = m_fused_moe.apply(

198 + a,

199 + w1,

200 + w2,

201 + topk_weights,

202 + topk_ids,

203 + activation=MoEActivation.SILU,

204 + apply_router_weight_on_input=False,

205 + expert_map=None,

206 + global_num_experts=w1.shape[0],

207 + )

208

209 # 0.039 only needed for M >= 8192

210 tol = 0.035 if M < 8192 else 0.039

@@ -252,23 +263,33 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed, monkeypatch)

263 w2_scale=w2_s,

264 block_shape=block_size,

265 )

266 + moe_config = make_dummy_moe_config()

267

256 - deep_gemm_experts = mk.FusedMoEModularKernel(

257 - prepare_finalize=MoEPrepareAndFinalizeNoEP(),

268 + deep_gemm_experts = mk.FusedMoEKernel(

269 + prepare_finalize=maybe_make_prepare_finalize(

270 + moe=moe_config,

271 + quant_config=quant_config,

272 + allow_new_interface=True,

273 + use_monolithic=False,

274 + ),

275 fused_experts=TritonOrDeepGemmExperts(

259 - moe_config=make_dummy_moe_config(),

276 + moe_config=moe_config,

277 quant_config=quant_config,

278 ),

279 inplace=False,

280 )

281

282 def deep_gemm_moe_fp8(a, w1, w2, w1_s, w2_s, topk_weights, topk_ids):

266 - return deep_gemm_experts(

283 + return deep_gemm_experts.apply(

284 hidden_states=a,

285 w1=w1,

286 w2=w2,

287 topk_weights=topk_weights,

288 topk_ids=topk_ids,

289 + global_num_experts=E,

290 + activation=MoEActivation.SILU,

291 + apply_router_weight_on_input=False,

292 + expert_map=False,

293 )

294

295 # Set the context to avoid lots of warning spam.

tests/kernels/moe/test_cutlass_moe.py

+37 -21

@@ -13,6 +13,9 @@ from vllm import _custom_ops as ops

13 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config

14 from vllm.model_executor.layers.fused_moe import fused_experts, fused_topk

15 from vllm.model_executor.layers.fused_moe.activation import MoEActivation

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

17 + maybe_make_prepare_finalize,

18 +)

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

20 FUSED_MOE_UNQUANTIZED_CONFIG,

21 FusedMoEQuantConfig,

@@ -22,9 +25,6 @@ from vllm.model_executor.layers.fused_moe.cutlass_moe import (

25 CutlassExpertsFp8,

26 run_cutlass_moe_fp8,

27 )

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

26 - MoEPrepareAndFinalizeNoEP,

27 -)

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

29 from vllm.platforms import current_platform

30 from vllm.utils.torch_utils import set_random_seed

@@ -197,20 +197,26 @@ def run_with_expert_maps(

197 for kwargs, new_quant_config in slice_experts():

198 w2 = kwargs["w2"]

199 a = kwargs["hidden_states"]

200 - kernel = mk.FusedMoEModularKernel(

201 - MoEPrepareAndFinalizeNoEP(),

200 + moe_config = make_dummy_moe_config(

201 + num_experts=w2.shape[0],

202 + hidden_dim=w2.shape[1],

203 + intermediate_size_per_partition=w2.shape[2],

204 + in_dtype=a.dtype,

205 + )

206 + kernel = mk.FusedMoEKernel(

207 + maybe_make_prepare_finalize(

208 + moe=moe_config,

209 + quant_config=new_quant_config,

210 + allow_new_interface=True,

211 + use_monolithic=False,

212 + ),

213 CutlassExpertsFp8(

203 - moe_config=make_dummy_moe_config(

204 - num_experts=w2.shape[0],

205 - hidden_dim=w2.shape[1],

206 - intermediate_size_per_partition=w2.shape[2],

207 - in_dtype=a.dtype,

208 - ),

214 + moe_config=moe_config,

215 quant_config=new_quant_config,

216 ),

217 inplace=False,

218 )

213 - out_tensor = out_tensor + kernel(**kwargs)

219 + out_tensor = out_tensor + kernel.apply(**kwargs)

220

221 return out_tensor

222

@@ -252,25 +258,35 @@ def run_8_bit(

258 "w2": moe_tensors.w2_q, # type: ignore[union-attr]

259 "topk_weights": topk_weights,

260 "topk_ids": topk_ids,

261 + "global_num_experts": moe_tensors.w1_q.shape[0], # type: ignore[union-attr]

262 + "activation": MoEActivation.SILU,

263 + "expert_map": None,

264 + "apply_router_weight_on_input": False,

265 }

266

267 num_experts = moe_tensors.w1.size(0) # type: ignore[attr-defined]

268 with_ep = num_local_experts is not None or num_local_experts == num_experts

269 if not with_ep:

260 - kernel = mk.FusedMoEModularKernel(

261 - MoEPrepareAndFinalizeNoEP(),

270 + moe_config = make_dummy_moe_config(

271 + num_experts=moe_tensors.w2_q.shape[0], # type: ignore[union-attr]

272 + hidden_dim=moe_tensors.w2_q.shape[1], # type: ignore[union-attr]

273 + intermediate_size_per_partition=moe_tensors.w2_q.shape[2], # type: ignore[union-attr]

274 + in_dtype=moe_tensors.a.dtype,

275 + )

276 + kernel = mk.FusedMoEKernel(

277 + maybe_make_prepare_finalize(

278 + moe=moe_config,

279 + quant_config=quant_config,

280 + allow_new_interface=True,

281 + use_monolithic=False,

282 + ),

283 CutlassExpertsFp8(

263 - moe_config=make_dummy_moe_config(

264 - num_experts=moe_tensors.w2_q.shape[0], # type: ignore[union-attr]

265 - hidden_dim=moe_tensors.w2_q.shape[1], # type: ignore[union-attr]

266 - intermediate_size_per_partition=moe_tensors.w2_q.shape[2], # type: ignore[union-attr]

267 - in_dtype=moe_tensors.a.dtype,

268 - ),

284 + moe_config=moe_config,

285 quant_config=quant_config,

286 ),

287 inplace=False,

288 )

273 - return kernel(**kwargs)

289 + return kernel.apply(**kwargs)

290

291 assert num_local_experts is not None

292 return run_with_expert_maps(

tests/kernels/moe/test_deepep_deepgemm_moe.py

+9 -9

@@ -22,7 +22,7 @@ from vllm.model_executor.layers.fused_moe.config import (

22 fp8_w8a8_moe_quant_config,

23 )

24 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts

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

25 +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel

26 from vllm.utils.deep_gemm import (

27 get_mk_alignment_for_contiguous_layout,

28 is_deep_gemm_e8m0_used,

@@ -170,7 +170,7 @@ def make_ll_modular_kernel(

170 q_dtype: torch.dtype | None,

171 test_config: TestConfig,

172 quant_config: FusedMoEQuantConfig,

173 -) -> FusedMoEModularKernel:

173 +) -> FusedMoEKernel:

174 assert test_config.low_latency

175 assert test_config.use_fp8_dispatch is not None

176

@@ -195,7 +195,7 @@ def make_ll_modular_kernel(

195 quant_config=quant_config,

196 moe_config=make_dummy_moe_config(),

197 )

198 - return FusedMoEModularKernel(

198 + return FusedMoEKernel(

199 prepare_finalize=a2a,

200 fused_experts=fused_experts,

201 inplace=False,

@@ -210,7 +210,7 @@ def make_ht_modular_kernel(

210 q_dtype: torch.dtype | None,

211 test_config: TestConfig,

212 quant_config: FusedMoEQuantConfig,

213 -) -> FusedMoEModularKernel:

213 +) -> FusedMoEKernel:

214 assert not test_config.low_latency

215 assert test_config.use_fp8_dispatch is None

216

@@ -228,7 +228,7 @@ def make_ht_modular_kernel(

228 moe_config=make_dummy_moe_config(),

229 quant_config=quant_config,

230 )

231 - return FusedMoEModularKernel(

231 + return FusedMoEKernel(

232 prepare_finalize=a2a,

233 fused_experts=fused_experts,

234 inplace=False,

@@ -242,11 +242,11 @@ def make_modular_kernel(

242 num_local_experts: int,

243 test_tensors: TestTensors,

244 quant_config: FusedMoEQuantConfig,

245 -) -> FusedMoEModularKernel:

245 +) -> FusedMoEKernel:

246 q_dtype = torch.float8_e4m3fn

247 test_config = test_tensors.config

248

249 - mk: FusedMoEModularKernel

249 + mk: FusedMoEKernel

250 # Make modular kernel

251 if test_config.low_latency:

252 max_tokens_per_rank = max(64, next_power_of_2(test_tensors.rank_tokens.size(0)))

@@ -307,7 +307,7 @@ def deepep_deepgemm_moe_impl(

307 )

308

309 # Make modular kernel

310 - mk: FusedMoEModularKernel = make_modular_kernel(

310 + mk: FusedMoEKernel = make_modular_kernel(

311 pg=pg,

312 pgi=pgi,

313 dp_size=dp_size,

@@ -319,7 +319,7 @@ def deepep_deepgemm_moe_impl(

319 with with_dp_metadata(

320 M=test_tensors.rank_tokens.size(0), world_size=pgi.world_size

321 ):

322 - out = mk.forward(

322 + out = mk.apply(

323 hidden_states=test_tensors.rank_tokens,

324 w1=w1,

325 w2=w2,

tests/kernels/moe/test_deepep_moe.py

+5 -5

@@ -20,7 +20,7 @@ from vllm.model_executor.layers.fused_moe.config import (

20 FusedMoEQuantConfig,

21 )

22 from vllm.model_executor.layers.fused_moe.fused_batched_moe import BatchedTritonExperts

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

23 +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel

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

25 per_token_group_quant_fp8,

26 )

@@ -135,7 +135,7 @@ def make_modular_kernel(

135 q_dtype: torch.dtype | None,

136 use_fp8_dispatch: bool,

137 quant_config: FusedMoEQuantConfig,

138 -) -> FusedMoEModularKernel:

138 +) -> FusedMoEKernel:

139 ht_args: DeepEPHTArgs | None = None

140 ll_args: DeepEPLLArgs | None = None

141

@@ -180,7 +180,7 @@ def make_modular_kernel(

180 quant_config=quant_config,

181 )

182

183 - mk = FusedMoEModularKernel(

183 + mk = FusedMoEKernel(

184 prepare_finalize=a2a,

185 fused_experts=fused_experts,

186 inplace=False,

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

242 )

243

244 # Make modular kernel

245 - mk: FusedMoEModularKernel = make_modular_kernel(

245 + mk: FusedMoEKernel = make_modular_kernel(

246 pg,

247 pgi,

248 low_latency_mode,

@@ -255,7 +255,7 @@ def deep_ep_moe_impl(

255 quant_config,

256 )

257

258 - out = mk.forward(

258 + out = mk.apply(

259 hidden_states=rank_tokens_chunk,

260 w1=w1,

261 w2=w2,

tests/kernels/moe/test_deepgemm.py

+20 -7

@@ -14,13 +14,16 @@ import torch

14 # vLLM fused-expert reference (Triton fallback + DeepGEMM option)

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

16 from tests.kernels.moe.utils import make_dummy_moe_config

17 +from vllm.model_executor.layers.fused_moe.activation import (

18 + MoEActivation,

19 +)

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

21 + maybe_make_prepare_finalize,

22 +)

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

24 fp8_w8a8_moe_quant_config,

25 )

26 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts

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

22 - MoEPrepareAndFinalizeNoEP,

23 -)

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

28 TritonOrDeepGemmExperts,

29 )

@@ -108,11 +111,17 @@ def run_single_case(m, n, k, topk, num_experts, block_size):

111 a1_scale=a1_scale,

112 block_shape=block_size,

113 )

114 + moe_config = make_dummy_moe_config()

115

112 - deep_gemm_experts = mk.FusedMoEModularKernel(

113 - prepare_finalize=MoEPrepareAndFinalizeNoEP(),

116 + deep_gemm_experts = mk.FusedMoEKernel(

117 + prepare_finalize=maybe_make_prepare_finalize(

118 + moe=moe_config,

119 + quant_config=quant_config,

120 + allow_new_interface=True,

121 + use_monolithic=False,

122 + ),

123 fused_experts=TritonOrDeepGemmExperts(

115 - moe_config=make_dummy_moe_config(),

124 + moe_config=moe_config,

125 quant_config=quant_config,

126 ),

127 inplace=False,

@@ -130,12 +139,16 @@ def run_single_case(m, n, k, topk, num_experts, block_size):

139 )

140

141 # DeepGemm

133 - out_deepgemm = deep_gemm_experts(

142 + out_deepgemm = deep_gemm_experts.apply(

143 hidden_states=tokens_bf16,

144 w1=w1,

145 w2=w2,

146 topk_weights=topk_weights,

147 topk_ids=topk_ids,

148 + global_num_experts=num_experts,

149 + activation=MoEActivation.SILU,

150 + apply_router_weight_on_input=False,

151 + expert_map=None,

152 )

153 diff = calc_diff(out_deepgemm, out_triton)

154 assert diff < 0.001, f"Diff exceeded 1%: {diff}"

tests/kernels/moe/test_flashinfer.py

+49 -21

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

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

9 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config

10 from vllm.model_executor.layers.fused_moe.activation import MoEActivation

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 FusedMoEParallelConfig,

@@ -15,16 +18,14 @@ from vllm.model_executor.layers.fused_moe.config import (

18 RoutingMethodType,

19 fp8_w8a8_moe_quant_config,

20 )

21 +from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import (

22 + TrtLlmFp8Experts,

23 +)

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

25 FlashInferExperts,

26 )

27 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts

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

23 - MoEPrepareAndFinalizeNoEP,

24 -)

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

26 - apply_fi_trtllm_fp8_per_tensor_moe,

27 - register_scales_for_trtllm_fp8_per_tensor_moe,

29 rotate_weights_for_fi_trtllm_fp8_per_tensor_moe,

30 swap_w13_to_w31,

31 )

@@ -115,6 +116,7 @@ class TestData:

116 e: int,

117 is_trtllm: bool,

118 activation: MoEActivation = MoEActivation.SILU,

119 + topk: int = 1,

120 ) -> "TestData":

121 is_gated = activation.is_gated

122

@@ -152,13 +154,6 @@ class TestData:

154 rotate_weights_for_fi_trtllm_fp8_per_tensor_moe(

155 layer.w13_weight, layer.w2_weight, is_gated

156 )

155 - register_scales_for_trtllm_fp8_per_tensor_moe(

156 - layer,

157 - layer.w13_weight_scale,

158 - layer.w13_input_scale,

159 - layer.w2_weight_scale,

160 - layer.w2_input_scale,

161 - )

157 layer.custom_routing_function = Llama4MoE.custom_routing_function

158 layer.routing_method_type = RoutingMethodType.Llama4

159 layer.renormalize = False

@@ -166,6 +161,21 @@ class TestData:

161 layer.ep_rank = 0

162 layer.local_num_experts = e

163

164 + layer.moe = FusedMoEConfig(

165 + num_experts=e,

166 + experts_per_token=topk,

167 + hidden_dim=k,

168 + intermediate_size_per_partition=n,

169 + num_local_experts=e,

170 + num_logical_experts=e,

171 + moe_parallel_config=layer.moe_parallel_config,

172 + in_dtype=hidden_states.dtype,

173 + is_act_and_mul=is_gated,

174 + routing_method=layer.routing_method_type,

175 + activation=activation,

176 + device=w13_quantized.device,

177 + )

178 +

179 return TestData(

180 hidden_states=hidden_states,

181 w13_quantized=w13_quantized,

@@ -230,16 +240,29 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph(

240 quant_config=quant_config,

241 )

242

233 - flashinfer_output = apply_fi_trtllm_fp8_per_tensor_moe(

234 - layer=td.layer,

243 + kernel = mk.FusedMoEKernel(

244 + maybe_make_prepare_finalize(

245 + moe=td.layer.moe,

246 + quant_config=quant_config,

247 + allow_new_interface=True,

248 + use_monolithic=True,

249 + ),

250 + TrtLlmFp8Experts(

251 + moe_config=td.layer.moe,

252 + quant_config=quant_config,

253 + ),

254 + )

255 +

256 + flashinfer_output = kernel.apply_monolithic(

257 hidden_states=td.hidden_states,

258 + w1=td.layer.w13_weight,

259 + w2=td.layer.w2_weight,

260 router_logits=score,

237 - routing_bias=None,

261 + activation=activation,

262 global_num_experts=e,

239 - top_k=topk,

240 - num_expert_group=None,

241 - topk_group=None,

263 + expert_map=None,

264 apply_router_weight_on_input=True,

265 + routed_scaling_factor=1.0,

266 )

267

268 check_accuracy(

@@ -329,8 +352,13 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(

352 routing_method=RoutingMethodType.TopK,

353 )

354

332 - kernel = mk.FusedMoEModularKernel(

333 - MoEPrepareAndFinalizeNoEP(),

355 + kernel = mk.FusedMoEKernel(

356 + maybe_make_prepare_finalize(

357 + moe=moe_config,

358 + quant_config=quant_config,

359 + allow_new_interface=True,

360 + use_monolithic=False,

361 + ),

362 FlashInferExperts(

363 moe_config=moe_config,

364 quant_config=quant_config,

@@ -338,7 +366,7 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(

366 inplace=False,

367 )

368

341 - flashinfer_cutlass_output = kernel(

369 + flashinfer_cutlass_output = kernel.apply(

370 td.hidden_states,

371 td.layer.w13_weight,

372 td.layer.w2_weight,

tests/kernels/moe/test_flashinfer_moe.py

+15 -7

@@ -14,6 +14,9 @@ from vllm import _custom_ops as ops

14 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config

15 from vllm.model_executor.layers.fused_moe import fused_topk

16 from vllm.model_executor.layers.fused_moe.activation import MoEActivation

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

18 + maybe_make_prepare_finalize,

19 +)

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

21 FusedMoEConfig,

22 FusedMoEParallelConfig,

@@ -23,10 +26,7 @@ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (

26 FlashInferExperts,

27 is_valid_flashinfer_cutlass_fused_moe,

28 )

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

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

28 - MoEPrepareAndFinalizeNoEP,

29 -)

29 +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel

30 from vllm.platforms import current_platform

31 from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe

32 from vllm.utils.torch_utils import set_random_seed

@@ -107,19 +107,27 @@ def test_flashinfer_fp4_moe_no_graph(

107 routing_method=RoutingMethodType.TopK,

108 )

109

110 - flashinfer_experts = FusedMoEModularKernel(

111 - MoEPrepareAndFinalizeNoEP(),

110 + flashinfer_experts = FusedMoEKernel(

111 + maybe_make_prepare_finalize(

112 + moe=moe_config,

113 + quant_config=quant_config,

114 + allow_new_interface=True,

115 + use_monolithic=False,

116 + ),

117 FlashInferExperts(moe_config=moe_config, quant_config=quant_config),

118 inplace=False,

119 )

120

116 - flashinfer_output = flashinfer_experts(

121 + flashinfer_output = flashinfer_experts.apply(

122 hidden_states=a,

123 w1=w1_q,

124 w2=w2_q,

125 topk_weights=topk_weights,

126 topk_ids=topk_ids,

127 activation=activation,

128 + global_num_experts=e,

129 + expert_map=None,

130 + apply_router_weight_on_input=False,

131 )

132

133 # Reference check:

tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py

+10 -11

@@ -221,16 +221,16 @@ def test_marlin_vs_trtllm_mxint4_moe_kimik2(monkeypatch, m, n, k, e, topk, group

221 )

222

223 marlin_output = fused_marlin_moe(

224 - a,

225 - w1_marlin,

226 - w2_marlin,

227 - None,

228 - None,

229 - w1_scales_marlin,

230 - w2_scales_marlin,

231 - None, # gating_output not needed when topk_weights/ids provided

232 - topk_weights,

233 - topk_ids,

224 + hidden_states=a,

225 + w1=w1_marlin,

226 + w2=w2_marlin,

227 + bias1=None,

228 + bias2=None,

229 + w1_scale=w1_scales_marlin,

230 + w2_scale=w2_scales_marlin,

231 + topk_weights=topk_weights,

232 + topk_ids=topk_ids,

233 + quant_type_id=scalar_types.uint4b8.id,

234 global_num_experts=e,

235 expert_map=None,

236 global_scale1=None,

@@ -244,7 +244,6 @@ def test_marlin_vs_trtllm_mxint4_moe_kimik2(monkeypatch, m, n, k, e, topk, group

244 w1_zeros=None,

245 w2_zeros=None,

246 input_dtype=dtype,

247 - quant_type_id=scalar_types.uint4b8.id,

247 is_k_full=True,

248 )

249

tests/kernels/moe/test_modular_kernel_combinations.py

+2 -3

@@ -168,7 +168,6 @@ FUSED_MOE_CHUNK_SIZEs = [None, 16]

168 def is_nyi_config(config: Config) -> bool:

169 # We know these configs to be legitimate. but still fail.

170 info = expert_info(config.fused_experts_type)

171 -

171 if info.needs_matching_quant:

172 # The triton kernels expect both per-act-token-quant and

173 # per-out-ch-quant or neither.

@@ -259,7 +258,7 @@ def test_modular_kernel_combinations_multigpu(

258 dtype: torch.dtype,

259 quant_config: TestMoEQuantConfig | None,

260 prepare_finalize_type: mk.FusedMoEPrepareAndFinalize,

262 - fused_experts_type: mk.FusedMoEPermuteExpertsUnpermute,

261 + fused_experts_type: mk.FusedMoEExperts,

262 chunk_size: int | None,

263 world_size: int,

264 pytestconfig,

@@ -301,7 +300,7 @@ def test_modular_kernel_combinations_singlegpu(

300 dtype: torch.dtype,

301 quant_config: TestMoEQuantConfig | None,

302 prepare_finalize_type: mk.FusedMoEPrepareAndFinalize,

304 - fused_experts_type: mk.FusedMoEPermuteExpertsUnpermute,

303 + fused_experts_type: mk.FusedMoEExperts,

304 chunk_size: int | None,

305 world_size: int,

306 pytestconfig,

tests/kernels/moe/test_modular_oai_triton_moe.py

+18 -10

@@ -7,6 +7,7 @@ Test modular OAI Triton MoE

7 import pytest

8 import torch

10 +from tests.utils import wait_for_gpu_memory_to_clear

11 from vllm.model_executor.layers.fused_moe.activation import MoEActivation

12 from vllm.utils.import_utils import has_triton_kernels

@@ -24,15 +25,15 @@ from triton_kernels.tensor_details import layout

25 from triton_kernels.testing import assert_close

27 from vllm.config import VllmConfig, set_current_vllm_config

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

29 + maybe_make_prepare_finalize,

30 +)

31 from vllm.model_executor.layers.fused_moe.config import mxfp4_w4a16_moe_quant_config

32 from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (

33 OAITritonExperts,

34 UnfusedOAITritonExperts,

35 )

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

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

34 - MoEPrepareAndFinalizeNoEP,

35 -)

36 +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel

37 from vllm.platforms import current_platform

38 from vllm.utils.torch_utils import set_random_seed

@@ -174,19 +175,25 @@ def oai_triton_moe_impl(

175 w1_scale=w1_scale,

176 w2_scale=w2_scale,

177 )

178 + moe_config = make_dummy_moe_config()

179

180 if unfused:

179 - fused_experts = UnfusedOAITritonExperts(make_dummy_moe_config(), quant_config)

181 + fused_experts = UnfusedOAITritonExperts(moe_config, quant_config)

182 else:

181 - fused_experts = OAITritonExperts(make_dummy_moe_config(), quant_config)

182 -

183 - mk = FusedMoEModularKernel(

184 - MoEPrepareAndFinalizeNoEP(),

183 + fused_experts = OAITritonExperts(moe_config, quant_config)

184 +

185 + mk = FusedMoEKernel(

186 + maybe_make_prepare_finalize(

187 + moe=moe_config,

188 + quant_config=quant_config,

189 + allow_new_interface=True,

190 + use_monolithic=False,

191 + ),

192 fused_experts,

193 inplace=False,

194 )

195

189 - return mk.forward(

196 + return mk.apply(

197 hidden_states=x,

198 w1=w1,

199 w2=w2,

@@ -217,6 +224,7 @@ def test_oai_triton_moe(

224 unfused: bool,

225 workspace_init,

226 ):

227 + wait_for_gpu_memory_to_clear(devices=[0], threshold_ratio=0.1)

228 set_random_seed(0)

229 (

230 w1,

tests/kernels/moe/test_moe.py

+6 -2

@@ -346,14 +346,16 @@ def test_fused_moe(

346 expert_map: torch.Tensor | None = None,

347 ) -> torch.Tensor:

348 topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)

349 - return m_fused_moe_fn(

349 + return m_fused_moe_fn.apply(

350 a,

351 w1,

352 w2,

353 topk_weights,

354 topk_ids,

355 + activation=MoEActivation.SILU,

356 global_num_experts=global_num_experts,

357 expert_map=expert_map,

358 + apply_router_weight_on_input=False,

359 )

360

361 fused_moe_fn = functools.partial(fused_moe, renormalize=False)

@@ -500,14 +502,16 @@ def test_naive_block_assignment_moe(

502 expert_map: torch.Tensor | None = None,

503 ) -> torch.Tensor:

504 topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)

503 - return m_fused_moe_fn(

505 + return m_fused_moe_fn.apply(

506 a,

507 w1,

508 w2,

509 topk_weights,

510 topk_ids,

511 + activation=MoEActivation.SILU,

512 global_num_experts=global_num_experts,

513 expert_map=expert_map,

514 + apply_router_weight_on_input=False,

515 )

516

517 fused_moe_fn = functools.partial(fused_moe, renormalize=False)

tests/kernels/moe/test_nvfp4_moe.py

+24 -8

@@ -15,12 +15,15 @@ from vllm import _custom_ops as ops

15 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config

16 from vllm.model_executor.layers.fused_moe import fused_topk

17 from vllm.model_executor.layers.fused_moe.activation import MoEActivation

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

19 + maybe_make_prepare_finalize,

20 +)

21 from vllm.model_executor.layers.fused_moe.config import nvfp4_moe_quant_config

22 from vllm.model_executor.layers.fused_moe.cutlass_moe import (

23 CutlassExpertsFp4,

24 )

25 from vllm.model_executor.layers.fused_moe.prepare_finalize import (

23 - MoEPrepareAndFinalizeNoEP,

26 + make_moe_prepare_and_finalize_no_dp_ep,

27 )

28 from vllm.platforms import current_platform

29 from vllm.utils.torch_utils import set_random_seed

@@ -89,22 +92,32 @@ def test_cutlass_fp4_moe_no_graph(

92 w1_scale=w1_blockscale,

93 w2_scale=w2_blockscale,

94 )

95 + moe_config = make_dummy_moe_config()

93 - kernel = mk.FusedMoEModularKernel(

94 - MoEPrepareAndFinalizeNoEP(),

97 + kernel = mk.FusedMoEKernel(

98 + maybe_make_prepare_finalize(

99 + moe=moe_config,

100 + quant_config=quant_config,

101 + allow_new_interface=True,

102 + use_monolithic=False,

103 + ),

104 CutlassExpertsFp4(

96 - moe_config=make_dummy_moe_config(),

105 + moe_config=moe_config,

106 quant_config=quant_config,

107 ),

108 inplace=False,

109 )

110

102 - cutlass_output = kernel(

111 + cutlass_output = kernel.apply(

112 hidden_states=a,

113 w1=w1_q,

114 w2=w2_q,

115 topk_weights=topk_weights,

116 topk_ids=topk_ids,

117 + global_num_experts=e,

118 + activation=mk.MoEActivation.SILU,

119 + apply_router_weight_on_input=False,

120 + expert_map=None,

121 )

122

123 # Reference check:

@@ -207,8 +220,8 @@ def test_cutlass_fp4_moe_swiglustep(

220 w2_scale=w2_blockscale,

221 )

222

210 - kernel = mk.FusedMoEModularKernel(

211 - MoEPrepareAndFinalizeNoEP(),

223 + kernel = mk.FusedMoEKernel(

224 + make_moe_prepare_and_finalize_no_dp_ep(use_monolithic=False),

225 CutlassExpertsFp4(

226 moe_config=make_dummy_moe_config(),

227 quant_config=quant_config,

@@ -216,13 +229,16 @@ def test_cutlass_fp4_moe_swiglustep(

229 inplace=False,

230 )

231

219 - cutlass_output = kernel(

232 + cutlass_output = kernel.apply(

233 hidden_states=a,

234 w1=w1_q,

235 w2=w2_q,

236 topk_weights=topk_weights,

237 topk_ids=topk_ids,

238 activation=MoEActivation.SWIGLUSTEP,

239 + global_num_experts=e,

240 + expert_map=None,

241 + apply_router_weight_on_input=False,

242 )

243

244 # Reference: dequantize everything and run torch_moe with swiglustep

tests/kernels/moe/utils.py

+41 -13

@@ -8,6 +8,9 @@ from tests.kernels.quant_utils import per_block_cast_to_int8

8 from tests.kernels.quantization.nvfp4_utils import FLOAT4_E2M1_MAX, FLOAT8_E4M3_MAX

9 from vllm.model_executor.layers.activation import SiluAndMul

10 from vllm.model_executor.layers.fused_moe.activation import MoEActivation

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 FusedMoEParallelConfig,

@@ -23,10 +26,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe import (

26 TritonExperts,

27 fused_experts,

28 )

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

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

28 - MoEPrepareAndFinalizeNoEP,

29 -)

29 +from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel

30 from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk

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

32 from vllm.utils.deep_gemm import per_block_cast_to_fp8

@@ -125,7 +125,9 @@ def batched_moe(

125 a2_scale=a2_scale,

126 )

127

128 - fused_experts = FusedMoEModularKernel(

128 + moe_config = make_dummy_moe_config()

129 +

130 + fused_experts = FusedMoEKernel(

131 BatchedPrepareAndFinalize(

132 max_num_tokens, num_dispatchers=1, num_local_experts=w1.shape[0], rank=0

133 ),

@@ -133,12 +135,22 @@ def batched_moe(

135 max_num_tokens=max_num_tokens,

136 num_dispatchers=1,

137 quant_config=quant_config,

136 - moe_config=make_dummy_moe_config(),

138 + moe_config=moe_config,

139 ),

140 inplace=False,

141 )

142

141 - return fused_experts(a, w1, w2, topk_weight, topk_ids)

143 + return fused_experts.apply(

144 + a,

145 + w1,

146 + w2,

147 + topk_weight,

148 + topk_ids,

149 + global_num_experts=w1.shape[0],

150 + activation=moe_config.activation,

151 + apply_router_weight_on_input=False,

152 + expert_map=None,

153 + )

154

155

156 def naive_batched_moe(

@@ -166,8 +178,9 @@ def naive_batched_moe(

178 a1_scale=a1_scale,

179 a2_scale=a2_scale,

180 )

181 + moe_config = make_dummy_moe_config()

182

170 - fused_experts = FusedMoEModularKernel(

183 + fused_experts = FusedMoEKernel(

184 BatchedPrepareAndFinalize(

185 max_num_tokens, num_dispatchers=1, num_local_experts=w1.shape[0], rank=0

186 ),

@@ -175,12 +188,22 @@ def naive_batched_moe(

188 max_num_tokens=max_num_tokens,

189 num_dispatchers=1,

190 quant_config=quant_config,

178 - moe_config=make_dummy_moe_config(),

191 + moe_config=moe_config,

192 ),

193 inplace=False,

194 )

195

183 - return fused_experts(a, w1, w2, topk_weight, topk_ids)

196 + return fused_experts.apply(

197 + a,

198 + w1,

199 + w2,

200 + topk_weight,

201 + topk_ids,

202 + global_num_experts=w1.shape[0],

203 + activation=moe_config.activation,

204 + apply_router_weight_on_input=False,

205 + expert_map=None,

206 + )

207

208

209 def chunk_scales(

@@ -581,9 +604,14 @@ def modular_triton_fused_moe(

604 moe_config: FusedMoEConfig,

605 quant_config: FusedMoEQuantConfig,

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

584 -) -> FusedMoEModularKernel:

585 - return FusedMoEModularKernel(

586 - MoEPrepareAndFinalizeNoEP(),

607 +) -> FusedMoEKernel:

608 + return FusedMoEKernel(

609 + maybe_make_prepare_finalize(

610 + moe=moe_config,

611 + quant_config=quant_config,

612 + allow_new_interface=True,

613 + use_monolithic=False,

614 + ),

615 TritonExperts(moe_config, quant_config),

616 shared_experts,

617 inplace=False,

tests/quantization/test_blackwell_moe.py

+80

@@ -127,6 +127,14 @@ def test_deepseek_fp8_block_moe_deep_gemm(monkeypatch: pytest.MonkeyPatch):

127 )

128

129

130 +def test_deepseek_fp8_block_moe_vllm_triton(monkeypatch: pytest.MonkeyPatch):

131 + can_initialize(

132 + "deepseek-ai/DeepSeek-V3.1",

133 + hf_overrides=HF_OVERRIDE_TEXT,

134 + extra_args=["--moe-backend=triton"],

135 + )

136 +

137 +

138 @pytest.mark.skip(

139 reason=(

140 "Known issue: lack of kernel support. "

@@ -149,6 +157,14 @@ def test_deepseek_fp8_block_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatc

157 )

158

159

160 +def test_deepseek_nvfp4_moe_flashinfer_vllm(monkeypatch: pytest.MonkeyPatch):

161 + can_initialize(

162 + "nvidia/DeepSeek-R1-0528-FP4-v2",

163 + hf_overrides=HF_OVERRIDE_TEXT,

164 + extra_args=["--moe-backend=cutlass"],

165 + )

166 +

167 +

168 def test_deepseek_nvfp4_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch):

169 can_initialize(

170 "nvidia/DeepSeek-R1-0528-FP4-v2",

@@ -200,3 +216,67 @@ def test_qwen3_next_bf16_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):

216 hf_overrides=HF_OVERRIDE_TEXT,

217 extra_args=["--moe-backend=flashinfer_trtllm"],

218 )

219 +

220 +

221 +## NemoTron ##

222 +

223 +

224 +def test_nemotron_fp8_moe_flashinfer_throughput(monkeypatch: pytest.MonkeyPatch):

225 + can_initialize(

226 + "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8",

227 + hf_overrides=HF_OVERRIDE_TEXT,

228 + extra_args=["--moe-backend=flashinfer_cutlass"],

229 + )

230 +

231 +

232 +@pytest.mark.skip(

233 + reason=(

234 + "FP8 MoE backend FLASHINFER_TRTLLM does not support the "

235 + "deployment configuration since kernel does not support "

236 + "no act_and_mul MLP layer."

237 + )

238 +)

239 +def test_nemotron_fp8_moe_flashinfer_latency(monkeypatch: pytest.MonkeyPatch):

240 + can_initialize(

241 + "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8",

242 + hf_overrides=HF_OVERRIDE_TEXT,

243 + extra_args=["--moe-backend=flashinfer_trtllm"],

244 + )

245 +

246 +

247 +@pytest.mark.skip(

248 + reason=(

249 + "FP8 MoE backend TRITON does not support the "

250 + "deployment configuration since kernel does not support "

251 + "no act_and_mul MLP layer."

252 + )

253 +)

254 +def test_nemotron_fp8_moe_vllm_triton(monkeypatch: pytest.MonkeyPatch):

255 + can_initialize(

256 + "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8",

257 + hf_overrides=HF_OVERRIDE_TEXT,

258 + extra_args=["--moe-backend=triton"],

259 + )

260 +

261 +

262 +def test_nemotron_fp4_moe_flashinfer_throughput(monkeypatch: pytest.MonkeyPatch):

263 + can_initialize(

264 + "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4",

265 + hf_overrides=HF_OVERRIDE_TEXT,

266 + extra_args=["--moe-backend=flashinfer_cutlass"],

267 + )

268 +

269 +

270 +@pytest.mark.skip(

271 + reason=(

272 + "FP4 MoE backend FLASHINFER_TRTLLM does not support the "

273 + "deployment configuration since kernel does not support "

274 + "hidden_dim % 512 != 0."

275 + )

276 +)

277 +def test_nemotron_fp4_moe_flashinfer_latency(monkeypatch: pytest.MonkeyPatch):

278 + can_initialize(

279 + "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4",

280 + hf_overrides=HF_OVERRIDE_TEXT,

281 + extra_args=["--moe-backend=flashinfer_trtllm"],

282 + )

vllm/lora/layers/fused_moe.py

+10 -9

@@ -32,10 +32,10 @@ from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (

32 UnfusedOAITritonExperts,

33 )

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

35 - FusedMoEModularKernel,

35 + FusedMoEKernel,

36 )

37 from vllm.model_executor.layers.fused_moe.prepare_finalize import (

38 - MoEPrepareAndFinalizeNoEP,

38 + MoEPrepareAndFinalizeNoDPEPModular,

39 )

41 from .utils import _get_lora_device, try_get_optimal_moe_lora_config

@@ -136,7 +136,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

136

137 if getattr(self.base_layer.quant_method, "supports_internal_mk", False):

138 # Use the existing modular kernel from the quant method

139 - m_fused_moe_fn = self.base_layer.quant_method.moe_mk

139 + m_fused_moe_fn = self.base_layer.quant_method.moe_kernel

140 # Don't let the kernel own shared experts so the runner can

141 # overlap them with routed experts via a separate CUDA stream.

142 m_fused_moe_fn.shared_experts = None

@@ -144,8 +144,8 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

144 # Create a new modular kernel via select_gemm_impl.

145 # Don't pass shared_experts to the kernel so the runner can

146 # overlap them with routed experts via a separate CUDA stream.

147 - prepare_finalize = MoEPrepareAndFinalizeNoEP()

148 - m_fused_moe_fn = FusedMoEModularKernel(

147 + prepare_finalize = MoEPrepareAndFinalizeNoDPEPModular()

148 + m_fused_moe_fn = FusedMoEKernel(

149 prepare_finalize,

150 self.base_layer.quant_method.select_gemm_impl(

151 prepare_finalize, self.base_layer

@@ -154,10 +154,11 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

154

155 if quant_config.use_mxfp4_w4a16:

156 assert isinstance(

157 - m_fused_moe_fn.fused_experts, (MarlinExperts, UnfusedOAITritonExperts)

157 + m_fused_moe_fn.impl.fused_experts,

158 + (MarlinExperts, UnfusedOAITritonExperts),

159 )

160 else:

160 - assert isinstance(m_fused_moe_fn.fused_experts, TritonExperts)

161 + assert isinstance(m_fused_moe_fn.impl.fused_experts, TritonExperts)

162

163 def fwd_decorator(layer, func):

164 def wrapper(*args, **kwargs):

@@ -337,9 +338,9 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):

338

339 return wrapper

340

340 - fused_experts = m_fused_moe_fn.fused_experts

341 + fused_experts = m_fused_moe_fn.impl.fused_experts

342

342 - m_fused_moe_fn.forward = fwd_decorator(self.base_layer, m_fused_moe_fn.forward)

343 + m_fused_moe_fn.apply = fwd_decorator(self.base_layer, m_fused_moe_fn.apply)

344 fused_experts.activation = act_decorator(

345 self.base_layer, fused_experts.activation

346 )

vllm/model_executor/layers/fused_moe/__init__.py

+4 -4

@@ -22,8 +22,8 @@ from vllm.model_executor.layers.fused_moe.layer import (

22 )

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

24 FusedMoEActivationFormat,

25 - FusedMoEPermuteExpertsUnpermute,

26 - FusedMoEPrepareAndFinalize,

25 + FusedMoEExpertsModular,

26 + FusedMoEPrepareAndFinalizeModular,

27 )

28 from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (

29 FusedMoERouter,

@@ -62,9 +62,9 @@ __all__ = [

62 "MoEActivation",

63 "UnquantizedFusedMoEMethod",

64 "FusedMoeWeightScaleSupported",

65 - "FusedMoEPermuteExpertsUnpermute",

65 + "FusedMoEExpertsModular",

66 "FusedMoEActivationFormat",

67 - "FusedMoEPrepareAndFinalize",

67 + "FusedMoEPrepareAndFinalizeModular",

68 "GateLinear",

69 "RoutingMethodType",

70 "SharedFusedMoE",

vllm/model_executor/layers/fused_moe/all2all_utils.py

+9 -6

@@ -21,8 +21,8 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (

21 FusedMoEPrepareAndFinalize,

22 )

23 from vllm.model_executor.layers.fused_moe.prepare_finalize import (

24 - MoEPrepareAndFinalizeNaiveEP,

25 - MoEPrepareAndFinalizeNoEP,

24 + make_moe_prepare_and_finalize_naive_dp_ep,

25 + make_moe_prepare_and_finalize_no_dp_ep,

26 )

27 from vllm.platforms import current_platform

28 from vllm.utils.import_utils import has_deep_ep, has_mori

@@ -77,6 +77,7 @@ def maybe_make_prepare_finalize(

77 quant_config: FusedMoEQuantConfig | None,

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

79 allow_new_interface: bool = False,

80 + use_monolithic: bool = False,

81 ) -> FusedMoEPrepareAndFinalize | None:

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

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

@@ -102,14 +103,15 @@ def maybe_make_prepare_finalize(

103 "Detected DP deployment with no --enable-expert-parallel. "

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

105 )

105 - return MoEPrepareAndFinalizeNaiveEP(

106 + return make_moe_prepare_and_finalize_naive_dp_ep(

107 is_sequence_parallel=moe.moe_parallel_config.is_sequence_parallel,

108 num_dispatchers=(

109 get_ep_group().device_communicator.all2all_manager.world_size

110 ),

111 + use_monolithic=use_monolithic,

112 )

113 else:

112 - return MoEPrepareAndFinalizeNoEP()

114 + return make_moe_prepare_and_finalize_no_dp_ep(use_monolithic)

115

116 all2all_manager = get_ep_group().device_communicator.all2all_manager

117 assert all2all_manager is not None

@@ -201,8 +203,9 @@ def maybe_make_prepare_finalize(

203 )

204

205 elif moe.use_naive_all2all_kernels and allow_new_interface:

204 - prepare_finalize = MoEPrepareAndFinalizeNaiveEP(

205 - is_sequence_parallel=(moe.moe_parallel_config.is_sequence_parallel),

206 + prepare_finalize = make_moe_prepare_and_finalize_naive_dp_ep(

207 + use_monolithic=use_monolithic,

208 + is_sequence_parallel=moe.moe_parallel_config.is_sequence_parallel,

209 num_dispatchers=all2all_manager.world_size,

210 )

211

vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py

+1 -1

@@ -261,7 +261,7 @@ def persistent_masked_m_silu_mul_quant(

261 return y_q, y_s

262

263

264 -class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

264 +class BatchedDeepGemmExperts(mk.FusedMoEExpertsModular):

265 def __init__(

266 self,

267 moe_config: FusedMoEConfig,

vllm/model_executor/layers/fused_moe/config.py

@@ -228,6 +228,7 @@ class FusedMoEQuantConfig:

228 _a2: FusedMoEQuantDesc

229 _w1: FusedMoEQuantDesc

230 _w2: FusedMoEQuantDesc

231 + is_nvfp4_scale_swizzled: bool = True

232

233 def __post_init__(self):

234 assert not self.per_act_token_quant or self.block_shape is None, (

@@ -475,6 +476,7 @@ class FusedMoEQuantConfig:

476 w1_zp: torch.Tensor | None = None,

477 w2_zp: torch.Tensor | None = None,

478 weight_dtype: torch.dtype | str | None = None,

479 + is_nvfp4_scale_swizzled: bool = True,

480 ) -> "FusedMoEQuantConfig":

481 """

482 General builder function for a FusedMoEQuantConfig.

@@ -504,6 +506,7 @@ class FusedMoEQuantConfig:

506 - w2_bias: Optional biases for w1 (GPT OSS Triton).

507 - w1_zp: Optional w1 zero points for int4/int8 quantization.

508 - w2_zp: Optional w2 zero points for int4/int8 quantization.

509 + - is_nvfp4_scale_swizzled: Whether to swizzle the nvfp4 scale swizzling.

510 """

511 assert not isinstance(quant_dtype, str) or quant_dtype in {

512 "nvfp4",

@@ -536,6 +539,7 @@ class FusedMoEQuantConfig:

539 _w2=FusedMoEQuantDesc(

540 weight_dtype, w_shape, w2_scale, g2_alphas, w2_zp, w2_bias

541 ),

542 + is_nvfp4_scale_swizzled=is_nvfp4_scale_swizzled,

543 )

544 assert quant_config.per_act_token_quant == per_act_token_quant

545 assert quant_config.per_out_ch_quant == per_out_ch_quant

@@ -737,6 +741,7 @@ def nvfp4_moe_quant_config(

741 w2_scale: torch.Tensor,

742 w1_bias: torch.Tensor | None = None,

743 w2_bias: torch.Tensor | None = None,

744 + is_nvfp4_scale_swizzled: bool = True,

745 ) -> FusedMoEQuantConfig:

746 """

747 Construct a quant config for mxfp4 activations and nvp4 weights.

@@ -754,6 +759,7 @@ def nvfp4_moe_quant_config(

759 per_act_token_quant=False,

760 per_out_ch_quant=False,

761 block_shape=None,

762 + is_nvfp4_scale_swizzled=is_nvfp4_scale_swizzled,

763 )

764

765

vllm/model_executor/layers/fused_moe/cutlass_moe.py

+7 -8

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

21 moe_unpermute,

22 )

23 from vllm.model_executor.layers.fused_moe.prepare_finalize import (

24 - MoEPrepareAndFinalizeNoEP,

24 + MoEPrepareAndFinalizeNoDPEPModular,

25 )

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

27 TopKWeightAndReduceDelegate,

@@ -262,7 +262,7 @@ def run_cutlass_moe_fp8(

262 )

263

264

265 -class CutlassExpertsFp8Base(mk.FusedMoEPermuteExpertsUnpermute):

265 +class CutlassExpertsFp8Base(mk.FusedMoEExpertsModular):

266 def __init__(

267 self,

268 moe_config: FusedMoEConfig,

@@ -661,7 +661,7 @@ def run_cutlass_moe_fp4(

661 return

662

663

664 -class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):

664 +class CutlassExpertsFp4(mk.FusedMoEExpertsModular):

665 """CUTLASS FP4 fused MoE expert implementation."""

666

667 @property

@@ -928,7 +928,7 @@ def run_cutlass_moe_w4a8_fp8(

928 )

929

930

931 -class CutlassExpertsW4A8Fp8(mk.FusedMoEPermuteExpertsUnpermute):

931 +class CutlassExpertsW4A8Fp8(mk.FusedMoEExpertsModular):

932 def __init__(

933 self,

934 out_dtype: torch.dtype | None,

@@ -1170,8 +1170,8 @@ def cutlass_moe_w4a8_fp8(

1170

1171 num_experts = global_num_experts if global_num_experts != -1 else w1_q.size(0)

1172

1173 - fn = mk.FusedMoEModularKernel(

1174 - MoEPrepareAndFinalizeNoEP(),

1173 + fn = mk.FusedMoEKernel(

1174 + MoEPrepareAndFinalizeNoDPEPModular(),

1175 CutlassExpertsW4A8Fp8(

1176 out_dtype=a.dtype,

1177 a_strides1=a_strides1,

@@ -1186,10 +1186,9 @@ def cutlass_moe_w4a8_fp8(

1186 quant_config=quant_config,

1187 group_size=group_size,

1188 ),

1189 - inplace=False,

1189 )

1190

1192 - return fn(

1191 + return fn.apply(

1192 a,

1193 w1_q,

1194 w2_q,

vllm/model_executor/layers/fused_moe/deep_gemm_moe.py

+1 -1

@@ -113,7 +113,7 @@ def _valid_deep_gemm(

113 return True

114

115

116 -class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

116 +class DeepGemmExperts(mk.FusedMoEExpertsModular):

117 """DeepGemm-based fused MoE expert implementation."""

118

119 def __init__(self, moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig):

vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py

+2 -1

@@ -25,7 +25,7 @@ from vllm.v1.worker.ubatching import (

25 )

28 -class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

28 +class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):

29 """

30 Prepare/Finalize using DeepEP High-Throughput kernels.

31 """

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

239 quant_dtype=quant_config.quant_dtype,

240 per_act_token_quant=False,

241 block_shape=quant_config.block_shape,

242 + is_fp4_scale_swizzled=quant_config.is_nvfp4_scale_swizzled,

243 )

244

245 return (

vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py

+2 -2

@@ -49,7 +49,7 @@ def dequant_fp8(

49 return (expert_x_fp32 * expert_x_scales).view(expert_x_fp8.size())

52 -class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

52 +class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):

53 """

54 Prepare/Finalize using DeepEP low-latency kernels.

55 """

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

119 # time. This setting is handled by post_init_setup.

120 self.use_ue8m0_dispatch = False

121

122 - def post_init_setup(self, fused_experts: mk.FusedMoEPermuteExpertsUnpermute):

122 + def post_init_setup(self, fused_experts: mk.FusedMoEExperts):

123 if not fused_experts.supports_packed_ue8m0_act_scales():

124 # Early exit.

125 return

vllm/model_executor/layers/fused_moe/experts/__init__.py

vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py new

+335

@@ -0,0 +1,335 @@

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.model_executor.layers.fused_moe.activation import MoEActivation

8 +from vllm.model_executor.layers.fused_moe.config import (

9 + FusedMoEConfig,

10 + FusedMoEParallelConfig,

11 + FusedMoEQuantConfig,

12 + RoutingMethodType,

13 +)

14 +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

15 + activation_to_flashinfer_int,

16 +)

17 +from vllm.model_executor.layers.quantization.utils.quant_utils import (

18 + QuantKey,

19 + kFp8Dynamic128Sym,

20 + kFp8Static128BlockSym,

21 + kFp8StaticTensorSym,

22 +)

23 +from vllm.platforms import current_platform

24 +

25 +

26 +class TrtLlmFp8Experts(mk.FusedMoEExpertsMonolithic):

27 + """

28 + Fp8 TRTLLM-Gen MoE kernels. Supports monolithic interface.

29 + """

30 +

31 + def __init__(

32 + self,

33 + moe_config: FusedMoEConfig,

34 + quant_config: FusedMoEQuantConfig,

35 + ):

36 + super().__init__(moe_config, quant_config)

37 +

38 + if moe_config.moe_parallel_config.use_ep and quant_config.is_per_tensor:

39 + raise NotImplementedError(

40 + "EP parallelism is not supported with TRTLLM"

41 + "per-tensor FP8 quantization."

42 + )

43 +

44 + self.routing_method_type = moe_config.routing_method

45 + self.topk = moe_config.experts_per_token

46 + self.intermediate_size_per_partition = (

47 + moe_config.intermediate_size_per_partition

48 + )

49 + self.hidden_dim = moe_config.hidden_dim

50 + self.local_num_experts = moe_config.num_local_experts

51 + self.ep_rank = moe_config.moe_parallel_config.ep_rank

52 +

53 + # Make additional scales for per-tensor interface.

54 + if self.quant_config.is_per_tensor:

55 + w1_scale = self.quant_config.w1_scale

56 + assert w1_scale is not None

57 + a1_scale = self.quant_config.a1_scale

58 + assert a1_scale is not None

59 + w2_scale = self.quant_config.w2_scale

60 + assert w2_scale is not None

61 + a2_scale = self.quant_config.a2_scale

62 + assert a2_scale is not None

63 +

64 + self._g1_alphas = (w1_scale * a1_scale).squeeze()

65 + self._g2_alphas = (w2_scale * a2_scale).squeeze()

66 + self._g1_scale_c = (

67 + self._g1_alphas / self.quant_config.a2_scale

68 + if moe_config.is_act_and_mul

69 + else torch.ones_like(self._g1_alphas) / self.quant_config.a2_scale

70 + )

71 +

72 + @staticmethod

73 + def activation_format() -> mk.FusedMoEActivationFormat:

74 + return mk.FusedMoEActivationFormat.Standard

75 +

76 + @staticmethod

77 + def _supports_current_device() -> bool:

78 + """Supports only Blackwell-family GPUs."""

79 + p = current_platform

80 + # Add check flashinfer trtllm is available

81 + return p.is_cuda() and p.is_device_capability_family(100)

82 +

83 + @staticmethod

84 + def _supports_no_act_and_mul() -> bool:

85 + """Does not support non-gated MoE (i.e. Nanotron-3-Nano)."""

86 + return True

87 +

88 + @staticmethod

89 + def _supports_quant_scheme(

90 + weight_key: QuantKey | None,

91 + activation_key: QuantKey | None,

92 + ) -> bool:

93 + """Supports Fp8 per-tensor and Fp8 block."""

94 + SUPPORTED_W_A = [

95 + (kFp8Static128BlockSym, kFp8Dynamic128Sym),

96 + (kFp8StaticTensorSym, kFp8StaticTensorSym),

97 + ]

98 + return (weight_key, activation_key) in SUPPORTED_W_A

99 +

100 + @staticmethod

101 + def _supports_activation(activation: MoEActivation) -> bool:

102 + """Supports only SiLU and RELU^2 non-gated activation."""

103 + return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]

104 +

105 + @staticmethod

106 + def _supports_routing_method(

107 + routing_method: RoutingMethodType,

108 + weight_key: QuantKey | None,

109 + activation_key: QuantKey | None,

110 + ) -> bool:

111 + """Monolithic kernels need to express router support."""

112 + # NOTE(dbari): TopK routing could also be enabled, but need to validate models

113 + # NOTE(dbari): Default is not implemented and should not be enabled until it is

114 + if (weight_key, activation_key) == (kFp8Static128BlockSym, kFp8Dynamic128Sym):

115 + # NOTE(rob): potentially allow others here. This is a conservative list.

116 + return routing_method in [

117 + RoutingMethodType.DeepSeekV3,

118 + RoutingMethodType.Renormalize,

119 + RoutingMethodType.RenormalizeNaive,

120 + ]

121 + elif (weight_key, activation_key) == (kFp8StaticTensorSym, kFp8StaticTensorSym):

122 + # NOTE(dbari): as above, potentially allow others here.

123 + return routing_method in [

124 + RoutingMethodType.DeepSeekV3,

125 + RoutingMethodType.Llama4,

126 + RoutingMethodType.Renormalize,

127 + RoutingMethodType.RenormalizeNaive,

128 + ]

129 + else:

130 + raise ValueError("Unsupported quantization scheme.")

131 +

132 + @staticmethod

133 + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:

134 + """Monolithic kernel so only use with naive DP/EP and TP."""

135 + return (

136 + not moe_parallel_config.use_all2all_kernels

137 + or moe_parallel_config.use_naive_all2all_kernels

138 + ) and not moe_parallel_config.enable_eplb

139 +

140 + @staticmethod

141 + def _supports_router_logits_dtype(

142 + router_logits_dtype: torch.dtype | None,

143 + routing_method: RoutingMethodType,

144 + ) -> bool:

145 + """

146 + The FlashInfer TRTLLM FP8 kernel expects bfloat16 router_logits by default.

147 + Only DeepSeekV3 routing supports float32 router_logits (which is converted

148 + internally in the kernel).

149 + """

150 + if router_logits_dtype == torch.float32:

151 + # Only DeepSeekV3 routing handles float32 logits

152 + # https://github.com/flashinfer-ai/flashinfer/issues/2469

153 + return routing_method == RoutingMethodType.DeepSeekV3

154 + return True

155 +

156 + def supports_chunking(self) -> bool:

157 + return False

158 +

159 + def supports_expert_map(self) -> bool:

160 + return False

161 +

162 + def _apply_per_block(

163 + self,

164 + hidden_states: torch.Tensor,

165 + w1: torch.Tensor,

166 + w2: torch.Tensor,

167 + router_logits: torch.Tensor,

168 + activation: MoEActivation,

169 + global_num_experts: int,

170 + expert_map: torch.Tensor | None,

171 + a1q_scale: torch.Tensor | None,

172 + apply_router_weight_on_input: bool,

173 + # grouped topk + fused topk bias parameters

174 + num_expert_group: int | None = None,

175 + e_score_correction_bias: torch.Tensor | None = None,

176 + routed_scaling_factor: float | None = None,

177 + topk_group: int | None = None,

178 + ) -> torch.Tensor:

179 + # Delay import for non-CUDA.

180 + import flashinfer

181 +

182 + assert not apply_router_weight_on_input

183 + assert activation == MoEActivation.SILU

184 +

185 + if e_score_correction_bias is not None:

186 + e_score_correction_bias = e_score_correction_bias.to(hidden_states.dtype)

187 +

188 + if self.routing_method_type == RoutingMethodType.DeepSeekV3:

189 + router_logits = router_logits.to(torch.float32)

190 +

191 + assert self.topk <= global_num_experts

192 + assert self.topk <= 10

193 + assert global_num_experts % 4 == 0

194 + assert self.quant_config.block_shape == [128, 128]

195 + # Routing kernel expects #experts <= #threads 512

196 + assert global_num_experts <= 512

197 +

198 + # Kernel requires transposed hidden state scales

199 + # TODO: fuse into the quant kernel.

200 + assert a1q_scale is not None

201 + a1q_scale_t = a1q_scale.t().contiguous()

202 +

203 + return flashinfer.fused_moe.trtllm_fp8_block_scale_moe(

204 + routing_logits=router_logits,

205 + routing_bias=e_score_correction_bias,

206 + hidden_states=hidden_states,

207 + hidden_states_scale=a1q_scale_t,

208 + gemm1_weights=w1,

209 + gemm1_weights_scale=self.quant_config.w1_scale,

210 + gemm2_weights=w2,

211 + gemm2_weights_scale=self.quant_config.w2_scale,

212 + num_experts=global_num_experts,

213 + top_k=self.topk,

214 + n_group=(num_expert_group or 0),

215 + topk_group=(topk_group or 0),

216 + intermediate_size=self.intermediate_size_per_partition,

217 + local_expert_offset=self.ep_rank * self.local_num_experts,

218 + local_num_experts=self.local_num_experts,

219 + routed_scaling_factor=routed_scaling_factor,

220 + routing_method_type=self.routing_method_type,

221 + use_shuffled_weight=False,

222 + )

223 +

224 + def _apply_per_tensor(

225 + self,

226 + hidden_states: torch.Tensor,

227 + w1: torch.Tensor,

228 + w2: torch.Tensor,

229 + router_logits: torch.Tensor,

230 + activation: MoEActivation,

231 + global_num_experts: int,

232 + expert_map: torch.Tensor | None,

233 + a1q_scale: torch.Tensor | None,

234 + apply_router_weight_on_input: bool,

235 + # grouped topk + fused topk bias parameters

236 + num_expert_group: int | None = None,

237 + e_score_correction_bias: torch.Tensor | None = None,

238 + routed_scaling_factor: float | None = None,

239 + topk_group: int | None = None,

240 + ) -> torch.Tensor:

241 + # Delay import for non-CUDA.

242 + import flashinfer

243 + from flashinfer.fused_moe.core import ActivationType

244 +

245 + # Confirm supported activation function.

246 + assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]

247 +

248 + activation_type = ActivationType(activation_to_flashinfer_int(activation))

249 +

250 + # Confirm Llama-4 routing is proper.

251 + if self.routing_method_type == RoutingMethodType.Llama4:

252 + assert apply_router_weight_on_input

253 + else:

254 + assert not apply_router_weight_on_input

255 +

256 + # The DeepSeekV3 routing method requires float32 router logits.

257 + if self.routing_method_type == RoutingMethodType.DeepSeekV3:

258 + router_logits = router_logits.to(torch.float32)

259 +

260 + out = flashinfer.fused_moe.trtllm_fp8_per_tensor_scale_moe(

261 + routing_logits=router_logits,

262 + routing_bias=e_score_correction_bias,

263 + hidden_states=hidden_states,

264 + gemm1_weights=w1,

265 + output1_scales_scalar=self._g1_scale_c,

266 + output1_scales_gate_scalar=self._g1_alphas,

267 + gemm2_weights=w2,

268 + output2_scales_scalar=self._g2_alphas,

269 + num_experts=global_num_experts,

270 + top_k=self.topk,

271 + n_group=num_expert_group or 0,

272 + topk_group=topk_group or 0,

273 + intermediate_size=self.intermediate_size_per_partition,

274 + local_expert_offset=self.ep_rank * self.local_num_experts,

275 + local_num_experts=self.local_num_experts,

276 + routed_scaling_factor=routed_scaling_factor,

277 + use_routing_scales_on_input=apply_router_weight_on_input,

278 + routing_method_type=self.routing_method_type,

279 + activation_type=activation_type,

280 + )

281 + return out

282 +

283 + def apply(

284 + self,

285 + hidden_states: torch.Tensor,

286 + w1: torch.Tensor,

287 + w2: torch.Tensor,

288 + router_logits: torch.Tensor,

289 + activation: MoEActivation,

290 + global_num_experts: int,

291 + expert_map: torch.Tensor | None,

292 + a1q_scale: torch.Tensor | None,

293 + apply_router_weight_on_input: bool,

294 + # grouped topk + fused topk bias parameters

295 + num_expert_group: int | None = None,

296 + e_score_correction_bias: torch.Tensor | None = None,

297 + routed_scaling_factor: float | None = None,

298 + topk_group: int | None = None,

299 + ) -> torch.Tensor:

300 + if self.quant_config.block_shape is not None:

301 + return self._apply_per_block(

302 + hidden_states,

303 + w1,

304 + w2,

305 + router_logits,

306 + activation,

307 + global_num_experts,

308 + expert_map,

309 + a1q_scale,

310 + apply_router_weight_on_input,

311 + num_expert_group=num_expert_group,

312 + e_score_correction_bias=e_score_correction_bias,

313 + routed_scaling_factor=routed_scaling_factor,

314 + topk_group=topk_group,

315 + )

316 + elif self.quant_config.is_per_tensor:

317 + return self._apply_per_tensor(

318 + hidden_states,

319 + w1,

320 + w2,

321 + router_logits,

322 + activation,

323 + global_num_experts,

324 + expert_map,

325 + a1q_scale,

326 + apply_router_weight_on_input,

327 + num_expert_group=num_expert_group,

328 + e_score_correction_bias=e_score_correction_bias,

329 + routed_scaling_factor=routed_scaling_factor,

330 + )

331 + else:

332 + raise NotImplementedError(

333 + "Only per-block and per-tensor quantization are supported in "

334 + f"{self.__class__.__name__}."

335 + )

vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py new

+326

@@ -0,0 +1,326 @@

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

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

3 +

4 +import flashinfer

5 +import torch

6 +

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

8 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation

9 +from vllm.model_executor.layers.fused_moe.config import (

10 + FusedMoEConfig,

11 + FusedMoEParallelConfig,

12 + FusedMoEQuantConfig,

13 + RoutingMethodType,

14 +)

15 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

16 + TopKWeightAndReduceNoOP,

17 +)

18 +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

19 + activation_to_flashinfer_int,

20 +)

21 +from vllm.model_executor.layers.quantization.utils.quant_utils import (

22 + QuantKey,

23 + kNvfp4Dynamic,

24 + kNvfp4Static,

25 +)

26 +from vllm.platforms import current_platform

27 +

28 +

29 +class TrtLlmNvFp4ExpertsBase:

30 + """

31 + NvFp4 TRTLLM-Gen MoE kernels. Supports modular and monolithic interface.

32 + """

33 +

34 + def __init__(

35 + self,

36 + moe_config: FusedMoEConfig,

37 + quant_config: FusedMoEQuantConfig,

38 + ):

39 + self.moe_config = moe_config

40 + self.quant_config = quant_config

41 +

42 + self.routing_method_type = self.moe_config.routing_method

43 + self.topk = moe_config.experts_per_token

44 + self.intermediate_size_per_partition = (

45 + moe_config.intermediate_size_per_partition

46 + )

47 + self.hidden_dim = moe_config.hidden_dim

48 + self.local_num_experts = moe_config.num_local_experts

49 + self.ep_rank = moe_config.moe_parallel_config.ep_rank

50 +

51 + assert self.quant_config.g1_alphas is not None

52 + assert self.quant_config.a2_gscale is not None

53 + if moe_config.is_act_and_mul:

54 + # g1_alpha_s = a13_scale * w13_scale_2

55 + # a2_gscale = (1 / a2_scale)

56 + # g1_scale_c = a13_scale * w13_scale_2 / a2_scale

57 + self.g1_scale_c = self.quant_config.g1_alphas * self.quant_config.a2_gscale

58 + else:

59 + self.g1_scale_c = (

60 + torch.ones_like(self.quant_config.a1_gscale)

61 + * self.quant_config.a2_gscale

62 + )

63 +

64 + @staticmethod

65 + def _supports_current_device() -> bool:

66 + """Supports only Blackwell-family GPUs."""

67 + p = current_platform

68 + return p.is_cuda() and p.is_device_capability_family(100)

69 +

70 + @staticmethod

71 + def _supports_no_act_and_mul() -> bool:

72 + """Supports non-gated MoE (i.e. Nemotron-Nano)."""

73 + return True

74 +

75 + @staticmethod

76 + def _supports_quant_scheme(

77 + weight_key: QuantKey | None,

78 + activation_key: QuantKey | None,

79 + ) -> bool:

80 + """Supports Nvfp4 quantization."""

81 + SUPPORTED_W_A = [

82 + (kNvfp4Static, kNvfp4Dynamic),

83 + ]

84 + return (weight_key, activation_key) in SUPPORTED_W_A

85 +

86 + @staticmethod

87 + def _supports_activation(activation: MoEActivation) -> bool:

88 + """Supports only SiLU and RELU^2 non-gated activation."""

89 + return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]

90 +

91 + @staticmethod

92 + def _supports_shape(hidden_dim: int) -> bool:

93 + """Requires hidden dim to be multiple of 512."""

94 + return hidden_dim % 512 == 0

95 +

96 + @staticmethod

97 + def activation_format() -> mk.FusedMoEActivationFormat:

98 + return mk.FusedMoEActivationFormat.Standard

99 +

100 + def supports_chunking(self) -> bool:

101 + return False

102 +

103 + def supports_expert_map(self) -> bool:

104 + return False

105 +

106 +

107 +class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModular):

108 + """

109 + Modular version of the implementation (just the experts).

110 + """

111 +

112 + @staticmethod

113 + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:

114 + """The modular implementation supports all parallel configs."""

115 + return True

116 +

117 + def workspace_shapes(

118 + self,

119 + M: int,

120 + N: int,

121 + K: int,

122 + topk: int,

123 + global_num_experts: int,

124 + local_num_experts: int,

125 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

126 + activation: MoEActivation,

127 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

128 + # The workspaces for this implementation are managed by flashinfer.

129 + workspace1 = (0,)

130 + workspace2 = (0,)

131 +

132 + # Hidden states are Nvfp4, packed into int8 dtype, so we

133 + # need to multiply K by 2 to get the output shape right.

134 + assert self.hidden_dim == K * 2

135 + output = (M, self.hidden_dim)

136 +

137 + return (workspace1, workspace2, output)

138 +

139 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:

140 + return TopKWeightAndReduceNoOP()

141 +

142 + def apply(

143 + self,

144 + output: torch.Tensor,

145 + hidden_states: torch.Tensor,

146 + w1: torch.Tensor,

147 + w2: torch.Tensor,

148 + topk_weights: torch.Tensor,

149 + topk_ids: torch.Tensor,

150 + activation: MoEActivation,

151 + global_num_experts: int,

152 + expert_map: torch.Tensor | None,

153 + a1q_scale: torch.Tensor | None,

154 + a2_scale: torch.Tensor | None,

155 + workspace13: torch.Tensor,

156 + workspace2: torch.Tensor,

157 + expert_tokens_meta: mk.ExpertTokensMetadata | None,

158 + apply_router_weight_on_input: bool,

159 + ):

160 + assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]

161 + assert a1q_scale is not None

162 + assert self.quant_config.w1_scale is not None

163 + assert self.quant_config.w2_scale is not None

164 +

165 + # Pack topk ids and weights into format expected by the kernel.

166 + packed_tensor = (topk_ids.to(torch.int32) << 16) | topk_weights.to(

167 + torch.bfloat16

168 + ).view(torch.int16)

169 +

170 + # trtllm_fp4_block_scale_routed_moe does not support autotuning

171 + # so skip this kernel during dummy run for autotuning.

172 + import vllm.utils.flashinfer as fi_utils

173 +

174 + if fi_utils._is_fi_autotuning:

175 + return hidden_states

176 +

177 + # Invoke kernel.

178 + flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe(

179 + topk_ids=packed_tensor,

180 + routing_bias=None,

181 + hidden_states=hidden_states,

182 + hidden_states_scale=a1q_scale.view(torch.float8_e4m3fn).reshape(

183 + *hidden_states.shape[:-1], -1

184 + ),

185 + gemm1_weights=w1,

186 + gemm1_weights_scale=self.quant_config.w1_scale.view(torch.float8_e4m3fn),

187 + gemm1_bias=None,

188 + gemm1_alpha=None,

189 + gemm1_beta=None,

190 + gemm1_clamp_limit=None,

191 + gemm2_weights=w2,

192 + gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn),

193 + gemm2_bias=None,

194 + output1_scale_scalar=self.g1_scale_c,

195 + output1_scale_gate_scalar=self.quant_config.g1_alphas,

196 + output2_scale_scalar=self.quant_config.g2_alphas,

197 + num_experts=global_num_experts,

198 + top_k=self.topk,

199 + n_group=0,

200 + topk_group=0,

201 + intermediate_size=self.intermediate_size_per_partition,

202 + local_expert_offset=self.ep_rank * self.local_num_experts,

203 + local_num_experts=self.local_num_experts,

204 + routed_scaling_factor=None,

205 + routing_method_type=1,

206 + do_finalize=True,

207 + activation_type=activation_to_flashinfer_int(activation),

208 + output=output,

209 + )

210 +

211 +

212 +class TrtLlmNvFp4ExpertsMonolithic(

213 + TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsMonolithic

214 +):

215 + """

216 + Monolithic version of the kernel (router + experts).

217 + """

218 +

219 + @staticmethod

220 + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:

221 + """The modular implementation should be used for the Dp/Ep or EPLB case."""

222 + return (

223 + not moe_parallel_config.use_all2all_kernels

224 + and not moe_parallel_config.enable_eplb

225 + )

226 +

227 + @staticmethod

228 + def _supports_routing_method(

229 + routing_method_type: RoutingMethodType,

230 + weight_key: QuantKey | None,

231 + activation_key: QuantKey | None,

232 + ) -> bool:

233 + # NOTE(rob): this is a conservative list.

234 + return routing_method_type in [

235 + RoutingMethodType.DeepSeekV3,

236 + RoutingMethodType.Renormalize,

237 + RoutingMethodType.RenormalizeNaive,

238 + RoutingMethodType.Llama4,

239 + ]

240 +

241 + @staticmethod

242 + def _supports_router_logits_dtype(

243 + router_logits_dtype: torch.dtype | None,

244 + routing_method: RoutingMethodType,

245 + ) -> bool:

246 + """

247 + The FlashInfer TRTLLM NvFp4 kernel expects bfloat16 router_logits by default.

248 + Only DeepSeekV3 routing supports float32 router_logits (which is converted

249 + internally in the kernel).

250 + """

251 + if router_logits_dtype == torch.float32:

252 + # Only DeepSeekV3 routing handles float32 logits

253 + # https://github.com/flashinfer-ai/flashinfer/issues/2469

254 + return routing_method == RoutingMethodType.DeepSeekV3

255 + return True

256 +

257 + def apply(

258 + self,

259 + hidden_states: torch.Tensor,

260 + w1: torch.Tensor,

261 + w2: torch.Tensor,

262 + router_logits: torch.Tensor,

263 + activation: MoEActivation,

264 + global_num_experts: int,

265 + expert_map: torch.Tensor | None,

266 + a1q_scale: torch.Tensor | None,

267 + apply_router_weight_on_input: bool,

268 + # grouped topk + fused topk bias parameters

269 + num_expert_group: int | None = None,

270 + e_score_correction_bias: torch.Tensor | None = None,

271 + routed_scaling_factor: float | None = None,

272 + topk_group: int | None = None,

273 + ) -> torch.Tensor:

274 + assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]

275 + assert a1q_scale is not None

276 + assert self.quant_config.w1_scale is not None

277 + assert self.quant_config.w2_scale is not None

278 + assert (

279 + apply_router_weight_on_input

280 + and self.routing_method_type == RoutingMethodType.Llama4

281 + ) or (

282 + not apply_router_weight_on_input

283 + and self.routing_method_type != RoutingMethodType.Llama4

284 + )

285 +

286 + # Prepare routing bias into kernel format.

287 + routing_bias = e_score_correction_bias

288 + if routing_bias is not None:

289 + routing_bias = routing_bias.to(torch.bfloat16)

290 + router_logits = (

291 + router_logits.to(torch.float32)

292 + if self.routing_method_type == RoutingMethodType.DeepSeekV3

293 + else router_logits

294 + )

295 +

296 + # Invoke kernel.

297 + return flashinfer.fused_moe.trtllm_fp4_block_scale_moe(

298 + routing_logits=router_logits,

299 + routing_bias=routing_bias,

300 + hidden_states=hidden_states,

301 + hidden_states_scale=a1q_scale.view(torch.float8_e4m3fn).reshape(

302 + *hidden_states.shape[:-1], -1

303 + ),

304 + gemm1_weights=w1,

305 + gemm1_weights_scale=self.quant_config.w1_scale.view(torch.float8_e4m3fn),

306 + gemm1_bias=None,

307 + gemm1_alpha=None,

308 + gemm1_beta=None,

309 + gemm1_clamp_limit=None,

310 + gemm2_weights=w2,

311 + gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn),

312 + gemm2_bias=None,

313 + output1_scale_scalar=self.g1_scale_c,

314 + output1_scale_gate_scalar=self.quant_config.g1_alphas,

315 + output2_scale_scalar=self.quant_config.g2_alphas,

316 + num_experts=global_num_experts,

317 + top_k=self.topk,

318 + n_group=(num_expert_group or 0),

319 + topk_group=(topk_group or 0),

320 + intermediate_size=self.intermediate_size_per_partition,

321 + local_expert_offset=self.ep_rank * self.local_num_experts,

322 + local_num_experts=self.local_num_experts,

323 + routed_scaling_factor=routed_scaling_factor,

324 + routing_method_type=self.routing_method_type,

325 + do_finalize=True,

326 + )[0]

vllm/model_executor/layers/fused_moe/fallback.py

+6 -6

@@ -11,13 +11,13 @@ from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig

11 from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey

14 -class FallbackExperts(mk.FusedMoEPermuteExpertsUnpermute, ABC):

14 +class FallbackExperts(mk.FusedMoEExpertsModular, ABC):

15 """Base class for runtime dispatching of expert implementations."""

17 def __init__(

18 self,

19 - experts: mk.FusedMoEPermuteExpertsUnpermute,

20 - fallback_experts: mk.FusedMoEPermuteExpertsUnpermute,

19 + experts: mk.FusedMoEExpertsModular,

20 + fallback_experts: mk.FusedMoEExpertsModular,

21 ):

22 super().__init__(

23 moe_config=experts.moe_config, quant_config=experts.quant_config

@@ -27,8 +27,8 @@ class FallbackExperts(mk.FusedMoEPermuteExpertsUnpermute, ABC):

28 @staticmethod

29 def get_clses() -> tuple[

30 - type[mk.FusedMoEPermuteExpertsUnpermute],

31 - type[mk.FusedMoEPermuteExpertsUnpermute],

30 + type[mk.FusedMoEExpertsModular],

31 + type[mk.FusedMoEExpertsModular],

32 ]:

33 """

34 Get the cls for the experts and fallback experts.

@@ -149,7 +149,7 @@ class FallbackExperts(mk.FusedMoEPermuteExpertsUnpermute, ABC):

149 hidden_states: torch.Tensor,

150 w1: torch.Tensor,

151 w2: torch.Tensor,

152 - ) -> mk.FusedMoEPermuteExpertsUnpermute:

152 + ) -> mk.FusedMoEExpertsModular:

153 raise NotImplementedError

154

155 def apply(

vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py

+3 -3

@@ -18,7 +18,7 @@ def get_local_sizes():

18 return get_forward_context().dp_metadata.get_chunk_sizes_across_dp_rank()

21 -class FlashInferA2APrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

21 +class FlashInferA2APrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):

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

24 def __init__(

@@ -185,8 +185,8 @@ def flashinfer_alltoall_dispatch(

185 ep_size,

186 )

187

188 - # Swizzle after the A2A if nvfp4.

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

188 + # Swizzle after the A2A if MoE kernel expects swizzled scales.

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

190 if x_sf.element_size() == 1:

191 x_sf = x_sf.view(torch.uint8)

192 x_sf = nvfp4_block_scale_interleave(x_sf)

vllm/model_executor/layers/fused_moe/flashinfer_cutedsl_moe.py

+1 -1

@@ -30,7 +30,7 @@ from vllm.utils.flashinfer import (

30 logger = init_logger(__name__)

33 -class FlashInferCuteDSLExperts(mk.FusedMoEPermuteExpertsUnpermute):

33 +class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):

34 def __init__(

35 self,

36 moe_config: FusedMoEConfig,

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py

+1 -1

@@ -60,7 +60,7 @@ def is_valid_flashinfer_cutlass_fused_moe(

60 return True

63 -class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):

63 +class FlashInferExperts(mk.FusedMoEExpertsModular):

64 def __init__(

65 self,

66 moe_config: mk.FusedMoEConfig,

vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py

-298

@@ -10,16 +10,6 @@ from vllm.model_executor.layers.fused_moe.config import (

10 FusedMoEParallelConfig,

11 RoutingMethodType,

12 )

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

14 -from vllm.model_executor.layers.quantization.utils.fp8_utils import (

15 - per_token_group_quant_fp8,

16 -)

17 -from vllm.model_executor.layers.quantization.utils.quant_utils import (

18 - QuantKey,

19 - kFp8Dynamic128Sym,

20 - kFp8Static128BlockSym,

21 - kFp8StaticTensorSym,

22 -)

13 from vllm.platforms import current_platform

14 from vllm.utils.torch_utils import direct_register_custom_op

@@ -39,49 +29,10 @@ def _supports_no_act_and_mul() -> bool:

29 return True

42 -def _supports_quant_scheme(

43 - weight_key: QuantKey | None,

44 - activation_key: QuantKey | None,

45 -) -> bool:

46 - """Supports Fp8 per-tensor and Fp8 block."""

47 - SUPPORTED_W_A = [

48 - (kFp8Static128BlockSym, kFp8Dynamic128Sym),

49 - (kFp8StaticTensorSym, kFp8StaticTensorSym),

50 - ]

51 - return (weight_key, activation_key) in SUPPORTED_W_A

52 -

53 -

32 def _supports_activation(activation: MoEActivation) -> bool:

33 return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]

58 -def _supports_routing_method(

59 - weight_key: QuantKey | None,

60 - activation_key: QuantKey | None,

61 - routing_method: RoutingMethodType,

62 -) -> bool:

63 - """Monolithic kernels need to express router support."""

64 - # NOTE(dbari): TopK routing could also be enabled, but need to validate models

65 - # NOTE(dbari): Default is not implemented and should not be enabled until it is

66 - if (weight_key, activation_key) == (kFp8Static128BlockSym, kFp8Dynamic128Sym):

67 - # NOTE(rob): potentially allow others here. This is a conservative list.

68 - return routing_method in [

69 - RoutingMethodType.DeepSeekV3,

70 - RoutingMethodType.Renormalize,

71 - RoutingMethodType.RenormalizeNaive,

72 - ]

73 - elif (weight_key, activation_key) == (kFp8StaticTensorSym, kFp8StaticTensorSym):

74 - # NOTE(dbari): as above, potentially allow others here.

75 - return routing_method in [

76 - RoutingMethodType.DeepSeekV3,

77 - RoutingMethodType.Llama4,

78 - RoutingMethodType.Renormalize,

79 - RoutingMethodType.RenormalizeNaive,

80 - ]

81 - else:

82 - raise ValueError("Unsupported quantization scheme.")

83 -

84 -

36 def _supports_routing_method_bf16(

37 routing_method: RoutingMethodType,

38 ) -> bool:

@@ -99,62 +50,6 @@ def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bo

50 return not moe_parallel_config.enable_eplb

102 -def _supports_router_logits_dtype(

103 - router_logits_dtype: torch.dtype | None,

104 - routing_method: RoutingMethodType,

105 -) -> bool:

106 - """

107 - The FlashInfer TRTLLM FP8 kernel expects bfloat16 router_logits by default.

108 - Only DeepSeekV3 routing supports float32 router_logits (which is converted

109 - internally in the kernel).

110 - """

111 - if router_logits_dtype == torch.float32:

112 - # Only DeepSeekV3 routing handles float32 logits

113 - # https://github.com/flashinfer-ai/flashinfer/issues/2469

114 - return routing_method == RoutingMethodType.DeepSeekV3

115 - return True

116 -

117 -

118 -def is_supported_config_trtllm_fp8(

119 - moe_config: FusedMoEConfig,

120 - weight_key: QuantKey | None,

121 - activation_key: QuantKey | None,

122 - activation_format: mk.FusedMoEActivationFormat,

123 -) -> tuple[bool, str | None]:

124 - """

125 - This method mirrors mk.FusedMoEPermuteExpertsUnpermute.is_supported_config

126 - """

127 -

128 - def _make_reason(reason: str) -> str:

129 - return f"kernel does not support {reason}"

130 -

131 - if not _supports_current_device():

132 - return False, _make_reason(f"current device {current_platform.device_name}")

133 - elif not (moe_config.is_act_and_mul or _supports_no_act_and_mul()):

134 - return False, _make_reason("no act_and_mul MLP layer")

135 - elif not _supports_activation(moe_config.activation):

136 - return False, _make_reason(f"{moe_config.activation} activation")

137 - elif not _supports_quant_scheme(weight_key, activation_key):

138 - return False, _make_reason(f"quantization scheme {weight_key}x{activation_key}")

139 - elif not _supports_parallel_config(moe_config.moe_parallel_config):

140 - return False, _make_reason(f"parallel config {moe_config.moe_parallel_config}")

141 - elif not _supports_routing_method(

142 - weight_key, activation_key, moe_config.routing_method

143 - ):

144 - return False, _make_reason(f"routing method {moe_config.routing_method}")

145 - elif activation_format != mk.FusedMoEActivationFormat.Standard:

146 - return False, _make_reason(f"activation format {activation_format}")

147 - elif not _supports_router_logits_dtype(

148 - moe_config.router_logits_dtype, moe_config.routing_method

149 - ):

150 - return False, _make_reason(

151 - "float32 router_logits with non-DeepSeekV3 routing "

152 - f"{moe_config.router_logits_dtype}x{moe_config.routing_method}"

153 - )

154 -

155 - return True, None

156 -

157 -

53 def is_supported_config_trtllm_bf16(

54 moe_config: FusedMoEConfig,

55 activation_format: mk.FusedMoEActivationFormat,

@@ -183,199 +78,6 @@ def is_supported_config_trtllm_bf16(

78 return True, None

186 -def flashinfer_fused_moe_blockscale_fp8(

187 - routing_logits: torch.Tensor,

188 - routing_bias: torch.Tensor | None,

189 - x: torch.Tensor,

190 - w13_weight: torch.Tensor,

191 - w13_weight_scale_inv: torch.Tensor,

192 - w2_weight: torch.Tensor,

193 - w2_weight_scale_inv: torch.Tensor,

194 - global_num_experts: int,

195 - top_k: int,

196 - num_expert_group: int | None,

197 - topk_group: int | None,

198 - intermediate_size: int,

199 - expert_offset: int,

200 - local_num_experts: int,

201 - block_shape: list[int],

202 - routing_method_type: int,

203 - routed_scaling: float | None = 1.0,

204 -) -> torch.Tensor:

205 - from vllm.utils.flashinfer import flashinfer_trtllm_fp8_block_scale_moe

206 -

207 - num_expert_group = num_expert_group if num_expert_group is not None else 0

208 - topk_group = topk_group if topk_group is not None else 0

209 - assert top_k <= global_num_experts

210 - assert top_k <= 10

211 - assert global_num_experts % 4 == 0

212 - assert block_shape == [128, 128]

213 - # Routing kernel expects #experts <= #threads 512

214 - assert global_num_experts <= 512

215 -

216 - # The DeepSeekV3 routing method requires float32 router logits.

217 - if routing_method_type == RoutingMethodType.DeepSeekV3:

218 - routing_logits = routing_logits.to(torch.float32)

219 -

220 - if routing_bias is not None:

221 - routing_bias = routing_bias.to(x.dtype)

222 -

223 - a_q, a_sf = per_token_group_quant_fp8(x, block_shape[1])

224 - # NOTE: scales of hidden states have to be transposed!

225 - a_sf_t = a_sf.t().contiguous()

226 - return flashinfer_trtllm_fp8_block_scale_moe(

227 - routing_logits=routing_logits,

228 - routing_bias=routing_bias,

229 - hidden_states=a_q,

230 - hidden_states_scale=a_sf_t,

231 - gemm1_weights=w13_weight,

232 - gemm1_weights_scale=w13_weight_scale_inv,

233 - gemm2_weights=w2_weight,

234 - gemm2_weights_scale=w2_weight_scale_inv,

235 - num_experts=global_num_experts,

236 - top_k=top_k,

237 - n_group=num_expert_group,

238 - topk_group=topk_group,

239 - intermediate_size=intermediate_size,

240 - local_expert_offset=expert_offset,

241 - local_num_experts=local_num_experts,

242 - routed_scaling_factor=routed_scaling,

243 - routing_method_type=routing_method_type,

244 - use_shuffled_weight=False,

245 - )

246 -

247 -

248 -def flashinfer_fused_moe_blockscale_fp8_fake(

249 - routing_logits: torch.Tensor,

250 - routing_bias: torch.Tensor | None,

251 - x: torch.Tensor,

252 - w13_weight: torch.Tensor,

253 - w13_weight_scale_inv: torch.Tensor,

254 - w2_weight: torch.Tensor,

255 - w2_weight_scale_inv: torch.Tensor,

256 - global_num_experts: int,

257 - top_k: int,

258 - num_expert_group: int,

259 - topk_group: int,

260 - intermediate_size: int,

261 - expert_offset: int,

262 - local_num_experts: int,

263 - block_shape: list[int],

264 - routing_method_type: int,

265 - routed_scaling: float = 1.0,

266 -) -> torch.Tensor:

267 - return torch.empty_like(x)

268 -

269 -

270 -# TODO(bnell): Does this really need to be a torch.op?

271 -direct_register_custom_op(

272 - op_name="flashinfer_fused_moe_blockscale_fp8",

273 - op_func=flashinfer_fused_moe_blockscale_fp8,

274 - fake_impl=flashinfer_fused_moe_blockscale_fp8_fake,

275 - tags=(torch.Tag.needs_fixed_stride_order,),

276 -)

277 -

278 -

279 -def fi_trtllm_fp8_per_tensor_moe(

280 - routing_logits: torch.Tensor,

281 - routing_bias: torch.Tensor | None,

282 - hidden_states: torch.Tensor,

283 - input_scale: torch.Tensor,

284 - gemm1_weights: torch.Tensor,

285 - gemm2_weights: torch.Tensor,

286 - output1_scales_scalar: torch.Tensor,

287 - output1_scales_gate_scalar: torch.Tensor,

288 - output2_scales_scalar: torch.Tensor,

289 - num_experts: int,

290 - top_k: int,

291 - num_expert_group: int | None,

292 - topk_group: int | None,

293 - intermediate_size: int,

294 - local_expert_offset: int,

295 - local_num_experts: int,

296 - use_routing_scales_on_input: bool,

297 - routing_method_type: int,

298 - activation_type: int,

299 - routed_scaling_factor: float = 1.0,

300 -) -> torch.Tensor:

301 - num_expert_group = num_expert_group if num_expert_group is not None else 0

302 - topk_group = topk_group if topk_group is not None else 0

303 -

304 - quant_hidden_states, _ = moe_kernel_quantize_input(

305 - hidden_states,

306 - input_scale,

307 - quant_dtype=torch.float8_e4m3fn,

308 - per_act_token_quant=False,

309 - )

310 -

311 - from flashinfer.fused_moe.core import ActivationType

312 -

313 - from vllm.utils.flashinfer import flashinfer_trtllm_fp8_per_tensor_scale_moe

314 -

315 - # The DeepSeekV3 routing method requires float32 router logits.

316 - if routing_method_type == RoutingMethodType.DeepSeekV3:

317 - routing_logits = routing_logits.to(torch.float32)

318 -

319 - return flashinfer_trtllm_fp8_per_tensor_scale_moe(

320 - routing_logits=routing_logits,

321 - routing_bias=routing_bias,

322 - hidden_states=quant_hidden_states,

323 - gemm1_weights=gemm1_weights,

324 - output1_scales_scalar=output1_scales_scalar,

325 - output1_scales_gate_scalar=output1_scales_gate_scalar,

326 - gemm2_weights=gemm2_weights,

327 - output2_scales_scalar=output2_scales_scalar,

328 - num_experts=num_experts,

329 - top_k=top_k,

330 - n_group=num_expert_group,

331 - topk_group=topk_group,

332 - intermediate_size=intermediate_size,

333 - local_expert_offset=local_expert_offset,

334 - local_num_experts=local_num_experts,

335 - routed_scaling_factor=routed_scaling_factor,

336 - use_routing_scales_on_input=use_routing_scales_on_input,

337 - routing_method_type=routing_method_type,

338 - # TODO: enum type Required for flashinfer==0.6.3, remove with update

339 - # https://github.com/flashinfer-ai/flashinfer/pull/2508

340 - activation_type=ActivationType(activation_type),

341 - )

342 -

343 -

344 -def fi_trtllm_fp8_per_tensor_moe_fake(

345 - routing_logits: torch.Tensor,

346 - routing_bias: torch.Tensor | None,

347 - hidden_states: torch.Tensor,

348 - input_scale: torch.Tensor,

349 - gemm1_weights: torch.Tensor,

350 - gemm2_weights: torch.Tensor,

351 - output1_scales_scalar: torch.Tensor,

352 - output1_scales_gate_scalar: torch.Tensor,

353 - output2_scales_scalar: torch.Tensor,

354 - num_experts: int,

355 - top_k: int,

356 - num_expert_group: int | None,

357 - topk_group: int | None,

358 - intermediate_size: int,

359 - local_expert_offset: int,

360 - local_num_experts: int,

361 - use_routing_scales_on_input: bool,

362 - routing_method_type: int,

363 - activation_type: int,

364 - routed_scaling_factor: float = 1.0,

365 -) -> torch.Tensor:

366 - return torch.empty_like(hidden_states)

367 -

368 -

369 -# TODO(bnell): Does this really need to be a torch.op?

370 -direct_register_custom_op(

371 - op_name="fi_trtllm_fp8_per_tensor_moe",

372 - op_func=fi_trtllm_fp8_per_tensor_moe,

373 - mutates_args=["hidden_states"],

374 - fake_impl=fi_trtllm_fp8_per_tensor_moe_fake,

375 - tags=(torch.Tag.needs_fixed_stride_order,),

376 -)

377 -

378 -

81 def flashinfer_fused_moe_bf16(

82 routing_logits: torch.Tensor,

83 routing_bias: torch.Tensor | None,

vllm/model_executor/layers/fused_moe/fused_batched_moe.py

+3 -3

@@ -489,7 +489,7 @@ def invoke_moe_batched_triton_kernel(

489 )

490

491

492 -class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

492 +class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):

493 """

494 A reference prepare/finalize class that reorganizes the tokens into

495 expert batched format, i.e. E x max_num_tokens x K. This is the format

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

645 )

646

647

648 -class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):

648 +class NaiveBatchedExperts(mk.FusedMoEExpertsModular):

649 """

650 A reference MoE expert class that operates on expert batched format,

651 i.e. E x max_num_tokens x K. This is the format that the batched

@@ -877,7 +877,7 @@ def batched_moe_kernel_quantize_input(

877 return A_q, A_q_scale

878

879

880 -class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

880 +class BatchedTritonExperts(mk.FusedMoEExpertsModular):

881 """

882 A Triton based MoE expert class that operates on expert batched format,

883 i.e. E x max_num_tokens x K. This is the format that the batched

vllm/model_executor/layers/fused_moe/fused_marlin_moe.py

+1 -1

@@ -526,7 +526,7 @@ def batched_fused_marlin_moe(

526 return output

527

528

529 -class MarlinExpertsBase(mk.FusedMoEPermuteExpertsUnpermute):

529 +class MarlinExpertsBase(mk.FusedMoEExpertsModular):

530 def __init__(

531 self,

532 moe_config: FusedMoEConfig,

vllm/model_executor/layers/fused_moe/fused_moe.py

+2 -2

@@ -1736,7 +1736,7 @@ def fused_experts_impl(

1736 intermediate_cache3 = cache13[: M * top_k_num * K].view(M, top_k_num, K)

1737

1738 # This needs separate memory since it's used concurrently with cache1

1739 - activation_out_dim = mk.FusedMoEPermuteExpertsUnpermute.adjust_N_for_activation(

1739 + activation_out_dim = mk.FusedMoEExpertsModular.adjust_N_for_activation(

1740 N, activation_enum

1741 )

1742 intermediate_cache2 = torch.empty(

@@ -1924,7 +1924,7 @@ def fused_experts_impl(

1924 return out_hidden_states

1925

1926

1927 -class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

1927 +class TritonExperts(mk.FusedMoEExpertsModular):

1928 """Triton-based fused MoE expert implementation."""

1929

1930 def __init__(

vllm/model_executor/layers/fused_moe/fused_moe_method_base.py

+24 -27

@@ -12,8 +12,8 @@ from vllm.model_executor.layers.fused_moe.config import (

12 FusedMoEQuantConfig,

13 )

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

15 - FusedMoEPermuteExpertsUnpermute,

16 - FusedMoEPrepareAndFinalize,

15 + FusedMoEExpertsModular,

16 + FusedMoEPrepareAndFinalizeModular,

17 )

18 from vllm.model_executor.layers.quantization.base_config import (

19 QuantizeMethodBase,

@@ -27,19 +27,21 @@ 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

30 + self.moe_kernel: mk.FusedMoEKernel | None = None

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

36 + return self.moe_kernel is not None

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

42 + return (

43 + self.moe_kernel is not None and self.moe_kernel.shared_experts is not None

44 + )

46 @abstractmethod

47 def create_weights(

@@ -66,35 +68,25 @@ class FusedMoEMethodBase(QuantizeMethodBase):

68 def maybe_make_prepare_finalize(

69 self,

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

69 - ) -> FusedMoEPrepareAndFinalize | None:

71 + ) -> FusedMoEPrepareAndFinalizeModular | None:

72 from .all2all_utils import maybe_make_prepare_finalize

72 - return maybe_make_prepare_finalize(

74 + pf = maybe_make_prepare_finalize(

75 self.moe, self.moe_quant_config, routing_tables

76 )

77 + assert pf is None or isinstance(pf, FusedMoEPrepareAndFinalizeModular)

78 + return pf

80 def select_gemm_impl(

81 self,

78 - prepare_finalize: FusedMoEPrepareAndFinalize,

82 + prepare_finalize: FusedMoEPrepareAndFinalizeModular,

83 layer: torch.nn.Module,

80 - ) -> FusedMoEPermuteExpertsUnpermute:

84 + ) -> FusedMoEExpertsModular:

85 # based on the all2all implementation, select the appropriate

86 # gemm implementation

83 - raise NotImplementedError(

84 - f"{self.__class__.__name__} must select appropriate gemm "

85 - "implementation based on the prepare_finalize"

86 - )

87 -

88 - def prepare_dp_allgather_tensor(

89 - self,

90 - layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821

91 - hidden_states: torch.Tensor,

92 - router_logits: torch.Tensor,

93 - ) -> tuple[torch.Tensor, list[torch.Tensor]]:

94 - """Hook to prepare tensors and extra tensors for DP allgather + EP dispatch."""

95 - raise NotImplementedError(

96 - "Method 'prepare_dp_allgather_tensor' is not implemented in "

97 - f"{self.__class__.__name__}."

87 + raise ValueError(

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

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

90 )

92 @abstractmethod

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

98 @property

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

100 + if self.moe_kernel is not None:

101 + return self.moe_kernel.prepare_finalize.topk_indices_dtype()

102 return None

103

104 @property

@@ -119,7 +111,12 @@ class FusedMoEMethodBase(QuantizeMethodBase):

111

112 @property

113 def is_monolithic(self) -> bool:

122 - return False

114 + if self.moe_kernel is None:

115 + if hasattr(self, "experts_cls"):

116 + return self.experts_cls.is_monolithic()

117 + else:

118 + return False

119 + return self.moe_kernel.is_monolithic

120

121 def apply(

122 self,

vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py

+9 -9

@@ -13,8 +13,8 @@ from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (

13 FusedMoEMethodBase,

14 )

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

16 - FusedMoEModularKernel,

17 - FusedMoEPrepareAndFinalize,

16 + FusedMoEKernel,

17 + FusedMoEPrepareAndFinalizeModular,

18 )

20 logger = init_logger(__name__)

@@ -26,15 +26,15 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):

26 # --8<-- [end:modular_fused_moe]

28 def __init__(

29 - self, old_quant_method: FusedMoEMethodBase, experts: FusedMoEModularKernel

29 + self, old_quant_method: FusedMoEMethodBase, moe_kernel: FusedMoEKernel

30 ):

31 super().__init__(old_quant_method.moe)

32 self.moe_quant_config = old_quant_method.moe_quant_config

33 - self.moe_mk = experts

33 + self.moe_kernel = moe_kernel

34 self.disable_expert_map = getattr(

35 old_quant_method,

36 "disable_expert_map",

37 - not self.moe_mk.supports_expert_map(),

37 + not self.moe_kernel.supports_expert_map(),

38 )

39 self.old_quant_method = old_quant_method

40 logger.debug("Swapping out %s", self.old_quant_method.__class__.__name__)

@@ -43,13 +43,13 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):

43 def make(

44 moe_layer: torch.nn.Module,

45 old_quant_method: FusedMoEMethodBase,

46 - prepare_finalize: FusedMoEPrepareAndFinalize,

46 + prepare_finalize: FusedMoEPrepareAndFinalizeModular,

47 shared_experts: torch.nn.Module | None,

48 inplace: bool = False,

49 ) -> "FusedMoEModularMethod":

50 return FusedMoEModularMethod(

51 old_quant_method,

52 - FusedMoEModularKernel(

52 + FusedMoEKernel(

53 prepare_finalize,

54 old_quant_method.select_gemm_impl(prepare_finalize, moe_layer),

55 shared_experts,

@@ -90,8 +90,8 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):

90 topk_ids: torch.Tensor,

91 shared_experts_input: torch.Tensor | None,

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

93 - assert self.moe_mk is not None

94 - return self.moe_mk(

93 + assert self.moe_kernel is not None

94 + return self.moe_kernel.apply(

95 hidden_states=x,

96 w1=layer.w13_weight,

97 w2=layer.w2_weight,

vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py

+1 -1

@@ -511,7 +511,7 @@ def make_routing_data(

511 return routing_data, gather_indx, scatter_indx

512

513

514 -class BaseOAITritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

514 +class BaseOAITritonExperts(mk.FusedMoEExpertsModular):

515 @staticmethod

516 def _supports_current_device() -> bool:

517 raise NotImplementedError(

vllm/model_executor/layers/fused_moe/modular_kernel.py

+518 -156

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

20 FusedMoEConfig,

21 FusedMoEParallelConfig,

22 FusedMoEQuantConfig,

23 + RoutingMethodType,

24 )

25 from vllm.model_executor.layers.fused_moe.utils import (

26 _resize_cache,

@@ -56,25 +57,25 @@ logger = init_logger(__name__)

57 # MoE kernel implementations.

58 #

59 # The following main classes are defined:

59 -# * FusedMoEPrepareAndFinalize - an abstract base class for preparation of MoE

60 +# * FusedMoEPrepareAndFinalizeModular - an abstract base class for preparation of MoE

61 # inputs (e.g. quantization, distribution) and finalization of Moe outputs.

62 # The prepare method must take care of any needed quantization and the

62 -# finalize method, informed by the FusedMoEPermuteExpertsUnpermute method,

63 +# finalize method, informed by the FusedMoEExpertsModular method,

64 # may apply weights and/or do the final reduction of the output.

64 -# * FusedMoEPermuteExpertsUnpermute - an abstract base class for the main fused

65 +# * FusedMoEExpertsModular - an abstract base class for the main fused

66 # MoE operation, i.e matmul + act_mul + optionally quant + matmul.

66 -# Some FusedMoEPermuteExpertsUnpermute implementations may choose to do

67 +# Some FusedMoEExpertsModular implementations may choose to do

68 # the weight application and/or reduction. The class communicates this

69 # to [Finalize] via a TopKWeightAndReduce object.

70 # * FusedMoEModularKernel - an interface class that combines a

70 -# FusedMoEPrepareAndFinalize and a FusedMoEPermuteExpertsUnpermute to

71 +# FusedMoEPrepareAndFinalizeModular and a FusedMoEExpertsModular to

72 # provide the standard fused MoE kernel interface.

73 # * TopKWeightAndReduce - A TopKWeightAndReduce implementation chosen

73 -# by the FusedMoEPermuteExpertsUnpermute implementation that is passed

74 +# by the FusedMoEExpertsModular implementation that is passed

75 # on to [Finalize].

76 #

77 # [Quantize-Prepare] and [Finalize] functionality are bundled into a single

77 -# class `FusedMoEPrepareAndFinalize` since they could use collective

78 +# class `FusedMoEPrepareAndFinalizeModular` since they could use collective

79 # communication mechanisms that need to be consistent.

80 #

@@ -155,25 +156,96 @@ PrepareResultType = tuple[

156 torch.Tensor | None,

157 ]

158

159 +#

160 +# PrepareResultType is a tuple of:

161 +# - quantized + dispatched a.

162 +# - quantized + dispatched a1_scales.

163 +# - dispatched router logits.

164 +#

165 +# See `prepare_monolithic` method below.

166 +#

167 +PrepareMonolithicResultType = tuple[

168 + torch.Tensor,

169 + torch.Tensor | None,

170 + torch.Tensor,

171 +]

172 +

173 ReceiverType = Callable[[], PrepareResultType]

174

175 +################################################################################

176 +# Prepare/Finalize

177 +################################################################################

178 +

179

161 -# TODO: pass FusedMoEParallelConfig in as ctor parameter?

180 class FusedMoEPrepareAndFinalize(ABC):

181 """

182 An abstract base class for the [Quantize-Prepare] and [Finalize] steps

183 described above.

184 +

185 + There are two variants of this class:

186 + * FusedMoEPrepareAndFinalizeModular - this operates on topk ids and weights

187 + * FusedMoEPrepareAndFinalizeMonolithic - the operates on router_logits

188 """

189

168 - def post_init_setup(self, fused_experts: "FusedMoEPermuteExpertsUnpermute"):

190 + def post_init_setup(self, fused_experts: "FusedMoEExperts"):

191 """

170 - Initialize FusedMoEPrepareAndFinalize settings that depend on

171 - FusedMoEPermuteExpertsUnpermute experts object.

172 - The FusedMoEPrepareAndFinalize implementations that have such

192 + Initialize FusedMoEPrepareAndFinalizeModular settings that depend on

193 + FusedMoEExpertsModular experts object.

194 + The FusedMoEPrepareAndFinalizeModular implementations that have such

195 dependencies may choose to override this function.

196 """

197 return

198

199 + @property

200 + @abstractmethod

201 + def activation_format(self) -> FusedMoEActivationFormat:

202 + """

203 + A property indicating the output format of the activations for the

204 + 'prepare' method.

205 + """

206 + raise NotImplementedError

207 +

208 + @abstractmethod

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

210 + """

211 + The PrepareFinalize All2All implementations generally constrain the

212 + dtype of the topk_ids they support. This function returns the

213 + required topk indices dtype so it can be respected.

214 + Return None if there are no such restrictions.

215 + """

216 + raise NotImplementedError

217 +

218 + @abstractmethod

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

220 + """

221 + Some PrepareFinalize All2All implementations are batched. Meaning,

222 + they can process only as set of tokens at a time. This

223 + function returns the batch size i.e the maximum number of tokens

224 + the implementation can process at a time.

225 + Return None if there are no such restrictions.

226 + """

227 + raise NotImplementedError

228 +

229 + @abstractmethod

230 + def num_dispatchers(self) -> int:

231 + raise NotImplementedError

232 +

233 + @abstractmethod

234 + def output_is_reduced(self) -> bool:

235 + """

236 + Indicates whether or not the output of finalize is reduced across all

237 + ranks.

238 + """

239 + raise NotImplementedError

240 +

241 +

242 +# TODO: pass FusedMoEParallelConfig in as ctor parameter?

243 +class FusedMoEPrepareAndFinalizeModular(FusedMoEPrepareAndFinalize):

244 + """

245 + An abstract base class for the [Quantize-Prepare] and [Finalize] steps

246 + described above for the Modular case.

247 + """

248 +

249 @abstractmethod

250 def prepare(

251 self,

@@ -198,7 +270,7 @@ class FusedMoEPrepareAndFinalize(ABC):

270 activations, before quantization + dispatching.

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

272 - defer_input_quant: Runtime parameter indicating whether or not to

201 - defer input quantization to the FusedMoEPermuteExpertsUnpermute

273 + defer input quantization to the FusedMoEExpertsModular

274 in cases where the compute kernel expects unquantized inputs

275

276 Returns a tuple of:

@@ -245,7 +317,7 @@ class FusedMoEPrepareAndFinalize(ABC):

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

318 activations, before quantization + dispatching.

319 - defer_input_quant: Runtime parameter indicating whether or not to

248 - defer input quantization to the FusedMoEPermuteExpertsUnpermute

320 + defer input quantization to the FusedMoEExpertsModular

321 in cases where the compute kernel expects unquantized inputs

322

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

@@ -338,56 +410,58 @@ class FusedMoEPrepareAndFinalize(ABC):

410 """

411 raise NotImplementedError

412

341 - @property

342 - @abstractmethod

343 - def activation_format(self) -> FusedMoEActivationFormat:

344 - """

345 - A property indicating the output format of the activations for the

346 - 'prepare' method.

347 - """

348 - raise NotImplementedError

413 +

414 +class FusedMoEPrepareAndFinalizeMonolithic(FusedMoEPrepareAndFinalize):

415 + """

416 + An abstract base class for the [Quantize-Prepare] and [Finalize] steps

417 + described above for the monolithic case.

418 + """

419

420 @abstractmethod

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

421 + def prepare(

422 + self,

423 + a1: torch.Tensor,

424 + router_logits: torch.Tensor,

425 + quant_config: FusedMoEQuantConfig,

426 + defer_input_quant: bool = False,

427 + ) -> PrepareMonolithicResultType:

428 """

353 - The PrepareFinalize All2All implementations generally constrain the

354 - dtype of the topk_ids they support. This function returns the

355 - required topk indices dtype so it can be respected.

356 - Return None if there are no such restrictions.

429 + Optional method for subclasses compatible with monolithic

430 + FusedMoEExpertsModular kernels.

431 +

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

433 + - a1: The (unquantized) input to the MoE layer.

434 + - quant_config: Quantization info provided by the fused experts.

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

436 + defer input quantization to the FusedMoEExpertsModular

437 +

438 + Returns a tuple of:

439 + - quantized + dispatched a.

440 + - Optional quantized + dispatched a1_scales.

441 """

442 raise NotImplementedError

443

444 @abstractmethod

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

445 + def finalize(self, fused_expert_output: torch.Tensor) -> torch.Tensor:

446 """

363 - Some PrepareFinalize All2All implementations are batched. Meaning,

364 - they can process only as set of tokens at a time. This

365 - function returns the batch size i.e the maximum number of tokens

366 - the implementation can process at a time.

367 - Return None if there are no such restrictions.

447 + Optional method for subclasses compatible with monolithic

448 + FusedMoEExpertsModular kernels.

449 +

450 + Perform any combine plus apply weights and perform a reduction on the

451 + fused experts output.

452 + - fused_expert_output: The unweighted, unreduced output of the fused

453 + experts, it will have (M, topk, K) shape.

454 """

455 raise NotImplementedError

456

371 - @abstractmethod

372 - def num_dispatchers(self) -> int:

373 - raise NotImplementedError

457

375 - @abstractmethod

376 - def output_is_reduced(self) -> bool:

377 - """

378 - Indicates whether or not the output of finalize is reduced across all

379 - ranks.

380 - """

381 - raise NotImplementedError

458 +################################################################################

459 +# Experts

460 +################################################################################

461

462

463 # TODO: add supported activations method (return string)

385 -class FusedMoEPermuteExpertsUnpermute(ABC):

386 - """

387 - An abstract base class for the [Permute-Experts-Unpermute] step described

388 - above.

389 - """

390 -

464 +class FusedMoEExperts(ABC):

465 def __init__(

466 self,

467 moe_config: FusedMoEConfig,

@@ -419,6 +493,10 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

493 self.max_num_tokens = max_num_tokens

494 self.num_dispatchers = num_dispatchers

495

496 + @staticmethod

497 + def is_monolithic() -> bool:

498 + raise NotImplementedError("Implemented by subclasses.")

499 +

500 @property

501 def expects_unquantized_inputs(self) -> bool:

502 """

@@ -439,49 +517,6 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

517 """

518 raise NotImplementedError

519

442 - def moe_problem_size(

443 - self,

444 - a1: torch.Tensor,

445 - w1: torch.Tensor,

446 - w2: torch.Tensor,

447 - topk_ids: torch.Tensor,

448 - ) -> tuple[int, int, int, int, int]:

449 - """

450 - Extract the MoE problem size from the given tensor arguments:

451 - - a: The hidden states, input to the MoE layer.

452 - - w1: The first set of expert weights.

453 - - w2: The second set of expert weights.

454 - - topk_ids: The topk ids.

455 -

456 - Note: extracting the problem shape from the weight and activation

457 - tensors is not obvious. It needs to be done this way specifically

458 - due to subtle issues with particular kernels, e.g. the int4 kernels

459 - divide the trailing dimension by two, so it's not "correct" to

460 - extract N or K from the trailing dimension of w1 or w2. Similarly,

461 - some kernels transpose the weights, so this needs to be kept in mind.

462 -

463 - Note: This implementation covers most cases. However, if experts

464 - require a specialized implementation, like MarlinExperts, they are free

465 - to override this function.

466 - """

467 - assert w1.dim() == 3 and w2.dim() == 3

468 - E, N, _ = w1.size()

469 - K = a1.size(-1)

470 -

471 - if a1.dim() == 2:

472 - # Make sure we are using the correct a1 (pre-permute).

473 - assert topk_ids.size(0) == a1.size(0), f"{topk_ids.size(0)} != {a1.size(0)}"

474 - M = a1.size(0)

475 - else:

476 - assert a1.dim() == 3

477 - assert a1.size(0) == E, f"{a1.size(0)} == {E}"

478 - M = a1.size(1) # This is max_num_tokens

479 -

480 - assert topk_ids.dim() == 2

481 - topk = topk_ids.size(1)

482 -

483 - return E, M, N, K, topk

484 -

520 #

521 # Various helpers for registering support for various features.

522 # Used by the oracle to select a particular kernel for a deployment.

@@ -489,7 +524,7 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

524

525 @staticmethod

526 def is_supported_config(

492 - cls: type["FusedMoEPermuteExpertsUnpermute"],

527 + cls: type["FusedMoEExperts"],

528 moe_config: FusedMoEConfig,

529 weight_key: QuantKey | None,

530 activation_key: QuantKey | None,

@@ -512,6 +547,21 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

547 return False, _make_reason(

548 f"parallel config {moe_config.moe_parallel_config}"

549 )

550 + elif not cls._supports_routing_method(

551 + moe_config.routing_method, weight_key, activation_key

552 + ):

553 + return False, _make_reason(f"routing method {moe_config.routing_method}")

554 + elif not cls._supports_router_logits_dtype(

555 + moe_config.router_logits_dtype,

556 + moe_config.routing_method,

557 + ):

558 + return False, _make_reason(

559 + f"router logits dtype {moe_config.router_logits_dtype}"

560 + )

561 + elif not cls._supports_shape(moe_config.hidden_dim):

562 + return False, _make_reason(

563 + f"{moe_config.hidden_dim} hidden dim is not supported"

564 + )

565 elif activation_format != cls.activation_format():

566 return False, _make_reason(f"{activation_format.value} activation format")

567 return True, None

@@ -554,10 +604,48 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

604 @abstractmethod

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

606 """

557 - Whether the kernel supports deployment in expert parallel.

607 + Whether the kernel supports deployment in particular parallel config.

608 +

609 + Can be overriden if a kernel does not support EP, SP or some other

610 + configuration.

611 """

612 raise NotImplementedError

613

614 + @staticmethod

615 + def _supports_routing_method(

616 + routing_method: RoutingMethodType,

617 + weight_key: QuantKey | None,

618 + activation_key: QuantKey | None,

619 + ) -> bool:

620 + """

621 + Whether the kernel supports a routing method (e.g. GroupedTopK).

622 +

623 + Can be overriden by monolithic kernels that execute the router

624 + in addition to the experts if certain routers are not supported.

625 + """

626 + return True

627 +

628 + @staticmethod

629 + def _supports_router_logits_dtype(

630 + router_logits_dtype: torch.dtype | None,

631 + routing_method: RoutingMethodType,

632 + ) -> bool:

633 + """

634 + Whether a kernel supports a particular dtype for router logits input.

635 +

636 + Can be overriden by monolithic kernels that execute the router

637 + in addition to the experts if certain dtypes are not supported.

638 + """

639 + return True

640 +

641 + @staticmethod

642 + def _supports_shape(hidden_dim: int) -> bool:

643 + """

644 + Whether a kernel supports a particular shape. Can be overridden if a kernel

645 + has specific shape requirements.

646 + """

647 + return True

648 +

649 #

650 # Various helpers for accessing quantization parameters from the

651 # quant_config.

@@ -654,6 +742,65 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

742 """

743 return False

744

745 + def enable_chunking(self):

746 + return (

747 + envs.VLLM_ENABLE_FUSED_MOE_ACTIVATION_CHUNKING and self.supports_chunking()

748 + )

749 +

750 +

751 +class FusedMoEExpertsModular(FusedMoEExperts):

752 + """

753 + An abstract base class for the [Permute-Experts-Unpermute] step described

754 + above.

755 + """

756 +

757 + @staticmethod

758 + def is_monolithic() -> bool:

759 + return False

760 +

761 + def moe_problem_size(

762 + self,

763 + a1: torch.Tensor,

764 + w1: torch.Tensor,

765 + w2: torch.Tensor,

766 + topk_ids: torch.Tensor,

767 + ) -> tuple[int, int, int, int, int]:

768 + """

769 + Extract the MoE problem size from the given tensor arguments:

770 + - a: The hidden states, input to the MoE layer.

771 + - w1: The first set of expert weights.

772 + - w2: The second set of expert weights.

773 + - topk_ids: The topk ids.

774 +

775 + Note: extracting the problem shape from the weight and activation

776 + tensors is not obvious. It needs to be done this way specifically

777 + due to subtle issues with particular kernels, e.g. the int4 kernels

778 + divide the trailing dimension by two, so it's not "correct" to

779 + extract N or K from the trailing dimension of w1 or w2. Similarly,

780 + some kernels transpose the weights, so this needs to be kept in mind.

781 +

782 + Note: This implementation covers most cases. However, if experts

783 + require a specialized implementation, like MarlinExperts, they are free

784 + to override this function.

785 + """

786 + assert w1.dim() == 3 and w2.dim() == 3

787 + E, N, _ = w1.size()

788 + K = a1.size(-1)

789 +

790 + if a1.dim() == 2:

791 + # Make sure we are using the correct a1 (pre-permute).

792 + assert topk_ids.size(0) == a1.size(0), f"{topk_ids.size(0)} != {a1.size(0)}"

793 + M = a1.size(0)

794 + else:

795 + assert a1.dim() == 3

796 + assert a1.size(0) == E, f"{a1.size(0)} == {E}"

797 + M = a1.size(1) # This is max_num_tokens

798 +

799 + assert topk_ids.dim() == 2

800 + topk = topk_ids.size(1)

801 +

802 + return E, M, N, K, topk

803 +

804 def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype:

805 """

806 Workspace type: The dtype to use for the workspace tensors.

@@ -726,11 +873,7 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

873 ) -> None:

874 apply_moe_activation(activation, output, input)

875

729 - def enable_chunking(self):

730 - return (

731 - envs.VLLM_ENABLE_FUSED_MOE_ACTIVATION_CHUNKING and self.supports_chunking()

732 - )

733 -

876 + @abstractmethod

877 def finalize_weight_and_reduce_impl(self) -> TopKWeightAndReduce:

878 raise NotImplementedError

879

@@ -791,6 +934,67 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

934 raise NotImplementedError

935

936

937 +class FusedMoEExpertsMonolithic(FusedMoEExperts):

938 + """

939 + An abstract base class for the [Permute-Experts-Unpermute] step described

940 + above, but with the monolithic interface (accepts router logits

941 + rather than topk ids and weights).

942 + """

943 +

944 + @staticmethod

945 + def _supports_routing_method(

946 + routing_method: RoutingMethodType,

947 + weight_key: QuantKey | None,

948 + activation_key: QuantKey | None,

949 + ) -> bool:

950 + """

951 + Whether the kernel supports a routing method (e.g. GroupedTopK).

952 +

953 + Monolithic kernels should explicitly opt-in to support.

954 + """

955 + raise NotImplementedError

956 +

957 + @staticmethod

958 + def _supports_router_logits_dtype(

959 + router_logits_dtype: torch.dtype | None,

960 + routing_method: RoutingMethodType,

961 + ) -> bool:

962 + """

963 + Whether the kernel supports a dtype for router logits.

964 +

965 + Modular kernels should opt-in to support.

966 + """

967 + raise NotImplementedError

968 +

969 + @staticmethod

970 + def is_monolithic() -> bool:

971 + return True

972 +

973 + def apply(

974 + self,

975 + hidden_states: torch.Tensor,

976 + w1: torch.Tensor,

977 + w2: torch.Tensor,

978 + router_logits: torch.Tensor,

979 + activation: MoEActivation,

980 + global_num_experts: int,

981 + expert_map: torch.Tensor | None,

982 + a1q_scale: torch.Tensor | None,

983 + apply_router_weight_on_input: bool,

984 + # grouped topk + fused topk bias parameters

985 + num_expert_group: int | None = None,

986 + e_score_correction_bias: torch.Tensor | None = None,

987 + routed_scaling_factor: float | None = None,

988 + topk_group: int | None = None,

989 + ) -> torch.Tensor:

990 + """

991 + Same as apply(), except uses router_logits as opposed

992 + to the topk_ids and topk_weights. This is useful for kernels

993 + with fused router and fused_experts (e.g. FLASHINFER_TRTLLM).

994 + """

995 + raise NotImplementedError

996 +

997 +

998 def _slice_scales(

999 scales: torch.Tensor | None, start: int, end: int

1000 ) -> torch.Tensor | None:

@@ -802,75 +1006,32 @@ def _slice_scales(

1006 return None

1007

1008

805 -@final

806 -class FusedMoEModularKernel(torch.nn.Module):

807 - """

808 - This class combines a FusedMoEPrepareAndFinalize instance and

809 - a FusedMoEPermuteExpertsUnpermute to provide an interface that

810 - is compatible with the `fused_experts` function in fused_moe.py.

1009 +################################################################################

1010 +# Kernel

1011 +################################################################################

1012

812 - It takes care of managing any required scratch space.

813 -

814 - Note: Instances of this class should only be used for a single model

815 - layer due to any layer specific state that may be used by the component

816 - objects.

817 - """

1013

1014 +@final

1015 +class FusedMoEKernelModularImpl:

1016 def __init__(

1017 self,

821 - prepare_finalize: FusedMoEPrepareAndFinalize,

822 - fused_experts: FusedMoEPermuteExpertsUnpermute,

1018 + prepare_finalize: FusedMoEPrepareAndFinalizeModular,

1019 + fused_experts: FusedMoEExpertsModular,

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

1021 moe_parallel_config: FusedMoEParallelConfig | None = None,

1022 inplace: bool = False,

1023 ):

827 - super().__init__()

1024 self.prepare_finalize = prepare_finalize

1025 self.fused_experts = fused_experts

1026 self.shared_experts = shared_experts

1027 + self.moe_parallel_config = moe_parallel_config

1028 self.inplace = inplace

832 -

833 - # prefer an explicit FusedMoEParallelConfig when available (from

834 - # FusedMoE layers / tests).

835 - # if not provided, assume this kernel is

836 - # running in a non-DP+EP context

837 - self.moe_parallel_config: FusedMoEParallelConfig | None = moe_parallel_config

1029 self.is_dp_ep = (

1030 moe_parallel_config is not None

1031 and moe_parallel_config.dp_size > 1

1032 and moe_parallel_config.use_ep

1033 )

1034

844 - self._post_init_setup()

845 - assert (

846 - prepare_finalize.activation_format == fused_experts.activation_format()

847 - ), (

848 - f"{prepare_finalize.__class__.__name__}."

849 - f"{prepare_finalize.activation_format} == "

850 - f"{fused_experts.__class__.__name__}."

851 - f"{fused_experts.activation_format()}"

852 - )

853 -

854 - def _post_init_setup(self):

855 - """

856 - Resolve any leftover setup dependencies between self.prepare_finalize

857 - and self.fused_experts here.

858 - """

859 - self.prepare_finalize.post_init_setup(self.fused_experts)

860 -

861 - def supports_expert_map(self) -> bool:

862 - """

863 - A flag indicating whether or not this class supports expert maps.

864 - """

865 - return self.fused_experts.supports_expert_map()

866 -

867 - def output_is_reduced(self) -> bool:

868 - """

869 - Indicates whether or not the output of fused MoE kernel

870 - is reduced across all ranks.

871 - """

872 - return self.prepare_finalize.output_is_reduced()

873 -

1035 def _chunk_info(self, M: int) -> tuple[int, int]:

1036 """

1037 Compute number of chunks and chunk size for given M.

@@ -919,7 +1080,7 @@ class FusedMoEModularKernel(torch.nn.Module):

1080 workspace_dtype = self.fused_experts.workspace_dtype(out_dtype)

1081

1082 # Force worst-case allocation in profiling run for

922 - # "mk.FusedMoEModularKernel.Standard" formats where this is only bounded

1083 + # "mk.FusedMoEKernel.Standard" formats where this is only bounded

1084 # by `VLLM_FUSED_MOE_CHUNK_SIZE` and may not be seen during profiling with

1085 # DP+EP due to the random token routing.

1086 is_profile_run = (

@@ -1313,13 +1474,13 @@ class FusedMoEModularKernel(torch.nn.Module):

1474 assert shared_output is not None

1475 return shared_output, output

1476

1316 - def forward(

1477 + def apply(

1478 self,

1479 hidden_states: torch.Tensor,

1480 w1: torch.Tensor,

1481 w2: torch.Tensor,

1321 - topk_weights: torch.Tensor,

1482 topk_ids: torch.Tensor,

1483 + topk_weights: torch.Tensor,

1484 activation: MoEActivation = MoEActivation.SILU,

1485 global_num_experts: int = -1,

1486 expert_map: torch.Tensor | None = None,

@@ -1334,8 +1495,7 @@ class FusedMoEModularKernel(torch.nn.Module):

1495 - hidden_states: (torch.Tensor): The input tensor to the MoE layer.

1496 - w1 (torch.Tensor): The first set of expert weights.

1497 - w2 (torch.Tensor): The second set of expert weights.

1337 - - topk_weights (torch.Tensor): The topk weights applied at the end of

1338 - the layer.

1498 + - topk_weights (torch.Tensor): The topk weights applied at the end of the layer.

1499 - topk_ids (torch.Tensor): A map of row to expert id.

1500 - activation (MoEActivation): The activation function to apply after the first

1501 MoE layer.

@@ -1354,7 +1514,6 @@ class FusedMoEModularKernel(torch.nn.Module):

1514 Returns:

1515 - torch.Tensor: The output tensor after applying the MoE layer.

1516 """

1357 -

1517 if self.inplace:

1518 assert self.shared_experts is None

1519 assert not disable_inplace()

@@ -1400,3 +1559,206 @@ class FusedMoEModularKernel(torch.nn.Module):

1559 apply_router_weight_on_input,

1560 shared_experts_input=shared_experts_input,

1561 )

1562 +

1563 +

1564 +@final

1565 +class FusedMoEKernelMonolithicImpl:

1566 + def __init__(

1567 + self,

1568 + prepare_finalize: FusedMoEPrepareAndFinalizeMonolithic,

1569 + fused_experts: FusedMoEExpertsMonolithic,

1570 + ):

1571 + self.prepare_finalize = prepare_finalize

1572 + self.fused_experts = fused_experts

1573 +

1574 + def apply(

1575 + self,

1576 + hidden_states: torch.Tensor,

1577 + w1: torch.Tensor,

1578 + w2: torch.Tensor,

1579 + router_logits: torch.Tensor,

1580 + activation: MoEActivation,

1581 + global_num_experts: int,

1582 + expert_map: torch.Tensor | None,

1583 + apply_router_weight_on_input: bool,

1584 + # grouped topk + fused topk bias parameters

1585 + num_expert_group: int | None = None,

1586 + e_score_correction_bias: torch.Tensor | None = None,

1587 + routed_scaling_factor: float | None = None,

1588 + topk_group: int | None = None,

1589 + ) -> torch.Tensor:

1590 + """

1591 + Same as forward(), except uses router_logits as opposed

1592 + to the topk_ids and topk_weights. This is used for kernels

1593 + that have fused router + experts (e.g. FLASHINFER_TRTLLM).

1594 + """

1595 +

1596 + # TODO(rob): add inplace support.

1597 + a1q, a1q_scale, router_logits = self.prepare_finalize.prepare(

1598 + hidden_states,

1599 + router_logits=router_logits,

1600 + quant_config=self.fused_experts.quant_config,

1601 + defer_input_quant=self.fused_experts.expects_unquantized_inputs,

1602 + )

1603 +

1604 + fused_out = self.fused_experts.apply(

1605 + hidden_states=a1q,

1606 + w1=w1,

1607 + w2=w2,

1608 + router_logits=router_logits,

1609 + activation=activation,

1610 + global_num_experts=global_num_experts,

1611 + expert_map=expert_map,

1612 + apply_router_weight_on_input=apply_router_weight_on_input,

1613 + a1q_scale=a1q_scale,

1614 + # grouped topk + fused topk bias parameters

1615 + num_expert_group=num_expert_group,

1616 + e_score_correction_bias=e_score_correction_bias,

1617 + routed_scaling_factor=routed_scaling_factor,

1618 + topk_group=topk_group,

1619 + )

1620 +

1621 + output = self.prepare_finalize.finalize(fused_out)

1622 +

1623 + return output

1624 +

1625 +

1626 +@final

1627 +class FusedMoEKernel:

1628 + def __init__(

1629 + self,

1630 + prepare_finalize: FusedMoEPrepareAndFinalize,

1631 + fused_experts: FusedMoEExperts,

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

1633 + moe_parallel_config: FusedMoEParallelConfig | None = None,

1634 + inplace: bool = False,

1635 + ):

1636 + super().__init__()

1637 + self.shared_experts = shared_experts # NOTE: check if we can remove

1638 +

1639 + # Initialize the implementation (monolithic or modular).

1640 + self.impl: FusedMoEKernelModularImpl | FusedMoEKernelMonolithicImpl

1641 + if isinstance(

1642 + prepare_finalize, FusedMoEPrepareAndFinalizeModular

1643 + ) and isinstance(fused_experts, FusedMoEExpertsModular):

1644 + self.impl = FusedMoEKernelModularImpl(

1645 + prepare_finalize,

1646 + fused_experts,

1647 + shared_experts,

1648 + moe_parallel_config,

1649 + inplace,

1650 + )

1651 +

1652 + elif isinstance(

1653 + prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic

1654 + ) and isinstance(fused_experts, FusedMoEExpertsMonolithic):

1655 + assert shared_experts is None

1656 + assert not inplace

1657 + self.impl = FusedMoEKernelMonolithicImpl(

1658 + prepare_finalize,

1659 + fused_experts,

1660 + )

1661 +

1662 + else:

1663 + raise ValueError(

1664 + "prepare_finalize and fused_experts must both be either monolithic "

1665 + f"or non-monolithic but got {prepare_finalize.__class__.__name__} "

1666 + f"and {fused_experts.__class__.__name__}"

1667 + )

1668 +

1669 + self._post_init_setup()

1670 +

1671 + @property

1672 + def is_monolithic(self) -> bool:

1673 + return isinstance(self.impl, FusedMoEKernelMonolithicImpl)

1674 +

1675 + @property

1676 + def prepare_finalize(self) -> FusedMoEPrepareAndFinalize:

1677 + return self.impl.prepare_finalize

1678 +

1679 + @property

1680 + def fused_experts(self) -> FusedMoEExperts:

1681 + return self.impl.fused_experts

1682 +

1683 + def _post_init_setup(self):

1684 + """

1685 + Resolve any leftover setup dependencies between self.prepare_finalize

1686 + and self.fused_experts here.

1687 + """

1688 + self.prepare_finalize.post_init_setup(self.impl.fused_experts)

1689 + assert (

1690 + self.prepare_finalize.activation_format

1691 + == self.fused_experts.activation_format()

1692 + )

1693 +

1694 + def supports_expert_map(self) -> bool:

1695 + """

1696 + A flag indicating whether or not this class supports expert maps.

1697 + """

1698 + return self.fused_experts.supports_expert_map()

1699 +

1700 + def output_is_reduced(self) -> bool:

1701 + """

1702 + Indicates whether or not the output of fused MoE kernel

1703 + is reduced across all ranks.

1704 + """

1705 + return self.prepare_finalize.output_is_reduced()

1706 +

1707 + def apply_monolithic(

1708 + self,

1709 + hidden_states: torch.Tensor,

1710 + w1: torch.Tensor,

1711 + w2: torch.Tensor,

1712 + router_logits: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

1713 + activation: MoEActivation,

1714 + global_num_experts: int,

1715 + expert_map: torch.Tensor | None,

1716 + apply_router_weight_on_input: bool,

1717 + # grouped topk + fused topk bias parameters

1718 + num_expert_group: int | None = None,

1719 + e_score_correction_bias: torch.Tensor | None = None,

1720 + routed_scaling_factor: float | None = None,

1721 + topk_group: int | None = None,

1722 + ) -> torch.Tensor:

1723 + assert isinstance(self.impl, FusedMoEKernelMonolithicImpl)

1724 + return self.impl.apply(

1725 + hidden_states=hidden_states,

1726 + w1=w1,

1727 + w2=w2,

1728 + router_logits=router_logits,

1729 + activation=activation,

1730 + global_num_experts=global_num_experts,

1731 + expert_map=expert_map,

1732 + apply_router_weight_on_input=apply_router_weight_on_input,

1733 + num_expert_group=num_expert_group,

1734 + e_score_correction_bias=e_score_correction_bias,

1735 + routed_scaling_factor=routed_scaling_factor,

1736 + topk_group=topk_group,

1737 + )

1738 +

1739 + def apply(

1740 + self,

1741 + hidden_states: torch.Tensor,

1742 + w1: torch.Tensor,

1743 + w2: torch.Tensor,

1744 + topk_weights: torch.Tensor,

1745 + topk_ids: torch.Tensor,

1746 + activation: MoEActivation,

1747 + global_num_experts: int,

1748 + expert_map: torch.Tensor | None,

1749 + apply_router_weight_on_input: bool,

1750 + shared_experts_input: torch.Tensor | None = None,

1751 + ) -> torch.Tensor:

1752 + assert isinstance(self.impl, FusedMoEKernelModularImpl)

1753 + return self.impl.apply(

1754 + hidden_states=hidden_states,

1755 + w1=w1,

1756 + w2=w2,

1757 + topk_weights=topk_weights,

1758 + topk_ids=topk_ids,

1759 + activation=activation,

1760 + global_num_experts=global_num_experts,

1761 + expert_map=expert_map,

1762 + apply_router_weight_on_input=apply_router_weight_on_input,

1763 + shared_experts_input=shared_experts_input,

1764 + )

vllm/model_executor/layers/fused_moe/mori_prepare_finalize.py

+1 -1

@@ -12,7 +12,7 @@ from vllm.platforms import current_platform

12 logger = init_logger(__name__)

15 -class MoriPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

15 +class MoriPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):

16 """

17 Prepare/Finalize using MoRI kernels.

18 """

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

+41 -89

@@ -18,13 +18,9 @@ from vllm.model_executor.layers.fused_moe.config import (

18 fp8_w8a8_moe_quant_config,

19 fp8_w8a16_moe_quant_config,

20 )

21 -from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe import (

22 - is_supported_config_trtllm_fp8,

23 -)

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

22 FlashinferMoeBackend,

23 get_flashinfer_moe_backend,

27 - make_fp8_moe_alpha_scales_for_fi,

24 prepare_fp8_moe_layer_for_fi,

25 )

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

@@ -103,9 +99,13 @@ def _get_priority_backends(

100 def backend_to_kernel_cls(

101 backend: Fp8MoeBackend,

106 -) -> type[mk.FusedMoEPermuteExpertsUnpermute]:

102 +) -> type[mk.FusedMoEExperts]:

103 if backend == Fp8MoeBackend.FLASHINFER_TRTLLM:

108 - raise NotImplementedError

104 + from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import ( # noqa: E501

105 + TrtLlmFp8Experts,

106 + )

107 +

108 + return TrtLlmFp8Experts

109

110 elif backend == Fp8MoeBackend.FLASHINFER_CUTLASS:

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

@@ -205,13 +205,11 @@ def select_fp8_moe_backend(

205 weight_key: QuantKey | None,

206 activation_key: QuantKey | None,

207 allow_vllm_cutlass: bool = False,

208 -) -> tuple[Fp8MoeBackend, type[mk.FusedMoEPermuteExpertsUnpermute] | None]:

208 +) -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts] | None]:

209 """

210 Select the primary FP8 MoE backend

211 Note: Shape-specific fallbacks may still occur at runtime.

212 """

213 - k_cls: type[mk.FusedMoEPermuteExpertsUnpermute] | None = None

214 -

213 if config.is_lora_enabled:

214 return Fp8MoeBackend.TRITON, backend_to_kernel_cls(Fp8MoeBackend.TRITON)

215

@@ -252,7 +250,7 @@ def select_fp8_moe_backend(

250 weight_key: QuantKey | None,

251 activation_key: QuantKey | None,

252 activation_format: mk.FusedMoEActivationFormat,

255 - ) -> tuple[Fp8MoeBackend, type[mk.FusedMoEPermuteExpertsUnpermute]]:

253 + ) -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]:

254 k_cls = backend_to_kernel_cls(backend)

255 supported, reason = k_cls.is_supported_config(

256 k_cls, config, weight_key, activation_key, activation_format

@@ -287,16 +285,6 @@ def select_fp8_moe_backend(

285 "vLLM CUTLASS FP8 MoE backend is disabled for this configuration."

286 )

287

290 - # Handle FLASHINFER_TRTLLM specially (no kernel class).

291 - if requested_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:

292 - supported, reason = is_supported_config_trtllm_fp8(

293 - config, weight_key, activation_key, activation_format

294 - )

295 - if supported:

296 - logger.info_once(_make_log_backend(requested_backend))

297 - return requested_backend, None

298 - raise ValueError(_make_log_unsupported(requested_backend, reason))

299 -

288 return _return_or_raise(

289 requested_backend, config, weight_key, activation_key, activation_format

290 )

@@ -311,51 +299,32 @@ def select_fp8_moe_backend(

299 elif envs.is_set("VLLM_FLASHINFER_MOE_BACKEND"):

300 # If user is explicit about backend, validate it.

301 fi_backend = get_flashinfer_moe_backend()

314 -

315 - if fi_backend == FlashinferMoeBackend.TENSORRT_LLM:

316 - backend = Fp8MoeBackend.FLASHINFER_TRTLLM

317 - supported, reason = is_supported_config_trtllm_fp8(

318 - config, weight_key, activation_key, activation_format

319 - )

320 - if supported:

321 - logger.info_once(_make_log_backend(backend))

322 - return backend, None

323 - else:

324 - raise ValueError(_make_log_unsupported(backend, reason))

325 -

326 - elif fi_backend == FlashinferMoeBackend.CUTLASS:

302 + if fi_backend == FlashinferMoeBackend.CUTLASS:

303 backend = Fp8MoeBackend.FLASHINFER_CUTLASS

328 - return _return_or_raise(

329 - backend, config, weight_key, activation_key, activation_format

330 - )

331 -

304 + elif fi_backend == FlashinferMoeBackend.TENSORRT_LLM:

305 + backend = Fp8MoeBackend.FLASHINFER_TRTLLM

306 else:

333 - assert fi_backend == FlashinferMoeBackend.CUTEDSL

334 - raise ValueError("FlashInfer MaskedGEMM not supported for FP8")

335 -

307 + raise ValueError(

308 + f"FlashInfer MOE backend {fi_backend} does not support FP8 MoE."

309 + )

310 + k_cls = backend_to_kernel_cls(backend)

311 + return _return_or_raise(

312 + backend, config, weight_key, activation_key, activation_format

313 + )

314 else:

315 # If the user is not explicit about the backend, try both.

316 for backend in [

317 Fp8MoeBackend.FLASHINFER_TRTLLM,

318 Fp8MoeBackend.FLASHINFER_CUTLASS,

319 ]:

342 - if backend == Fp8MoeBackend.FLASHINFER_TRTLLM:

343 - k_cls = None

344 - supported, reason = is_supported_config_trtllm_fp8(

345 - config,

346 - weight_key,

347 - activation_key,

348 - activation_format,

349 - )

350 - else:

351 - k_cls = backend_to_kernel_cls(backend)

352 - supported, reason = k_cls.is_supported_config(

353 - k_cls,

354 - config,

355 - weight_key,

356 - activation_key,

357 - activation_format,

358 - )

320 + k_cls = backend_to_kernel_cls(backend)

321 + supported, reason = k_cls.is_supported_config(

322 + k_cls,

323 + config,

324 + weight_key,

325 + activation_key,

326 + activation_format,

327 + )

328

329 if supported:

330 logger.info_once(_make_log_backend(backend), scope="local")

@@ -408,23 +377,14 @@ def select_fp8_moe_backend(

377

378 # Select kernels in order of backend.

379 for backend in AVAILABLE_BACKENDS:

411 - if backend == Fp8MoeBackend.FLASHINFER_TRTLLM:

412 - k_cls = None

413 - supported, reason = is_supported_config_trtllm_fp8(

414 - config,

415 - weight_key,

416 - activation_key,

417 - activation_format,

418 - )

419 - else:

420 - k_cls = backend_to_kernel_cls(backend)

421 - supported, reason = k_cls.is_supported_config(

422 - k_cls,

423 - config,

424 - weight_key,

425 - activation_key,

426 - activation_format,

427 - )

380 + k_cls = backend_to_kernel_cls(backend)

381 + supported, reason = k_cls.is_supported_config(

382 + k_cls,

383 + config,

384 + weight_key,

385 + activation_key,

386 + activation_format,

387 + )

388

389 if supported:

390 logger.info_once(_make_log_backend(backend), scope="local")

@@ -510,7 +470,7 @@ def make_fp8_moe_quant_config(

470 block_shape: list[int] | None = None,

471 per_act_token_quant: bool = False,

472 per_out_ch_quant: bool = False,

513 -) -> FusedMoEQuantConfig | None:

473 +) -> FusedMoEQuantConfig:

474 """

475 Create FusedMoEQuantConfig for the specified FP8 Backend.

476 The FusedMoEQuantConfig holds the scales that are used

@@ -523,9 +483,6 @@ def make_fp8_moe_quant_config(

483 In a future PR, we will have this function should be

484 a method of the modular kernel itself.

485 """

526 - # TRTLLM does not use Modular Kernel abstraction yet.

527 - if fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:

528 - return None

486

487 # MARLIN is mixed precision W8A16 config.

488 if fp8_backend == Fp8MoeBackend.MARLIN:

@@ -539,12 +496,6 @@ def make_fp8_moe_quant_config(

496 # (alpha = w_scale * a_scale) and inverse a2 scale.

497 if fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS and block_shape is None:

498 assert a1_scale is not None and a2_scale is not None

542 - g1_alphas, g2_alphas = make_fp8_moe_alpha_scales_for_fi(

543 - w1_scale,

544 - a1_scale,

545 - w2_scale,

546 - a2_scale,

547 - )

499 return fp8_w8a8_moe_quant_config(

500 w1_scale=w1_scale,

501 w2_scale=w2_scale,

@@ -552,8 +503,8 @@ def make_fp8_moe_quant_config(

503 a2_scale=a2_scale,

504 a1_gscale=(1.0 / a1_scale),

505 a2_gscale=(1.0 / a2_scale),

555 - g1_alphas=g1_alphas,

556 - g2_alphas=g2_alphas,

506 + g1_alphas=(w1_scale * a1_scale).squeeze(),

507 + g2_alphas=(w2_scale * a2_scale).squeeze(),

508 )

509 # All other backends use normal config.

510 return fp8_w8a8_moe_quant_config(

@@ -570,17 +521,18 @@ def make_fp8_moe_quant_config(

521 def make_fp8_moe_kernel(

522 moe_quant_config: FusedMoEQuantConfig,

523 moe_config: FusedMoEConfig,

573 - experts_cls: type[mk.FusedMoEPermuteExpertsUnpermute],

524 + experts_cls: type[mk.FusedMoEExperts],

525 fp8_backend: Fp8MoeBackend,

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

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

577 -) -> mk.FusedMoEModularKernel:

528 +) -> mk.FusedMoEKernel:

529 # Create Prepare/Finalize.

530 prepare_finalize = maybe_make_prepare_finalize(

531 moe=moe_config,

532 quant_config=moe_quant_config,

533 routing_tables=routing_tables,

534 allow_new_interface=True,

535 + use_monolithic=issubclass(experts_cls, mk.FusedMoEExpertsMonolithic),

536 )

537 assert prepare_finalize is not None

538

@@ -605,7 +557,7 @@ def make_fp8_moe_kernel(

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

558 # if using all2all (for SBO). bnell is making this explicit in

559 # the new MoE runner class.

608 - kernel = mk.FusedMoEModularKernel(

560 + kernel = mk.FusedMoEKernel(

561 prepare_finalize,

562 experts,

563 shared_experts=(

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

+52 -84

@@ -19,7 +19,6 @@ from vllm.model_executor.layers.fused_moe.config import (

19 nvfp4_w4a16_moe_quant_config,

20 )

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

22 - is_supported_config_trtllm,

22 prepare_nvfp4_moe_layer_for_fi_or_cutlass,

23 )

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

@@ -67,39 +66,46 @@ def is_global_sf_supported_for_nvfp4_backend(backend: NvFp4MoeBackend) -> bool:

67 def backend_to_kernel_cls(

68 backend: NvFp4MoeBackend,

70 -) -> type[mk.FusedMoEPermuteExpertsUnpermute]:

69 +) -> list[type[mk.FusedMoEExperts]]:

70 if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM:

72 - raise NotImplementedError(

73 - "FLASHINFER_TRTLLM doesn't support Modular Kernel Interface"

71 + from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (

72 + TrtLlmNvFp4ExpertsModular,

73 + TrtLlmNvFp4ExpertsMonolithic,

74 )

76 + # NOTE: prefer Monolthic > Modular, so return Monolithic first.

77 + return [

78 + TrtLlmNvFp4ExpertsMonolithic,

79 + TrtLlmNvFp4ExpertsModular,

80 + ]

81 +

82 elif backend == NvFp4MoeBackend.FLASHINFER_CUTLASS:

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

84 FlashInferExperts,

85 )

81 - return FlashInferExperts

87 + return [FlashInferExperts]

89 elif backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL:

90 from vllm.model_executor.layers.fused_moe.flashinfer_cutedsl_moe import (

91 FlashInferCuteDSLExperts,

92 )

88 - return FlashInferCuteDSLExperts

94 + return [FlashInferCuteDSLExperts]

96 elif backend == NvFp4MoeBackend.VLLM_CUTLASS:

97 from vllm.model_executor.layers.fused_moe.cutlass_moe import (

98 CutlassExpertsFp4,

99 )

100

95 - return CutlassExpertsFp4

101 + return [CutlassExpertsFp4]

102

103 elif backend == NvFp4MoeBackend.MARLIN:

104 from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (

105 MarlinExperts,

106 )

107

102 - return MarlinExperts

108 + return [MarlinExperts]

109 else:

110 raise ValueError(f"Unknown NvFP4 MoE backend: {backend.value}")

111

@@ -125,7 +131,7 @@ def select_nvfp4_moe_backend(

131 config: FusedMoEConfig,

132 weight_key: QuantKey | None,

133 activation_key: QuantKey | None,

128 -) -> tuple[NvFp4MoeBackend, type[mk.FusedMoEPermuteExpertsUnpermute] | None]:

134 +) -> tuple[NvFp4MoeBackend, type[mk.FusedMoEExperts]]:

135 """

136 Select the primary NvFP4 MoE backend

137 Note: Shape-specific fallbacks may still occur at runtime.

@@ -175,29 +181,21 @@ def select_nvfp4_moe_backend(

181 weight_key: QuantKey | None,

182 activation_key: QuantKey | None,

183 activation_format: mk.FusedMoEActivationFormat,

178 - ) -> tuple[NvFp4MoeBackend, type[mk.FusedMoEPermuteExpertsUnpermute]]:

179 - k_cls = backend_to_kernel_cls(backend)

180 - supported, reason = k_cls.is_supported_config(

181 - k_cls, config, weight_key, activation_key, activation_format

182 - )

183 - if supported:

184 - logger.info_once(_make_log_backend(backend))

185 - return backend, k_cls

184 + ) -> tuple[NvFp4MoeBackend, type[mk.FusedMoEExperts]]:

185 + for k_cls in backend_to_kernel_cls(backend):

186 + supported, reason = k_cls.is_supported_config(

187 + k_cls, config, weight_key, activation_key, activation_format

188 + )

189 + if supported:

190 + logger.info_once(_make_log_backend(backend))

191 + return backend, k_cls

192 +

193 raise ValueError(_make_log_unsupported(backend, reason))

194

195 # Handle explicit moe_backend from user.

196 runner_backend = config.moe_backend

197 if runner_backend != "auto":

198 requested_backend = map_nvfp4_backend(runner_backend)

192 - if requested_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM:

193 - supported, reason = is_supported_config_trtllm(

194 - config, weight_key, activation_key, activation_format

195 - )

196 - if supported:

197 - logger.info_once(_make_log_backend(requested_backend))

198 - return requested_backend, None

199 - raise ValueError(_make_log_unsupported(requested_backend, reason))

200 -

199 return _return_or_raise(

200 requested_backend, config, weight_key, activation_key, activation_format

201 )

@@ -210,36 +208,14 @@ def select_nvfp4_moe_backend(

208

209 elif envs.is_set("VLLM_FLASHINFER_MOE_BACKEND"):

210 # If user is explicit about backend, validate it.

213 - fi_backend = get_flashinfer_moe_backend()

214 -

215 - if fi_backend == FlashinferMoeBackend.TENSORRT_LLM:

216 - backend = NvFp4MoeBackend.FLASHINFER_TRTLLM

217 - supported, reason = is_supported_config_trtllm(

218 - config, weight_key, activation_key, activation_format

219 - )

220 - if supported:

221 - logger.info_once(_make_log_backend(backend))

222 - return backend, None

223 - else:

224 - raise ValueError(_make_log_unsupported(backend, reason))

225 - else:

226 - backend = fi_2_vllm_backend_map[fi_backend]

227 - return _return_or_raise(

228 - backend, config, weight_key, activation_key, activation_format

229 - )

211 + backend = fi_2_vllm_backend_map[get_flashinfer_moe_backend()]

212 + return _return_or_raise(

213 + backend, config, weight_key, activation_key, activation_format

214 + )

215 else:

216 # If the user is not explicit about the backend, try each.

217 for backend in FLASHINFER_NVFP4_MOE_BACKENDS:

233 - if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM:

234 - k_cls = None

235 - supported, reason = is_supported_config_trtllm(

236 - config,

237 - weight_key,

238 - activation_key,

239 - activation_format,

240 - )

241 - else:

242 - k_cls = backend_to_kernel_cls(backend)

218 + for k_cls in backend_to_kernel_cls(backend):

219 supported, reason = k_cls.is_supported_config(

220 k_cls,

221 config,

@@ -247,13 +223,13 @@ def select_nvfp4_moe_backend(

223 activation_key,

224 activation_format,

225 )

250 - if supported:

251 - logger.info_once(_make_log_backend(backend), scope="local")

252 - return backend, None

253 - else:

254 - logger.debug_once(

255 - _make_log_unsupported(backend, reason), scope="local"

256 - )

226 + if supported:

227 + logger.info_once(_make_log_backend(backend), scope="local")

228 + return backend, k_cls

229 + else:

230 + logger.debug_once(

231 + _make_log_unsupported(backend, reason), scope="local"

232 + )

233

234 raise NotImplementedError(

235 "Found VLLM_USE_FLASHINFER_MOE_FP4=1, but no "

@@ -268,16 +244,7 @@ def select_nvfp4_moe_backend(

244

245 # Select kernels in order of backend.

246 for backend in AVAILABLE_BACKENDS:

271 - if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM:

272 - k_cls = None # type: ignore[assignment]

273 - supported, reason = is_supported_config_trtllm(

274 - config,

275 - weight_key,

276 - activation_key,

277 - activation_format,

278 - )

279 - else:

280 - k_cls = backend_to_kernel_cls(backend)

247 + for k_cls in backend_to_kernel_cls(backend):

248 supported, reason = k_cls.is_supported_config(

249 k_cls,

250 config,

@@ -286,11 +253,11 @@ def select_nvfp4_moe_backend(

253 activation_format,

254 )

255

289 - if supported:

290 - logger.info_once(_make_log_backend(backend), scope="local")

291 - return backend, k_cls

292 - else:

293 - logger.debug_once(_make_log_unsupported(backend, reason), scope="local")

256 + if supported:

257 + logger.info_once(_make_log_backend(backend), scope="local")

258 + return backend, k_cls

259 + else:

260 + logger.debug_once(_make_log_unsupported(backend, reason), scope="local")

261

262 raise NotImplementedError(

263 "No NvFp4 MoE backend supports the deployment configuration."

@@ -398,12 +365,8 @@ def make_nvfp4_moe_quant_config(

365 w2_scale_2: torch.Tensor,

366 a13_scale: torch.Tensor,

367 a2_scale: torch.Tensor,

401 -) -> FusedMoEQuantConfig | None:

402 - UNSUPPORTED = [NvFp4MoeBackend.FLASHINFER_TRTLLM]

403 - if backend in UNSUPPORTED:

404 - return None

405 -

406 - elif backend == NvFp4MoeBackend.MARLIN:

368 +) -> FusedMoEQuantConfig:

369 + if backend == NvFp4MoeBackend.MARLIN:

370 return nvfp4_w4a16_moe_quant_config(

371 g1_alphas=w13_scale_2,

372 g2_alphas=w2_scale_2,

@@ -420,22 +383,27 @@ def make_nvfp4_moe_quant_config(

383 a2_gscale=(1.0 / a2_scale),

384 w1_scale=w13_scale,

385 w2_scale=w2_scale,

386 + # NOTE(rob): this is a hack until the MoE kernels

387 + # create their own quant configs. TRTLLM kernel

388 + # does not accept swizzled input quant scales.

389 + is_nvfp4_scale_swizzled=(backend != NvFp4MoeBackend.FLASHINFER_TRTLLM),

390 )

391

392

393 def make_nvfp4_moe_kernel(

394 moe_quant_config: FusedMoEQuantConfig,

395 moe_config: FusedMoEConfig,

429 - experts_cls: type[mk.FusedMoEPermuteExpertsUnpermute],

396 + experts_cls: type[mk.FusedMoEExperts],

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

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

432 -) -> mk.FusedMoEModularKernel:

399 +) -> mk.FusedMoEKernel:

400 # Create Prepare/Finalize.

401 prepare_finalize = maybe_make_prepare_finalize(

402 moe=moe_config,

403 quant_config=moe_quant_config,

404 routing_tables=routing_tables,

405 allow_new_interface=True,

406 + use_monolithic=issubclass(experts_cls, mk.FusedMoEExpertsMonolithic),

407 )

408 assert prepare_finalize is not None

409

@@ -460,7 +428,7 @@ def make_nvfp4_moe_kernel(

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

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

430 # the new MoE runner class.

463 - kernel = mk.FusedMoEModularKernel(

431 + kernel = mk.FusedMoEKernel(

432 prepare_finalize,

433 experts,

434 shared_experts=(

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

+10 -10

@@ -19,7 +19,7 @@ from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe import (

19 is_supported_config_trtllm_bf16,

20 )

21 from vllm.model_executor.layers.fused_moe.prepare_finalize import (

22 - MoEPrepareAndFinalizeNoEP,

22 + MoEPrepareAndFinalizeNoDPEPModular,

23 )

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

25 swap_w13_to_w31,

@@ -209,7 +209,7 @@ def make_unquantized_moe_kernel(

209 backend: UnquantizedMoeBackend,

210 quant_config: FusedMoEQuantConfig,

211 moe_config: FusedMoEConfig,

212 -) -> mk.FusedMoEModularKernel | None:

212 +) -> mk.FusedMoEKernel | None:

213 if backend in UNSUPPORTED_BACKEND:

214 return None

215

@@ -218,8 +218,8 @@ def make_unquantized_moe_kernel(

218 FlashInferExperts,

219 )

220

221 - kernel = mk.FusedMoEModularKernel(

222 - MoEPrepareAndFinalizeNoEP(),

221 + kernel = mk.FusedMoEKernel(

222 + MoEPrepareAndFinalizeNoDPEPModular(),

223 FlashInferExperts(

224 moe_config=moe_config,

225 quant_config=quant_config,

@@ -232,8 +232,8 @@ def make_unquantized_moe_kernel(

232 AiterExperts,

233 )

234

235 - kernel = mk.FusedMoEModularKernel(

236 - MoEPrepareAndFinalizeNoEP(),

235 + kernel = mk.FusedMoEKernel(

236 + MoEPrepareAndFinalizeNoDPEPModular(),

237 AiterExperts(

238 moe_config=moe_config,

239 quant_config=quant_config,

@@ -243,8 +243,8 @@ def make_unquantized_moe_kernel(

243 elif backend == UnquantizedMoeBackend.TRITON:

244 from vllm.model_executor.layers.fused_moe import TritonExperts

245

246 - kernel = mk.FusedMoEModularKernel(

247 - MoEPrepareAndFinalizeNoEP(),

246 + kernel = mk.FusedMoEKernel(

247 + MoEPrepareAndFinalizeNoDPEPModular(),

248 TritonExperts(

249 moe_config=moe_config,

250 quant_config=quant_config,

@@ -254,8 +254,8 @@ def make_unquantized_moe_kernel(

254 elif backend == UnquantizedMoeBackend.XPU:

255 from vllm.model_executor.layers.fused_moe import XPUExperts

256

257 - kernel = mk.FusedMoEModularKernel(

258 - MoEPrepareAndFinalizeNoEP(),

257 + kernel = mk.FusedMoEKernel(

258 + MoEPrepareAndFinalizeNoDPEPModular(),

259 XPUExperts(

260 moe_config=moe_config,

261 quant_config=quant_config,

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

-209

@@ -1,209 +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_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 -

16 -

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

24 - self.is_sequence_parallel = is_sequence_parallel

25 - self._num_dispatchers = num_dispatchers

26 -

27 - @property

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

29 - return mk.FusedMoEActivationFormat.Standard

30 -

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

32 - return None

33 -

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

35 - return None

36 -

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 - """MoE prepare and finalize without expert parallelism."""

135 -

136 - @property

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

138 - return mk.FusedMoEActivationFormat.Standard

139 -

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

141 - return None

142 -

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

144 - return None

145 -

146 - def num_dispatchers(self) -> int:

147 - return 1

148 -

149 - def output_is_reduced(self) -> bool:

150 - return False

151 -

152 - def prepare(

153 - self,

154 - a1: torch.Tensor,

155 - topk_weights: torch.Tensor,

156 - topk_ids: torch.Tensor,

157 - num_experts: int,

158 - expert_map: torch.Tensor | None,

159 - apply_router_weight_on_input: bool,

160 - quant_config: FusedMoEQuantConfig,

161 - defer_input_quant: bool = False,

162 - ) -> mk.PrepareResultType:

163 - if apply_router_weight_on_input:

164 - topk = topk_ids.size(1)

165 - # TODO: this only works for topK=1, will need to update for topK>1

166 - assert topk == 1, (

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

168 - )

169 - # Note: do not use inplace for shared experts overlap

170 - a1 = a1 * topk_weights.to(a1.dtype)

171 -

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

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

174 - if defer_input_quant:

175 - return a1, None, None, None, None

176 -

177 - input_sf = (

178 - quant_config.a1_gscale

179 - if quant_config.use_nvfp4_w4a4

180 - else quant_config.a1_scale

181 - )

182 - a1q, a1q_scale = moe_kernel_quantize_input(

183 - a1,

184 - input_sf,

185 - quant_config.quant_dtype,

186 - quant_config.per_act_token_quant,

187 - quant_config.block_shape,

188 - )

189 -

190 - return a1q, a1q_scale, None, None, None

191 -

192 - def finalize(

193 - self,

194 - output: torch.Tensor,

195 - fused_expert_output: torch.Tensor,

196 - topk_weights: torch.Tensor,

197 - topk_ids: torch.Tensor,

198 - apply_router_weight_on_input: bool,

199 - weight_and_reduce_impl: mk.TopKWeightAndReduce,

200 - ) -> None:

201 - if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):

202 - weight_and_reduce_impl = TopKWeightAndReduceContiguous()

203 - weight_and_reduce_impl.apply(

204 - output=output,

205 - fused_expert_output=fused_expert_output,

206 - topk_weights=topk_weights,

207 - topk_ids=topk_ids,

208 - apply_router_weight_on_input=apply_router_weight_on_input,

209 - )

vllm/model_executor/layers/fused_moe/prepare_finalize/__init__.py new

+22

@@ -0,0 +1,22 @@

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

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

3 +

4 +from vllm.model_executor.layers.fused_moe.prepare_finalize.naive_dp_ep import (

5 + MoEPrepareAndFinalizeNaiveDPEPModular,

6 + MoEPrepareAndFinalizeNaiveDPEPMonolithic,

7 + make_moe_prepare_and_finalize_naive_dp_ep,

8 +)

9 +from vllm.model_executor.layers.fused_moe.prepare_finalize.no_dp_ep import (

10 + MoEPrepareAndFinalizeNoDPEPModular,

11 + MoEPrepareAndFinalizeNoDPEPMonolithic,

12 + make_moe_prepare_and_finalize_no_dp_ep,

13 +)

14 +

15 +__all__ = [

16 + "MoEPrepareAndFinalizeNaiveDPEPMonolithic",

17 + "MoEPrepareAndFinalizeNaiveDPEPModular",

18 + "make_moe_prepare_and_finalize_naive_dp_ep",

19 + "MoEPrepareAndFinalizeNoDPEPMonolithic",

20 + "MoEPrepareAndFinalizeNoDPEPModular",

21 + "make_moe_prepare_and_finalize_no_dp_ep",

22 +]

vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py new

+253

@@ -0,0 +1,253 @@

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

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

3 +import torch

4 +

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

6 +from vllm.distributed import get_ep_group

7 +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

8 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

9 + TopKWeightAndReduceContiguous,

10 + TopKWeightAndReduceDelegate,

11 +)

12 +from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input

13 +from vllm.utils.flashinfer import nvfp4_block_scale_interleave

14 +

15 +

16 +def _quantize_and_setup_dispatch(

17 + a1: torch.Tensor,

18 + quant_config: FusedMoEQuantConfig,

19 + defer_input_quant: bool = False,

20 +) -> tuple[torch.Tensor, list[torch.Tensor] | None]:

21 + # Defer input quantization to the MoE kernel.

22 + if defer_input_quant:

23 + a1q = a1

24 + a1q_scale = None

25 + else:

26 + input_sf = (

27 + quant_config.a1_gscale

28 + if quant_config.use_nvfp4_w4a4

29 + else quant_config.a1_scale

30 + )

31 +

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

33 + # which makes the scales tensor different shape than

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

35 + # delay the swizzling until after the A2A.

36 + a1q, a1q_scale = a1q, a1q_scale = moe_kernel_quantize_input(

37 + a1,

38 + input_sf,

39 + quant_dtype=quant_config.quant_dtype,

40 + per_act_token_quant=quant_config.per_act_token_quant,

41 + block_shape=quant_config.block_shape,

42 + is_fp4_scale_swizzled=False,

43 + )

44 +

45 + # Skip gathering scales if we have static quantization

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

47 + # if quantization is deferred.

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

49 + scales = None if skip_gather_scales else [a1q_scale]

50 +

51 + return a1q, scales

52 +

53 +

54 +def _unwrap_scale_and_prepare_for_moe(

55 + scales: list[torch.Tensor] | None,

56 + quant_config: FusedMoEQuantConfig,

57 +) -> torch.Tensor:

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

59 + a1q_scale = scales[0]

60 + # Apply swizzling after a2a if the MoE kernel needs it.

61 + if quant_config.quant_dtype == "nvfp4" and quant_config.is_nvfp4_scale_swizzled:

62 + assert a1q_scale is not None

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

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

65 + a1q_scale = nvfp4_block_scale_interleave(a1q_scale)

66 +

67 + return a1q_scale

68 +

69 +

70 +class MoEPrepareAndFinalizeNaiveDPEPModular(mk.FusedMoEPrepareAndFinalizeModular):

71 + """

72 + Naive Prepare/Finalize for Dp/Ep case for Modular Kernels.

73 +

74 + Uses Torch AR/RS or AR for dispatch/combine operations, applied

75 + to the topk weights and ids.

76 + """

77 +

78 + def __init__(

79 + self,

80 + is_sequence_parallel: bool = False,

81 + num_dispatchers: int = 1,

82 + ) -> None:

83 + super().__init__()

84 + self.is_sequence_parallel = is_sequence_parallel

85 + self._num_dispatchers = num_dispatchers

86 +

87 + @property

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

89 + return mk.FusedMoEActivationFormat.Standard

90 +

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

92 + return None

93 +

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

95 + return None

96 +

97 + def num_dispatchers(self) -> int:

98 + return self._num_dispatchers

99 +

100 + def output_is_reduced(self) -> bool:

101 + return False

102 +

103 + def prepare(

104 + self,

105 + a1: torch.Tensor,

106 + topk_weights: torch.Tensor,

107 + topk_ids: torch.Tensor,

108 + num_experts: int,

109 + expert_map: torch.Tensor | None,

110 + apply_router_weight_on_input: bool,

111 + quant_config: FusedMoEQuantConfig,

112 + defer_input_quant: bool = False,

113 + ) -> mk.PrepareResultType:

114 + """Quantize and Dispatch Topk Weights and Topk Ids."""

115 +

116 + if apply_router_weight_on_input:

117 + topk = topk_ids.size(1)

118 + assert topk == 1, (

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

120 + )

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

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

123 +

124 + a1q, scales = _quantize_and_setup_dispatch(a1, quant_config, defer_input_quant)

125 +

126 + res = get_ep_group().dispatch(

127 + a1q,

128 + topk_weights,

129 + topk_ids,

130 + is_sequence_parallel=self.is_sequence_parallel,

131 + extra_tensors=scales,

132 + )

133 +

134 + if scales is None:

135 + a1q, topk_weights, topk_ids = res

136 + a1q_scale = None

137 + else:

138 + a1q, topk_weights, topk_ids, scales = res

139 + a1q_scale = _unwrap_scale_and_prepare_for_moe(scales, quant_config)

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 isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):

153 + weight_and_reduce_impl = TopKWeightAndReduceContiguous()

154 +

155 + out = weight_and_reduce_impl.apply(

156 + output=None,

157 + fused_expert_output=fused_expert_output,

158 + topk_weights=topk_weights,

159 + topk_ids=topk_ids,

160 + apply_router_weight_on_input=apply_router_weight_on_input,

161 + )

162 +

163 + output.copy_(

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

165 + )

166 +

167 +

168 +class MoEPrepareAndFinalizeNaiveDPEPMonolithic(mk.FusedMoEPrepareAndFinalizeMonolithic):

169 + """

170 + Naive Prepare/Finalize for Dp/Ep case for Modular Kernels.

171 +

172 + Uses Torch AR/RS or AR for dispatch/combine operations, applied

173 + to the router logits (the MoE kernel runs the router internally).

174 + """

175 +

176 + def __init__(

177 + self,

178 + is_sequence_parallel: bool = False,

179 + num_dispatchers: int = 1,

180 + ) -> None:

181 + super().__init__()

182 + self.is_sequence_parallel = is_sequence_parallel

183 + self._num_dispatchers = num_dispatchers

184 +

185 + @property

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

187 + return mk.FusedMoEActivationFormat.Standard

188 +

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

190 + return None

191 +

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

193 + return None

194 +

195 + def num_dispatchers(self) -> int:

196 + return self._num_dispatchers

197 +

198 + def output_is_reduced(self) -> bool:

199 + return False

200 +

201 + def prepare(

202 + self,

203 + a1: torch.Tensor,

204 + router_logits: torch.Tensor,

205 + quant_config: FusedMoEQuantConfig,

206 + defer_input_quant: bool = False,

207 + ) -> mk.PrepareMonolithicResultType:

208 + """Quantize and Dispatch Router Logits."""

209 +

210 + a1q, scales = _quantize_and_setup_dispatch(a1, quant_config, defer_input_quant)

211 +

212 + res = get_ep_group().dispatch_router_logits(

213 + a1q,

214 + router_logits,

215 + is_sequence_parallel=self.is_sequence_parallel,

216 + extra_tensors=scales,

217 + )

218 +

219 + if scales is None:

220 + a1q, router_logits = res

221 + a1q_scale = None

222 + else:

223 + a1q, router_logits, scales = res

224 + a1q_scale = _unwrap_scale_and_prepare_for_moe(scales, quant_config)

225 +

226 + return a1q, a1q_scale, router_logits

227 +

228 + def finalize(

229 + self,

230 + fused_expert_output: torch.Tensor,

231 + ) -> torch.Tensor:

232 + out = get_ep_group().combine(

233 + fused_expert_output, is_sequence_parallel=self.is_sequence_parallel

234 + )

235 + return out

236 +

237 +

238 +def make_moe_prepare_and_finalize_naive_dp_ep(

239 + use_monolithic: bool,

240 + is_sequence_parallel: bool = False,

241 + num_dispatchers: int = 1,

242 +) -> MoEPrepareAndFinalizeNaiveDPEPModular | MoEPrepareAndFinalizeNaiveDPEPMonolithic:

243 + return (

244 + MoEPrepareAndFinalizeNaiveDPEPMonolithic(

245 + is_sequence_parallel=is_sequence_parallel,

246 + num_dispatchers=num_dispatchers,

247 + )

248 + if use_monolithic

249 + else MoEPrepareAndFinalizeNaiveDPEPModular(

250 + is_sequence_parallel=is_sequence_parallel,

251 + num_dispatchers=num_dispatchers,

252 + )

253 + )

vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py new

+141

@@ -0,0 +1,141 @@

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

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

3 +import torch

4 +

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

6 +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

7 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

8 + TopKWeightAndReduceContiguous,

9 + TopKWeightAndReduceDelegate,

10 +)

11 +from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input

12 +

13 +

14 +def _quantize_input(

15 + a1: torch.Tensor,

16 + quant_config: FusedMoEQuantConfig,

17 + defer_input_quant: bool = False,

18 +) -> tuple[torch.Tensor, torch.Tensor | None]:

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

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

21 + if defer_input_quant:

22 + return a1, None

23 +

24 + input_sf = (

25 + quant_config.a1_gscale if quant_config.use_nvfp4_w4a4 else quant_config.a1_scale

26 + )

27 + a1q, a1q_scale = moe_kernel_quantize_input(

28 + a1,

29 + input_sf,

30 + quant_dtype=quant_config.quant_dtype,

31 + per_act_token_quant=quant_config.per_act_token_quant,

32 + block_shape=quant_config.block_shape,

33 + is_fp4_scale_swizzled=quant_config.is_nvfp4_scale_swizzled,

34 + )

35 +

36 + return a1q, a1q_scale

37 +

38 +

39 +class MoEPrepareAndFinalizeNoDPEPModular(mk.FusedMoEPrepareAndFinalizeModular):

40 + @property

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

42 + return mk.FusedMoEActivationFormat.Standard

43 +

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

45 + return None

46 +

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

48 + return None

49 +

50 + def num_dispatchers(self) -> int:

51 + return 1

52 +

53 + def output_is_reduced(self) -> bool:

54 + return False

55 +

56 + def prepare(

57 + self,

58 + a1: torch.Tensor,

59 + topk_weights: torch.Tensor,

60 + topk_ids: torch.Tensor,

61 + num_experts: int,

62 + expert_map: torch.Tensor | None,

63 + apply_router_weight_on_input: bool,

64 + quant_config: FusedMoEQuantConfig,

65 + defer_input_quant: bool = False,

66 + ) -> mk.PrepareResultType:

67 + if apply_router_weight_on_input:

68 + topk = topk_ids.size(1)

69 + # TODO: this only works for topK=1, will need to update for topK>1

70 + assert topk == 1, (

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

72 + )

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

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

75 +

76 + a1q, a1q_scale = _quantize_input(a1, quant_config, defer_input_quant)

77 +

78 + return a1q, a1q_scale, None, None, None

79 +

80 + def finalize(

81 + self,

82 + output: torch.Tensor,

83 + fused_expert_output: torch.Tensor,

84 + topk_weights: torch.Tensor,

85 + topk_ids: torch.Tensor,

86 + apply_router_weight_on_input: bool,

87 + weight_and_reduce_impl: mk.TopKWeightAndReduce,

88 + ) -> None:

89 + if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):

90 + weight_and_reduce_impl = TopKWeightAndReduceContiguous()

91 + weight_and_reduce_impl.apply(

92 + output=output,

93 + fused_expert_output=fused_expert_output,

94 + topk_weights=topk_weights,

95 + topk_ids=topk_ids,

96 + apply_router_weight_on_input=apply_router_weight_on_input,

97 + )

98 +

99 +

100 +class MoEPrepareAndFinalizeNoDPEPMonolithic(mk.FusedMoEPrepareAndFinalizeMonolithic):

101 + @property

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

103 + return mk.FusedMoEActivationFormat.Standard

104 +

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

106 + return None

107 +

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

109 + return None

110 +

111 + def num_dispatchers(self) -> int:

112 + return 1

113 +

114 + def output_is_reduced(self) -> bool:

115 + return False

116 +

117 + def prepare(

118 + self,

119 + a1: torch.Tensor,

120 + router_logits: torch.Tensor,

121 + quant_config: FusedMoEQuantConfig,

122 + defer_input_quant: bool = False,

123 + ) -> mk.PrepareMonolithicResultType:

124 + a1q, a1q_scale = _quantize_input(a1, quant_config, defer_input_quant)

125 + return a1q, a1q_scale, router_logits

126 +

127 + def finalize(

128 + self,

129 + fused_expert_output: torch.Tensor,

130 + ) -> torch.Tensor:

131 + return fused_expert_output

132 +

133 +

134 +def make_moe_prepare_and_finalize_no_dp_ep(

135 + use_monolithic: bool,

136 +) -> MoEPrepareAndFinalizeNoDPEPModular | MoEPrepareAndFinalizeNoDPEPMonolithic:

137 + return (

138 + MoEPrepareAndFinalizeNoDPEPMonolithic()

139 + if use_monolithic

140 + else MoEPrepareAndFinalizeNoDPEPModular()

141 + )

vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py

+1 -1

@@ -292,7 +292,7 @@ def rocm_aiter_fused_experts(

292 )

293

294

295 -class AiterExperts(mk.FusedMoEPermuteExpertsUnpermute):

295 +class AiterExperts(mk.FusedMoEExpertsModular):

296 @property

297 def expects_unquantized_inputs(self) -> bool:

298 return True

vllm/model_executor/layers/fused_moe/router/base_router.py

+1 -1

@@ -64,7 +64,7 @@ if current_platform.is_cuda_alike():

65 # TODO(bowen): When using `FusedMoEModularKernel`, this

66 # can be done in a more unified way, since

67 - # `FusedMoEPrepareAndFinalize` will return the expert

67 + # `FusedMoEPrepareAndFinalizeModular` will return the expert

68 # token count, in some cases directly from the kernel.

69 # However, now there are many code paths not using

70 # the modular kernel, e.g. calling `fused_experts`,

vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py

+16 -53

@@ -320,8 +320,8 @@ class DefaultMoERunner(MoERunner):

320 """

321 assert self.quant_method is not None

322 return (

323 - self.quant_method.moe_mk is not None

324 - and self.quant_method.moe_mk.output_is_reduced()

323 + self.quant_method.moe_kernel is not None

324 + and self.quant_method.moe_kernel.output_is_reduced()

325 )

326

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

@@ -640,45 +640,6 @@ class DefaultMoERunner(MoERunner):

640 )

641

642 with sp_ctx:

643 - extra_tensors = None

644 - if do_naive_dispatch_combine:

645 - post_quant_allgather = (

646 - self.quant_method is not None

647 - and self.moe_config.dp_size > 1

648 - and self.moe_config.use_ep

649 - and getattr(self.quant_method, "do_post_quant_allgather", False)

650 - )

651 - if post_quant_allgather:

652 - hidden_states_to_dispatch, extra_tensors = (

653 - self.quant_method.prepare_dp_allgather_tensor(

654 - layer, hidden_states, router_logits

655 - )

656 - )

657 - else:

658 - hidden_states_to_dispatch = hidden_states

659 -

660 - dispatch_res = get_ep_group().dispatch_router_logits(

661 - hidden_states_to_dispatch,

662 - router_logits,

663 - self.moe_config.is_sequence_parallel,

664 - extra_tensors=extra_tensors,

665 - )

666 - if extra_tensors is not None:

667 - (

668 - orig_hidden_states,

669 - router_logits,

670 - extra_tensors_combined,

671 - ) = dispatch_res

672 - hidden_states_combined = (

673 - orig_hidden_states,

674 - extra_tensors_combined[0],

675 - )

676 - else:

677 - hidden_states_combined, router_logits = dispatch_res

678 - orig_hidden_states = hidden_states_combined

679 - else:

680 - orig_hidden_states = hidden_states

681 -

643 # Run shared experts before matrix multiply.

644 # because matrix multiply maybe modify the hidden_states.

645 if has_separate_shared_experts and not use_shared_experts_stream:

@@ -688,6 +649,17 @@ class DefaultMoERunner(MoERunner):

649 )

650 shared_output = self.shared_experts(shared_input)

651

652 + # For naive dispatch/combine Dp/Ep, dispatch the hidden states and

653 + # router logits to all experts.

654 + # NOTE: this will be removed once all kernels are migrated into the

655 + # MoEKernel framework.

656 + if do_naive_dispatch_combine:

657 + hidden_states, router_logits = get_ep_group().dispatch_router_logits(

658 + hidden_states,

659 + router_logits,

660 + self.moe_config.is_sequence_parallel,

661 + )

662 +

663 # NOTE: Similar with DP, PCP also needs dispatch and combine. For

664 # simplicity, AgRsAll2All was added separately for PCP here. Maybe

665 # we should modify All2AllManager abstract to better support PCP.

@@ -701,31 +673,22 @@ class DefaultMoERunner(MoERunner):

673 dim=0,

674 )

675

704 - # TODO(bnell): deal with fp4 flashinfer tuple hidden states hack (#30014).

705 - # Figure out nicer way to do this.

706 - if do_naive_dispatch_combine:

707 - x = hidden_states_combined

708 - x_orig = orig_hidden_states

709 - else:

710 - x = hidden_states

711 - x_orig = hidden_states

712 -

676 # Matrix multiply.

677 if self.quant_method.is_monolithic:

678 final_hidden_states = self.quant_method.apply_monolithic(

679 layer=layer,

717 - x=x,

680 + x=hidden_states,

681 router_logits=router_logits,

682 )

683 else:

684 topk_weights, topk_ids = self.router.select_experts(

722 - hidden_states=x_orig,

685 + hidden_states=hidden_states,

686 router_logits=router_logits,

687 )

688

689 final_hidden_states = self.quant_method.apply(

690 layer=layer,

728 - x=x, # The type signture of this is wrong due to the hack.

691 + x=hidden_states,

692 topk_weights=topk_weights,

693 topk_ids=topk_ids,

694 shared_experts_input=shared_input,

vllm/model_executor/layers/fused_moe/topk_weight_and_reduce.py

+2 -2

@@ -10,7 +10,7 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk

11 class TopKWeightAndReduceDelegate(mk.TopKWeightAndReduce):

12 """

13 - Useful in the case when some FusedMoEPermuteExpertsUnpermute

13 + Useful in the case when some FusedMoEExpertsModular

14 implementation does not perform weight application and reduction

15 but cannot address the needs of all the compatible PrepareAndFinalize

16 implementations.

@@ -62,7 +62,7 @@ class TopKWeightAndReduceNoOP(mk.TopKWeightAndReduce):

62 if output is None:

63 return fused_expert_output

65 - # MoEPrepareAndFinalizeNoEP needs the output to be in the `output`

65 + # MoEPrepareAndFinalizeNoDPEPModular needs the output to be in the `output`

66 # tensor.

67 assert output.size() == fused_expert_output.size(), (

68 "output shape is expected to match the fused_expert_output shape. "

vllm/model_executor/layers/fused_moe/triton_cutlass_moe.py

+3 -3

@@ -32,8 +32,8 @@ class TritonOrCutlassExperts(FallbackExperts):

33 @staticmethod

34 def get_clses() -> tuple[

35 - type[mk.FusedMoEPermuteExpertsUnpermute],

36 - type[mk.FusedMoEPermuteExpertsUnpermute],

35 + type[mk.FusedMoEExpertsModular],

36 + type[mk.FusedMoEExpertsModular],

37 ]:

38 return (CutlassExpertsFp8, TritonExperts)

@@ -77,7 +77,7 @@ class TritonOrCutlassExperts(FallbackExperts):

77 hidden_states: torch.Tensor,

78 w1: torch.Tensor,

79 w2: torch.Tensor,

80 - ) -> mk.FusedMoEPermuteExpertsUnpermute:

80 + ) -> mk.FusedMoEExpertsModular:

81 # Small batch fallback for sm100.

82 if self.is_sm100 and hidden_states.shape[0] <= 8:

83 return self.fallback_experts

vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py

+3 -3

@@ -32,8 +32,8 @@ class TritonOrDeepGemmExperts(FallbackExperts):

33 @staticmethod

34 def get_clses() -> tuple[

35 - type[mk.FusedMoEPermuteExpertsUnpermute],

36 - type[mk.FusedMoEPermuteExpertsUnpermute],

35 + type[mk.FusedMoEExpertsModular],

36 + type[mk.FusedMoEExpertsModular],

37 ]:

38 return (DeepGemmExperts, TritonExperts)

@@ -79,7 +79,7 @@ class TritonOrDeepGemmExperts(FallbackExperts):

79 hidden_states: torch.Tensor,

80 w1: torch.Tensor,

81 w2: torch.Tensor,

82 - ) -> mk.FusedMoEPermuteExpertsUnpermute:

82 + ) -> mk.FusedMoEExpertsModular:

83 if is_deep_gemm_e8m0_used() or _valid_deep_gemm(hidden_states, w1, w2):

84 return self.experts

85 else:

vllm/model_executor/layers/fused_moe/trtllm_moe.py

+1 -1

@@ -18,7 +18,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

18 )

21 -class TrtLlmGenExperts(mk.FusedMoEPermuteExpertsUnpermute):

21 +class TrtLlmGenExperts(mk.FusedMoEExpertsModular):

22 """TensorRT-LLM-based fused MoE expert implementation."""

24 def __init__(

vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py

+7 -7

@@ -24,8 +24,8 @@ from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (

24 )

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

26 FusedMoEActivationFormat,

27 - FusedMoEPermuteExpertsUnpermute,

28 - FusedMoEPrepareAndFinalize,

27 + FusedMoEExpertsModular,

28 + FusedMoEPrepareAndFinalizeModular,

29 )

30 from vllm.model_executor.layers.fused_moe.oracle.unquantized import (

31 UnquantizedMoeBackend,

@@ -70,7 +70,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

70 self.rocm_aiter_moe_enabled = (

71 rocm_aiter_ops.is_fused_moe_enabled() and moe.is_act_and_mul

72 )

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

73 + self.kernel: mk.FusedMoEKernel | None = None

74 self._is_monolithic = (

75 current_platform.is_cpu()

76 or self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM

@@ -107,7 +107,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

107 def maybe_make_prepare_finalize(

108 self,

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

110 - ) -> FusedMoEPrepareAndFinalize | None:

110 + ) -> FusedMoEPrepareAndFinalizeModular | None:

111 if self.unquantized_backend == UnquantizedMoeBackend.AITER:

112 return None

113 else:

@@ -115,9 +115,9 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

115

116 def select_gemm_impl(

117 self,

118 - prepare_finalize: FusedMoEPrepareAndFinalize,

118 + prepare_finalize: FusedMoEPrepareAndFinalizeModular,

119 layer: torch.nn.Module,

120 - ) -> FusedMoEPermuteExpertsUnpermute:

120 + ) -> FusedMoEExpertsModular:

121 assert self.moe_quant_config is not None

122 if (

123 prepare_finalize.activation_format

@@ -325,7 +325,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

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

326 assert self.kernel is not None

327

328 - return self.kernel(

328 + return self.kernel.apply(

329 hidden_states=x,

330 w1=layer.w13_weight,

331 w2=layer.w2_weight,

vllm/model_executor/layers/fused_moe/xpu_fused_moe.py

+1 -1

@@ -23,7 +23,7 @@ if current_platform.is_xpu():

23 from vllm_xpu_kernels.fused_moe_interface import xpu_fused_moe

26 -class XPUExperts(mk.FusedMoEPermuteExpertsUnpermute):

26 +class XPUExperts(mk.FusedMoEExpertsModular):

27 def __init__(

28 self,

29 moe_config: FusedMoEConfig,

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

+71 -168

@@ -19,8 +19,8 @@ from vllm.logger import init_logger

19 from vllm.model_executor.layers.fused_moe import (

20 FusedMoE,

21 FusedMoEActivationFormat,

22 + FusedMoEExpertsModular,

23 FusedMoEMethodBase,

23 - FusedMoEPermuteExpertsUnpermute,

24 FusedMoeWeightScaleSupported,

25 UnquantizedFusedMoEMethod,

26 )

@@ -40,7 +40,6 @@ from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (

40 fused_marlin_moe,

41 )

42 from vllm.model_executor.layers.fused_moe.oracle.fp8 import (

43 - Fp8MoeBackend,

43 convert_to_fp8_moe_kernel_format,

44 make_fp8_moe_kernel,

45 make_fp8_moe_quant_config,

@@ -59,18 +58,11 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compress

58 WNA16_SUPPORTED_BITS,

59 WNA16_SUPPORTED_TYPES_MAP,

60 )

62 -from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (

63 - flashinfer_trtllm_fp4_moe,

64 - flashinfer_trtllm_fp4_routed_moe,

65 -)

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

62 flashinfer_trtllm_mxint4_moe,

63 is_flashinfer_mxint4_moe_available,

64 prepare_static_weights_for_trtllm_mxint4_moe,

65 )

71 -from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

72 - apply_fi_trtllm_fp8_per_tensor_moe,

73 -)

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

67 process_fp8_input_tensor_strategy_moe,

68 process_fp8_weight_tensor_strategy_moe,

@@ -336,7 +328,7 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):

328

329 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

330 if self.moe_quant_config is not None:

339 - self.moe_mk = make_nvfp4_moe_kernel(

331 + self.moe_kernel = make_nvfp4_moe_kernel(

332 moe_quant_config=self.moe_quant_config,

333 moe_config=self.moe,

334 experts_cls=self.experts_cls,

@@ -352,8 +344,8 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):

344 topk_ids: torch.Tensor,

345 shared_experts_input: torch.Tensor | None,

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

355 - assert self.moe_mk is not None

356 - return self.moe_mk(

347 + assert self.moe_kernel is not None

348 + return self.moe_kernel.apply(

349 x,

350 layer.w13_weight,

351 layer.w2_weight,

@@ -562,43 +554,27 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

554 layer.w13_input_scale = a13_scale

555 layer.w2_input_scale = a2_scale

556

565 - # Setup modular kernel for TP case and naive DP/EP case.

566 - # In non-naive DP/EP case, we will create a ModularKernelMethod.

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

568 - # in both cases.

557 + # Setup modular kernel.

558 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

570 - if self.moe_quant_config:

571 - assert self.experts_cls is not None

572 - self.moe_mk = make_nvfp4_moe_kernel(

573 - moe_quant_config=self.moe_quant_config,

574 - moe_config=self.moe,

575 - experts_cls=self.experts_cls,

576 - shared_experts=layer.shared_experts,

577 - routing_tables=layer._maybe_init_expert_routing_tables(),

578 - )

559 + assert self.experts_cls is not None

560 + self.moe_kernel = make_nvfp4_moe_kernel(

561 + moe_quant_config=self.moe_quant_config,

562 + moe_config=self.moe,

563 + experts_cls=self.experts_cls,

564 + shared_experts=layer.shared_experts,

565 + routing_tables=layer._maybe_init_expert_routing_tables(),

566 + )

567

568 def maybe_make_prepare_finalize(

569 self,

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

583 - ) -> mk.FusedMoEPrepareAndFinalize | None:

584 - raise ValueError(

585 - f"{self.__class__.__name__} uses the new modular kernel initialization "

586 - "logic. This function should not be called."

587 - )

588 -

589 - def select_gemm_impl(

590 - self,

591 - prepare_finalize: mk.FusedMoEPrepareAndFinalize,

592 - layer: torch.nn.Module,

593 - ) -> mk.FusedMoEPermuteExpertsUnpermute:

571 + ) -> mk.FusedMoEPrepareAndFinalizeModular | None:

572 raise ValueError(

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

574 "logic. This function should not be called."

575 )

576

599 - def get_fused_moe_quant_config(

600 - self, layer: torch.nn.Module

601 - ) -> FusedMoEQuantConfig | None:

577 + def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:

578 return make_nvfp4_moe_quant_config(

579 backend=self.nvfp4_backend,

580 w13_scale=layer.w13_weight_scale,

@@ -609,13 +585,6 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

585 a2_scale=layer.w2_input_scale,

586 )

587

612 - @property

613 - def is_monolithic(self) -> bool:

614 - return (

615 - self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM

616 - and not self.moe.moe_parallel_config.enable_eplb

617 - )

618 -

588 def apply_monolithic(

589 self,

590 layer: FusedMoE,

@@ -623,24 +592,20 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

592 router_logits: torch.Tensor,

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

594 assert self.is_monolithic

626 - assert layer.activation == MoEActivation.SILU, (

627 - f"Only SiLU activation is supported, not {layer.activation}."

628 - )

629 - assert (

630 - self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM

631 - and not layer.enable_eplb

632 - )

633 - return flashinfer_trtllm_fp4_moe(

634 - layer=layer,

635 - x=x,

636 - router_logits=router_logits,

637 - top_k=layer.top_k,

595 + assert self.moe_kernel is not None

596 + return self.moe_kernel.apply_monolithic(

597 + x,

598 + layer.w13_weight,

599 + layer.w2_weight,

600 + router_logits,

601 activation=layer.activation,

602 global_num_experts=layer.global_num_experts,

603 + expert_map=layer.expert_map,

604 + apply_router_weight_on_input=layer.apply_router_weight_on_input,

605 num_expert_group=layer.num_expert_group,

606 topk_group=layer.topk_group,

642 - custom_routing_function=layer.custom_routing_function,

607 e_score_correction_bias=layer.e_score_correction_bias,

608 + routed_scaling_factor=layer.routed_scaling_factor,

609 )

610

611 def apply(

@@ -651,34 +616,19 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

616 topk_ids: torch.Tensor,

617 shared_experts_input: torch.Tensor | None,

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

654 - assert not self.is_monolithic

655 -

656 - # EPLB path

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

658 - assert layer.enable_eplb

659 - return flashinfer_trtllm_fp4_routed_moe(

660 - layer=layer,

661 - x=x,

662 - topk_ids=topk_ids,

663 - topk_weights=topk_weights,

664 - top_k=layer.top_k,

665 - activation=layer.activation,

666 - global_num_experts=layer.global_num_experts,

667 - )

668 - else:

669 - assert self.moe_mk is not None

670 - return self.moe_mk(

671 - x,

672 - layer.w13_weight,

673 - layer.w2_weight,

674 - topk_weights,

675 - topk_ids,

676 - activation=layer.activation,

677 - global_num_experts=layer.global_num_experts,

678 - expert_map=layer.expert_map,

679 - apply_router_weight_on_input=layer.apply_router_weight_on_input,

680 - shared_experts_input=shared_experts_input,

681 - )

619 + assert self.moe_kernel is not None

620 + return self.moe_kernel.apply(

621 + x,

622 + layer.w13_weight,

623 + layer.w2_weight,

624 + topk_weights,

625 + topk_ids,

626 + activation=layer.activation,

627 + global_num_experts=layer.global_num_experts,

628 + expert_map=layer.expert_map,

629 + apply_router_weight_on_input=layer.apply_router_weight_on_input,

630 + shared_experts_input=shared_experts_input,

631 + )

632

633

634 class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

@@ -966,7 +916,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

916 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

917 if self.moe_quant_config:

918 assert self.experts_cls is not None

969 - self.moe_mk = make_fp8_moe_kernel(

919 + self.moe_kernel = make_fp8_moe_kernel(

920 moe_quant_config=self.moe_quant_config,

921 moe_config=self.moe,

922 fp8_backend=self.fp8_backend,

@@ -978,94 +928,47 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

928 def maybe_make_prepare_finalize(

929 self,

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

981 - ) -> mk.FusedMoEPrepareAndFinalize | None:

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 select_gemm_impl(

988 - self,

989 - prepare_finalize: mk.FusedMoEPrepareAndFinalize,

990 - layer: torch.nn.Module,

991 - ) -> mk.FusedMoEPermuteExpertsUnpermute:

931 + ) -> mk.FusedMoEPrepareAndFinalizeModular | None:

932 raise ValueError(

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

934 "logic. This function should not be called."

935 )

936

997 - def get_fused_moe_quant_config(

998 - self, layer: torch.nn.Module

999 - ) -> FusedMoEQuantConfig | None:

1000 - w1_scale = layer.w13_weight_scale

1001 - w2_scale = layer.w2_weight_scale

1002 - a1_scale = layer.w13_input_scale

1003 - a2_scale = layer.w2_input_scale

1004 -

937 + def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:

938 + is_per_token = self.input_quant.strategy == QuantizationStrategy.TOKEN

939 return make_fp8_moe_quant_config(

940 fp8_backend=self.fp8_backend,

1007 - w1_scale=w1_scale,

1008 - w2_scale=w2_scale,

1009 - a1_scale=a1_scale,

1010 - a2_scale=a2_scale,

1011 - per_act_token_quant=(

1012 - self.input_quant.strategy == QuantizationStrategy.TOKEN

1013 - ),

1014 - per_out_ch_quant=(self.input_quant.strategy == QuantizationStrategy.TOKEN),

941 + w1_scale=layer.w13_weight_scale,

942 + w2_scale=layer.w2_weight_scale,

943 + a1_scale=layer.w13_input_scale,

944 + a2_scale=layer.w2_input_scale,

945 + per_act_token_quant=is_per_token,

946 + per_out_ch_quant=is_per_token,

947 block_shape=self.weight_block_size,

948 )

949

1018 - @property

1019 - def is_monolithic(self) -> bool:

1020 - return self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM

1021 -

950 def apply_monolithic(

951 self,

952 layer: FusedMoE,

953 x: torch.Tensor,

954 router_logits: torch.Tensor,

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

1028 - assert self.is_monolithic

1029 - assert self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM

1030 - assert layer.activation == MoEActivation.SILU, (

1031 - f"Only SiLU activation is supported, not {layer.activation}."

956 + assert self.moe_kernel is not None

957 + return self.moe_kernel.apply_monolithic(

958 + x,

959 + layer.w13_weight,

960 + layer.w2_weight,

961 + router_logits,

962 + activation=layer.activation,

963 + global_num_experts=layer.global_num_experts,

964 + expert_map=layer.expert_map,

965 + apply_router_weight_on_input=layer.apply_router_weight_on_input,

966 + num_expert_group=layer.num_expert_group,

967 + topk_group=layer.topk_group,

968 + e_score_correction_bias=layer.e_score_correction_bias,

969 + routed_scaling_factor=layer.routed_scaling_factor,

970 )

971

1034 - if self.block_quant:

1035 - import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401

1036 -

1037 - return torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8(

1038 - routing_logits=router_logits,

1039 - routing_bias=layer.e_score_correction_bias,

1040 - x=x,

1041 - w13_weight=layer.w13_weight,

1042 - w13_weight_scale_inv=layer.w13_weight_scale,

1043 - w2_weight=layer.w2_weight,

1044 - w2_weight_scale_inv=layer.w2_weight_scale,

1045 - global_num_experts=layer.global_num_experts,

1046 - top_k=layer.top_k,

1047 - num_expert_group=layer.num_expert_group,

1048 - topk_group=layer.topk_group,

1049 - intermediate_size=layer.intermediate_size_per_partition,

1050 - expert_offset=layer.ep_rank * layer.local_num_experts,

1051 - local_num_experts=layer.local_num_experts,

1052 - block_shape=self.weight_block_size,

1053 - routing_method_type=layer.routing_method_type,

1054 - routed_scaling=layer.routed_scaling_factor,

1055 - )

1056 - else:

1057 - return apply_fi_trtllm_fp8_per_tensor_moe(

1058 - layer=layer,

1059 - hidden_states=x,

1060 - router_logits=router_logits,

1061 - routing_bias=layer.e_score_correction_bias,

1062 - global_num_experts=layer.global_num_experts,

1063 - top_k=layer.top_k,

1064 - num_expert_group=layer.num_expert_group,

1065 - topk_group=layer.topk_group,

1066 - apply_router_weight_on_input=layer.apply_router_weight_on_input,

1067 - )

1068 -

972 def apply(

973 self,

974 layer: FusedMoE,

@@ -1075,8 +978,8 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

978 shared_experts_input: torch.Tensor | None,

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

980 assert not self.is_monolithic

1078 - assert self.moe_mk is not None

1079 - return self.moe_mk(

981 + assert self.moe_kernel is not None

982 + return self.moe_kernel.apply(

983 x,

984 layer.w13_weight,

985 layer.w2_weight,

@@ -1652,9 +1555,9 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

1555

1556 def select_gemm_impl(

1557 self,

1655 - prepare_finalize: mk.FusedMoEPrepareAndFinalize,

1558 + prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular,

1559 layer: torch.nn.Module,

1657 - ) -> mk.FusedMoEPermuteExpertsUnpermute:

1560 + ) -> mk.FusedMoEExpertsModular:

1561 assert self.num_bits == 4, "only supporting w4"

1562 layer.w13_weight = layer.w13_weight_packed

1563 layer.w2_weight = layer.w2_weight_packed

@@ -1943,9 +1846,9 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):

1846

1847 def select_gemm_impl(

1848 self,

1946 - prepare_finalize: mk.FusedMoEPrepareAndFinalize,

1849 + prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular,

1850 layer: torch.nn.Module,

1948 - ) -> mk.FusedMoEPermuteExpertsUnpermute:

1851 + ) -> mk.FusedMoEExpertsModular:

1852 if self.moe.is_lora_enabled:

1853 assert self.moe_quant_config is not None

1854 from vllm.triton_utils import HAS_TRITON

@@ -2527,7 +2430,7 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):

2430 def maybe_make_prepare_finalize(

2431 self,

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

2530 - ) -> mk.FusedMoEPrepareAndFinalize | None:

2433 + ) -> mk.FusedMoEPrepareAndFinalizeModular | None:

2434 return super().maybe_make_prepare_finalize(routing_tables)

2435

2436 def get_fused_moe_quant_config(

@@ -2548,9 +2451,9 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):

2451

2452 def select_gemm_impl(

2453 self,

2551 - prepare_finalize: mk.FusedMoEPrepareAndFinalize,

2454 + prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular,

2455 layer: torch.nn.Module,

2553 - ) -> mk.FusedMoEPermuteExpertsUnpermute:

2456 + ) -> mk.FusedMoEExpertsModular:

2457 assert self.moe_quant_config is not None

2458 assert (

2459 prepare_finalize.activation_format == FusedMoEActivationFormat.Standard

@@ -2558,7 +2461,7 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):

2461

2462 from vllm.model_executor.layers.fused_moe import CutlassExpertsW4A8Fp8

2463

2561 - experts: FusedMoEPermuteExpertsUnpermute

2464 + experts: FusedMoEExpertsModular

2465

2466 logger.debug("CutlassExpertsW4A8Fp8(%s)", self.__class__.__name__)

2467 experts = CutlassExpertsW4A8Fp8(

vllm/model_executor/layers/quantization/fp8.py

+19 -78

@@ -23,17 +23,13 @@ from vllm.model_executor.layers.batch_invariant import (

23 from vllm.model_executor.layers.fused_moe import (

24 FusedMoE,

25 FusedMoEMethodBase,

26 - FusedMoEPermuteExpertsUnpermute,

27 - FusedMoEPrepareAndFinalize,

26 FusedMoeWeightScaleSupported,

29 - MoEActivation,

27 )

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

29 FusedMoEQuantConfig,

30 )

31 from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod

32 from vllm.model_executor.layers.fused_moe.oracle.fp8 import (

36 - Fp8MoeBackend,

33 convert_to_fp8_moe_kernel_format,

34 make_fp8_moe_kernel,

35 make_fp8_moe_quant_config,

@@ -50,9 +46,6 @@ from vllm.model_executor.layers.quantization.base_config import (

46 QuantizeMethodBase,

47 )

48 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,

55 -)

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

50 W8A8BlockFp8LinearOp,

51 create_fp8_input_scale,

@@ -860,14 +853,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):

853 replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_scale)

854 replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_scale)

855

863 - # Setup modular kernel for TP case and naive DP/EP case.

864 - # In non-naive DP/EP case, we will create a ModularKernelMethod.

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

866 - # in both cases.

856 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

857 if self.moe_quant_config:

858 assert self.experts_cls is not None

870 - self.moe_mk = make_fp8_moe_kernel(

859 + self.moe_kernel = make_fp8_moe_kernel(

860 moe_quant_config=self.moe_quant_config,

861 moe_config=self.moe,

862 fp8_backend=self.fp8_backend,

@@ -930,29 +919,13 @@ class Fp8MoEMethod(FusedMoEMethodBase):

919 def maybe_make_prepare_finalize(

920 self,

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

933 - ) -> mk.FusedMoEPrepareAndFinalize | None:

934 - raise ValueError(

935 - f"{self.__class__.__name__} uses the new modular kernel initialization "

936 - "logic. This function should not be called."

937 - )

938 -

939 - def select_gemm_impl(

940 - self,

941 - prepare_finalize: FusedMoEPrepareAndFinalize,

942 - layer: torch.nn.Module,

943 - ) -> FusedMoEPermuteExpertsUnpermute:

922 + ) -> mk.FusedMoEPrepareAndFinalizeModular | None:

923 raise ValueError(

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

925 "logic. This function should not be called."

926 )

927

949 - def get_fused_moe_quant_config(

950 - self, layer: torch.nn.Module

951 - ) -> FusedMoEQuantConfig | None:

952 - # TRTLLM does not use Modular Kernel.

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

954 - return None

955 -

928 + def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:

929 w1_scale = getattr(layer, f"w13_{self.weight_scale_name}")

930 w2_scale = getattr(layer, f"w2_{self.weight_scale_name}")

931 a1_scale = layer.w13_input_scale

@@ -983,10 +956,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):

956 def supports_eplb(self) -> bool:

957 return True

958

986 - @property

987 - def is_monolithic(self) -> bool:

988 - return self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM

989 -

959 def apply_monolithic(

960 self,

961 layer: FusedMoE,

@@ -994,50 +963,22 @@ class Fp8MoEMethod(FusedMoEMethodBase):

963 router_logits: torch.Tensor,

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

965 assert self.is_monolithic

997 - assert self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM

998 -

999 - # TODO(rob): convert this to MK.

1000 - if layer.enable_eplb:

1001 - raise NotImplementedError("EPLB not supported for `Fp8MoEMethod` yet.")

1002 - assert layer.activation == MoEActivation.SILU, (

1003 - f"Expected 'silu' activation but got {layer.activation}"

966 + assert self.moe_kernel is not None

967 + return self.moe_kernel.apply_monolithic(

968 + x,

969 + layer.w13_weight,

970 + layer.w2_weight,

971 + router_logits,

972 + activation=layer.activation,

973 + global_num_experts=layer.global_num_experts,

974 + expert_map=layer.expert_map,

975 + apply_router_weight_on_input=layer.apply_router_weight_on_input,

976 + num_expert_group=layer.num_expert_group,

977 + topk_group=layer.topk_group,

978 + e_score_correction_bias=layer.e_score_correction_bias,

979 + routed_scaling_factor=layer.routed_scaling_factor,

980 )

981

1006 - if self.block_quant:

1007 - import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401

1008 -

1009 - return torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8(

1010 - routing_logits=router_logits,

1011 - routing_bias=layer.e_score_correction_bias,

1012 - x=x,

1013 - w13_weight=layer.w13_weight,

1014 - w13_weight_scale_inv=layer.w13_weight_scale_inv,

1015 - w2_weight=layer.w2_weight,

1016 - w2_weight_scale_inv=layer.w2_weight_scale_inv,

1017 - global_num_experts=layer.global_num_experts,

1018 - top_k=layer.top_k,

1019 - num_expert_group=layer.num_expert_group,

1020 - topk_group=layer.topk_group,

1021 - intermediate_size=layer.intermediate_size_per_partition,

1022 - expert_offset=layer.ep_rank * layer.local_num_experts,

1023 - local_num_experts=layer.local_num_experts,

1024 - block_shape=self.weight_block_size,

1025 - routing_method_type=layer.routing_method_type,

1026 - routed_scaling=layer.routed_scaling_factor,

1027 - )

1028 - else:

1029 - return apply_fi_trtllm_fp8_per_tensor_moe(

1030 - layer=layer,

1031 - hidden_states=x,

1032 - router_logits=router_logits,

1033 - routing_bias=layer.e_score_correction_bias,

1034 - global_num_experts=layer.global_num_experts,

1035 - top_k=layer.top_k,

1036 - num_expert_group=layer.num_expert_group,

1037 - topk_group=layer.topk_group,

1038 - apply_router_weight_on_input=layer.apply_router_weight_on_input,

1039 - )

1040 -

982 def apply(

983 self,

984 layer: FusedMoE,

@@ -1046,9 +987,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):

987 topk_ids: torch.Tensor,

988 shared_experts_input: torch.Tensor | None,

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

1049 - assert self.moe_mk is not None

990 assert not self.is_monolithic

1051 - return self.moe_mk(

991 + assert self.moe_kernel is not None

992 + return self.moe_kernel.apply(

993 x,

994 layer.w13_weight,

995 layer.w2_weight,

vllm/model_executor/layers/quantization/modelopt.py

+63 -166

@@ -13,7 +13,6 @@ from vllm.model_executor.kernels.linear import (

13 init_fp8_linear_kernel,

14 )

15 from vllm.model_executor.layers.attention import Attention, MLAAttention

16 -from vllm.model_executor.layers.fused_moe.activation import MoEActivation

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

17 FusedMoEConfig,

18 FusedMoEQuantConfig,

@@ -24,14 +23,12 @@ from vllm.model_executor.layers.fused_moe.layer import (

23 FusedMoeWeightScaleSupported,

24 )

25 from vllm.model_executor.layers.fused_moe.oracle.fp8 import (

27 - Fp8MoeBackend,

26 convert_to_fp8_moe_kernel_format,

27 make_fp8_moe_kernel,

28 make_fp8_moe_quant_config,

29 select_fp8_moe_backend,

30 )

31 from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (

34 - NvFp4MoeBackend,

32 convert_to_nvfp4_moe_kernel_format,

33 is_global_sf_supported_for_nvfp4_backend,

34 make_nvfp4_moe_kernel,

@@ -49,13 +46,6 @@ from vllm.model_executor.layers.quantization.base_config import (

46 QuantizeMethodBase,

47 )

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

52 -from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (

53 - flashinfer_trtllm_fp4_moe,

54 - flashinfer_trtllm_fp4_routed_moe,

55 -)

56 -from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

57 - apply_fi_trtllm_fp8_per_tensor_moe,

58 -)

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

50 W8A8BlockFp8LinearOp,

51 process_fp8_input_tensor_strategy_moe,

@@ -746,7 +736,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

736 def maybe_make_prepare_finalize(

737 self,

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

749 - ) -> mk.FusedMoEPrepareAndFinalize | None:

739 + ) -> mk.FusedMoEPrepareAndFinalizeModular | None:

740 raise ValueError(

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

742 "logic. This function should not be called."

@@ -754,9 +744,9 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

744

745 def select_gemm_impl(

746 self,

757 - prepare_finalize: mk.FusedMoEPrepareAndFinalize,

747 + prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular,

748 layer: torch.nn.Module,

759 - ) -> mk.FusedMoEPermuteExpertsUnpermute:

749 + ) -> mk.FusedMoEExpertsModular:

750 raise ValueError(

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

752 "logic. This function should not be called."

@@ -871,16 +861,15 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

861

862 # Setup modular kernel.

863 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

874 - if self.moe_quant_config:

875 - assert self.experts_cls is not None

876 - self.moe_mk = make_fp8_moe_kernel(

877 - moe_quant_config=self.moe_quant_config,

878 - moe_config=self.moe,

879 - fp8_backend=self.fp8_backend,

880 - experts_cls=self.experts_cls,

881 - routing_tables=layer._maybe_init_expert_routing_tables(),

882 - shared_experts=layer.shared_experts,

883 - )

864 + assert self.experts_cls is not None

865 + self.moe_kernel = make_fp8_moe_kernel(

866 + moe_quant_config=self.moe_quant_config,

867 + moe_config=self.moe,

868 + fp8_backend=self.fp8_backend,

869 + experts_cls=self.experts_cls,

870 + routing_tables=layer._maybe_init_expert_routing_tables(),

871 + shared_experts=layer.shared_experts,

872 + )

873

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

875 w13 = layer.w13_weight

@@ -913,9 +902,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

902 layer, w13, w2, w13_scale, w2_scale, w13_input_scale, w2_input_scale

903 )

904

916 - def get_fused_moe_quant_config(

917 - self, layer: torch.nn.Module

918 - ) -> FusedMoEQuantConfig | None:

905 + def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:

906 w1_scale = layer.w13_weight_scale

907 w2_scale = layer.w2_weight_scale

908 a1_scale = layer.w13_input_scale

@@ -929,10 +916,6 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

916 a2_scale=a2_scale,

917 )

918

932 - @property

933 - def is_monolithic(self) -> bool:

934 - return self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM

935 -

919 def apply_monolithic(

920 self,

921 layer: FusedMoE,

@@ -940,28 +923,20 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

923 router_logits: torch.Tensor,

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

925 assert self.is_monolithic

943 - assert self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM

944 - if layer.enable_eplb:

945 - raise NotImplementedError(

946 - "EPLB not supported for FlashInfer TRTLLM FP8 MoE Backend."

947 - )

948 - # TODO(rob): this validation should happen at kernel selection

949 - # time in the oracle rather than here.

950 - SUPPORTED_ACTIVATIONS = [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]

951 - assert layer.activation in SUPPORTED_ACTIVATIONS, (

952 - f"Only {SUPPORTED_ACTIVATIONS} activations are supported for FlashInfer "

953 - f"TRTLLM FP4 MoE, {layer.activation} found instead."

954 - )

955 - return apply_fi_trtllm_fp8_per_tensor_moe(

956 - layer=layer,

957 - hidden_states=x,

958 - router_logits=router_logits,

959 - routing_bias=layer.e_score_correction_bias,

926 + assert self.moe_kernel is not None

927 + return self.moe_kernel.apply_monolithic(

928 + x,

929 + layer.w13_weight,

930 + layer.w2_weight,

931 + router_logits,

932 + activation=layer.activation,

933 global_num_experts=layer.global_num_experts,

961 - top_k=layer.top_k,

934 + expert_map=layer.expert_map,

935 + apply_router_weight_on_input=layer.apply_router_weight_on_input,

936 num_expert_group=layer.num_expert_group,

937 topk_group=layer.topk_group,

964 - apply_router_weight_on_input=layer.apply_router_weight_on_input,

938 + e_score_correction_bias=layer.e_score_correction_bias,

939 + routed_scaling_factor=layer.routed_scaling_factor,

940 )

941

942 def apply(

@@ -973,25 +948,13 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

948 shared_experts_input: torch.Tensor | None,

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

950 assert not self.is_monolithic

976 -

977 - # TODO(rob): this validation should happen at kernel selection

978 - # time in the oracle rather than here.

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

980 - assert layer.activation in (

981 - MoEActivation.SILU,

982 - MoEActivation.RELU2_NO_MUL,

983 - ), (

984 - "Expected activation to be in ('silu', 'relu2_no_mul'),"

985 - f"but got {layer.activation}"

986 - )

987 -

988 - assert self.moe_mk is not None

989 - return self.moe_mk(

990 - hidden_states=x,

991 - w1=layer.w13_weight,

992 - w2=layer.w2_weight,

993 - topk_weights=topk_weights,

994 - topk_ids=topk_ids,

951 + assert self.moe_kernel is not None

952 + return self.moe_kernel.apply(

953 + x,

954 + layer.w13_weight,

955 + layer.w2_weight,

956 + topk_weights,

957 + topk_ids,

958 activation=layer.activation,

959 global_num_experts=layer.global_num_experts,

960 expert_map=layer.expert_map,

@@ -1235,17 +1198,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1198 def maybe_make_prepare_finalize(

1199 self,

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

1238 - ) -> mk.FusedMoEPrepareAndFinalize | None:

1239 - raise ValueError(

1240 - f"{self.__class__.__name__} uses the new modular kernel initialization "

1241 - "logic. This function should not be called."

1242 - )

1243 -

1244 - def select_gemm_impl(

1245 - self,

1246 - prepare_finalize: mk.FusedMoEPrepareAndFinalize,

1247 - layer: torch.nn.Module,

1248 - ) -> mk.FusedMoEPermuteExpertsUnpermute:

1201 + ) -> mk.FusedMoEPrepareAndFinalizeModular | None:

1202 raise ValueError(

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

1204 "logic. This function should not be called."

@@ -1420,51 +1373,18 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1373 replace_parameter(layer, "w2_weight_scale_2", w2_scale_2)

1374 replace_parameter(layer, "w2_input_scale", a2_scale)

1375

1423 - # Setup modular kernel for TP case and naive DP/EP case.

1424 - # In non-naive DP/EP case, we will create a ModularKernelMethod.

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

1426 - # in both cases.

1376 + # Setup modular kernel.

1377 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

1428 - if self.moe_quant_config:

1429 - assert self.experts_cls is not None

1430 - self.moe_mk = make_nvfp4_moe_kernel(

1431 - moe_quant_config=self.moe_quant_config,

1432 - moe_config=self.moe,

1433 - experts_cls=self.experts_cls,

1434 - shared_experts=layer.shared_experts,

1435 - routing_tables=layer._maybe_init_expert_routing_tables(),

1436 - )

1437 -

1438 - @property

1439 - def do_post_quant_allgather(self):

1440 - return self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM

1441 -

1442 - def prepare_dp_allgather_tensor(

1443 - self,

1444 - layer: FusedMoE,

1445 - hidden_states: torch.Tensor,

1446 - router_logits: torch.Tensor,

1447 - ) -> tuple[torch.Tensor, list[torch.Tensor]]:

1448 - """Optionally prepare extra tensors to carry through DP allgather/EP."""

1449 - if self.nvfp4_backend != NvFp4MoeBackend.FLASHINFER_TRTLLM:

1450 - raise RuntimeError(

1451 - "prepare_dp_allgather_tensor is only supported for "

1452 - "FlashInfer TRTLLM NVFP4 MoE backend."

1453 - )

1454 -

1455 - import flashinfer

1456 -

1457 - hidden_states_fp4, hidden_states_sf = flashinfer.fp4_quantize(

1458 - hidden_states,

1459 - layer.a1_gscale,

1460 - is_sf_swizzled_layout=False,

1378 + assert self.experts_cls is not None

1379 + self.moe_kernel = make_nvfp4_moe_kernel(

1380 + moe_quant_config=self.moe_quant_config,

1381 + moe_config=self.moe,

1382 + experts_cls=self.experts_cls,

1383 + shared_experts=layer.shared_experts,

1384 + routing_tables=layer._maybe_init_expert_routing_tables(),

1385 )

1462 - extra_tensors: list[torch.Tensor] = [hidden_states_sf]

1463 - return hidden_states_fp4, extra_tensors

1386

1465 - def get_fused_moe_quant_config(

1466 - self, layer: torch.nn.Module

1467 - ) -> FusedMoEQuantConfig | None:

1387 + def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:

1388 return make_nvfp4_moe_quant_config(

1389 backend=self.nvfp4_backend,

1390 w13_scale=layer.w13_weight_scale,

@@ -1479,13 +1399,6 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1399 def supports_eplb(self) -> bool:

1400 return True

1401

1482 - @property

1483 - def is_monolithic(self) -> bool:

1484 - return (

1485 - self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM

1486 - and not self.moe.moe_parallel_config.enable_eplb

1487 - )

1488 -

1402 def apply_monolithic(

1403 self,

1404 layer: FusedMoE,

@@ -1493,22 +1406,20 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1406 router_logits: torch.Tensor,

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

1408 assert self.is_monolithic

1496 - assert (

1497 - self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM

1498 - and not layer.enable_eplb

1499 - )

1500 -

1501 - return flashinfer_trtllm_fp4_moe(

1502 - layer=layer,

1503 - x=x,

1504 - router_logits=router_logits,

1505 - top_k=layer.top_k,

1409 + assert self.moe_kernel is not None

1410 + return self.moe_kernel.apply_monolithic(

1411 + x,

1412 + layer.w13_weight,

1413 + layer.w2_weight,

1414 + router_logits,

1415 activation=layer.activation,

1416 global_num_experts=layer.global_num_experts,

1417 + expert_map=layer.expert_map,

1418 + apply_router_weight_on_input=layer.apply_router_weight_on_input,

1419 num_expert_group=layer.num_expert_group,

1420 topk_group=layer.topk_group,

1510 - custom_routing_function=layer.custom_routing_function,

1421 e_score_correction_bias=layer.e_score_correction_bias,

1422 + routed_scaling_factor=layer.routed_scaling_factor,

1423 )

1424

1425 def apply(

@@ -1520,33 +1431,19 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1431 shared_experts_input: torch.Tensor | None,

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

1433 assert not self.is_monolithic

1523 -

1524 - # EPLB path

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

1526 - assert layer.enable_eplb

1527 - return flashinfer_trtllm_fp4_routed_moe(

1528 - layer=layer,

1529 - x=x,

1530 - topk_ids=topk_ids,

1531 - topk_weights=topk_weights,

1532 - top_k=layer.top_k,

1533 - activation=layer.activation,

1534 - global_num_experts=layer.global_num_experts,

1535 - )

1536 - else:

1537 - assert self.moe_mk is not None

1538 - return self.moe_mk(

1539 - hidden_states=x,

1540 - w1=layer.w13_weight,

1541 - w2=layer.w2_weight,

1542 - topk_weights=topk_weights,

1543 - topk_ids=topk_ids,

1544 - activation=layer.activation,

1545 - global_num_experts=layer.global_num_experts,

1546 - expert_map=layer.expert_map,

1547 - apply_router_weight_on_input=layer.apply_router_weight_on_input,

1548 - shared_experts_input=shared_experts_input,

1549 - )

1434 + assert self.moe_kernel is not None

1435 + return self.moe_kernel.apply(

1436 + x,

1437 + layer.w13_weight,

1438 + layer.w2_weight,

1439 + topk_weights,

1440 + topk_ids,

1441 + activation=layer.activation,

1442 + global_num_experts=layer.global_num_experts,

1443 + expert_map=layer.expert_map,

1444 + apply_router_weight_on_input=layer.apply_router_weight_on_input,

1445 + shared_experts_input=shared_experts_input,

1446 + )

1447

1448

1449 ModelOptNvFp4Config.LinearMethodCls = ModelOptNvFp4LinearMethod

vllm/model_executor/layers/quantization/mxfp4.py

+7 -7

@@ -266,7 +266,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

266 )

267 self._cache_permute_indices: dict[torch.Size, torch.Tensor] = {}

268 # Initialized in process_weights_after_loading for CUTLASS/SM90 backends

269 - self.moe_mk: mk.FusedMoEModularKernel | None = None

269 + self.moe_kernel: mk.FusedMoEKernel | None = None

270

271 def create_weights(

272 self,

@@ -440,7 +440,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

440 )

441 assert prepare_finalize is not None

442

443 - self.moe_mk = mk.FusedMoEModularKernel(

443 + self.moe_kernel = mk.FusedMoEKernel(

444 prepare_finalize,

445 MarlinExperts(

446 self.moe,

@@ -789,7 +789,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

789 )

790 assert prepare_finalize is not None

791

792 - self.moe_mk = mk.FusedMoEModularKernel(

792 + self.moe_kernel = mk.FusedMoEKernel(

793 prepare_finalize,

794 FlashInferExperts(

795 moe_config=self.moe,

@@ -954,9 +954,9 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

954

955 def select_gemm_impl(

956 self,

957 - prepare_finalize: mk.FusedMoEPrepareAndFinalize,

957 + prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular,

958 layer: torch.nn.Module,

959 - ) -> mk.FusedMoEPermuteExpertsUnpermute:

959 + ) -> mk.FusedMoEExpertsModular:

960 if (

961 prepare_finalize.activation_format

962 == mk.FusedMoEActivationFormat.BatchedExperts

@@ -1043,8 +1043,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

1043 or self.mxfp4_backend == Mxfp4Backend.MARLIN

1044 )

1045

1046 - assert self.moe_mk is not None

1047 - return self.moe_mk(

1046 + assert self.moe_kernel is not None

1047 + return self.moe_kernel.apply(

1048 hidden_states=x,

1049 w1=layer.w13_weight,

1050 w2=layer.w2_weight,

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

+13 -283

@@ -6,28 +6,18 @@ from typing import TYPE_CHECKING

7 import torch

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

10 -from vllm import _custom_ops as ops

9 +import vllm.envs as envs

10 from vllm.logger import init_logger

12 -from vllm.model_executor.layers.fused_moe.activation import MoEActivation

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

14 - FusedMoEConfig,

15 - FusedMoEParallelConfig,

16 - RoutingMethodType,

17 -)

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

19 - activation_to_flashinfer_int,

12 align_fp4_moe_weights_for_fi,

13 )

14 from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (

15 swizzle_blockscale,

16 )

25 -from vllm.model_executor.layers.quantization.utils.quant_utils import (

26 - QuantKey,

27 - kNvfp4Dynamic,

28 - kNvfp4Static,

29 -)

17 from vllm.platforms import current_platform

18 +from vllm.utils.flashinfer import (

19 + has_flashinfer_cutlass_fused_moe,

20 +)

22 if TYPE_CHECKING:

23 from vllm.model_executor.layers.fused_moe.layer import FusedMoE

@@ -42,92 +32,15 @@ __all__ = [

32 "reorder_w1w3_to_w3w1",

33 ]

45 -#

46 -# Methods used by the oracle for kernel selection.

47 -#

48 -

49 -

50 -def _supports_current_device() -> bool:

51 - """Supports only Blackwell-family GPUs."""

52 - p = current_platform

53 - return p.is_cuda() and p.is_device_capability_family(100)

54 -

55 -

56 -def _supports_no_act_and_mul() -> bool:

57 - """Supports non-gated MoE."""

58 - return True

59 -

60 -

61 -def _supports_quant_scheme(

62 - weight_key: QuantKey | None,

63 - activation_key: QuantKey | None,

64 -) -> bool:

65 - """Supports Nvfp4 quantization."""

66 - SUPPORTED_W_A = [

67 - (kNvfp4Static, kNvfp4Dynamic),

68 - ]

69 - return (weight_key, activation_key) in SUPPORTED_W_A

70 -

71 -

72 -def _supports_activation(activation: MoEActivation) -> bool:

73 - return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]

74 -

75 -

76 -def _supports_routing_method(

77 - routing_method: RoutingMethodType,

78 -) -> bool:

79 - """Monolithic kernels need to express router support."""

80 - # NOTE(rob): potentially allow others here. This is a conservative list.

81 - return routing_method in [

82 - RoutingMethodType.DeepSeekV3,

83 - RoutingMethodType.Renormalize,

84 - RoutingMethodType.RenormalizeNaive,

85 - RoutingMethodType.Llama4,

86 - ]

87 -

88 -

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

90 - """

91 - TRTLLM is a monolithic kernel that requires dispatch_router_logits() for

92 - the naive dispatch/combine path. DeepEP HT only implements dispatch() for

93 - the modular kernel path, so TRTLLM is incompatible with DeepEP HT.

94 - """

95 - return not moe_parallel_config.use_deepep_ht_kernels

96 -

97 -

98 -def is_supported_config_trtllm(

99 - moe_config: FusedMoEConfig,

100 - weight_key: QuantKey | None,

101 - activation_key: QuantKey | None,

102 - activation_format: mk.FusedMoEActivationFormat,

103 -) -> tuple[bool, str | None]:

104 - """

105 - This method mirrors mk.FusedMoEPermuteExpertsUnpermute.is_supported_config

106 - """

107 -

108 - def _make_reason(reason: str) -> str:

109 - return f"kernel does not support {reason}"

110 -

111 - if not _supports_current_device():

112 - return False, _make_reason(f"current device {current_platform.device_name}")

113 - elif not (moe_config.is_act_and_mul or _supports_no_act_and_mul()):

114 - return False, _make_reason("no act_and_mul MLP layer")

115 - elif not _supports_activation(moe_config.activation):

116 - return False, _make_reason(f"{moe_config.activation} activation")

117 - elif not _supports_quant_scheme(weight_key, activation_key):

118 - return False, _make_reason(f"quantization scheme {weight_key}x{activation_key}")

119 - elif not _supports_parallel_config(moe_config.moe_parallel_config):

120 - return False, _make_reason(f"parallel config {moe_config.moe_parallel_config}")

121 - elif not _supports_routing_method(moe_config.routing_method):

122 - return False, _make_reason(f"routing method {moe_config.routing_method}")

123 - elif activation_format != mk.FusedMoEActivationFormat.Standard:

124 - return False, _make_reason(f"activation format {activation_format}")

125 - elif moe_config.hidden_dim % 512 != 0:

126 - return False, _make_reason(

127 - f"hidden_dim must be divisible by 512, found {moe_config.hidden_dim}"

128 - )

130 - return True, None

36 +def is_flashinfer_fp4_cutlass_moe_available() -> bool:

37 + """Return `True` when FlashInfer CUTLASS NV-FP4 kernels can be used."""

38 + return (

39 + envs.VLLM_USE_FLASHINFER_MOE_FP4

40 + and has_flashinfer_cutlass_fused_moe()

41 + and current_platform.is_cuda()

42 + and current_platform.has_device_capability(100)

43 + )

46 def reorder_w1w3_to_w3w1(

@@ -276,190 +189,6 @@ def prepare_static_weights_for_trtllm_fp4_moe(

189 )

190

191

279 -def flashinfer_trtllm_fp4_moe(

280 - layer: torch.nn.Module,

281 - x: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

282 - router_logits: torch.Tensor,

283 - top_k: int,

284 - activation: MoEActivation,

285 - global_num_experts: int,

286 - num_expert_group: int | None,

287 - topk_group: int | None,

288 - custom_routing_function: object | None,

289 - e_score_correction_bias: torch.Tensor | None,

290 -) -> torch.Tensor:

291 - """

292 - Apply FlashInfer TensorRT-LLM FP4 MoE kernel.

293 -

294 - Args:

295 - layer: The MoE layer with weights and scales

296 - x: Input tensor

297 - router_logits: Router logits for expert selection

298 - top_k: Number of experts to select per token

299 - activation: Activation function to use

300 - global_num_experts: Total number of experts across all ranks

301 - num_expert_group: Number of expert groups (for grouped routing)

302 - topk_group: Top-k within each group

303 - custom_routing_function: Custom routing function (e.g., Llama4)

304 - e_score_correction_bias: Optional routing bias correction

305 -

306 - Returns:

307 - Output tensor from the MoE layer

308 - """

309 - import flashinfer

310 -

311 - from vllm.model_executor.models.llama4 import Llama4MoE

312 -

313 - SUPPORTED_ACTIVATIONS = [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]

314 - assert activation in SUPPORTED_ACTIVATIONS, (

315 - f"Only {SUPPORTED_ACTIVATIONS} activations are supported for FlashInfer "

316 - f"TRTLLM FP4 MoE, {activation} found instead."

317 - )

318 -

319 - # Quantize input to FP4

320 - if isinstance(x, tuple):

321 - hidden_states_fp4, hidden_states_scale_linear_fp4 = x

322 - else:

323 - # hidden_states is the already quantized

324 - (hidden_states_fp4, hidden_states_scale_linear_fp4) = ops.scaled_fp4_quant(

325 - x, layer.a1_gscale, is_sf_swizzled_layout=False

326 - )

327 -

328 - # Determine routing method type

329 - use_llama4_routing = custom_routing_function is Llama4MoE.custom_routing_function

330 - routing_method_type = layer.routing_method_type

331 - if use_llama4_routing:

332 - routing_method_type = flashinfer.RoutingMethodType.Llama4

333 -

334 - # Cast to Fp32 (required by kernel).

335 - router_logits = (

336 - router_logits.to(torch.float32)

337 - if routing_method_type == RoutingMethodType.DeepSeekV3

338 - else router_logits

339 - )

340 -

341 - # Determine activation type

342 - activation_type = activation_to_flashinfer_int(layer.activation)

343 -

344 - # Call TRT-LLM FP4 block-scale MoE kernel

345 - out = flashinfer.fused_moe.trtllm_fp4_block_scale_moe(

346 - routing_logits=router_logits,

347 - routing_bias=e_score_correction_bias,

348 - hidden_states=hidden_states_fp4,

349 - hidden_states_scale=hidden_states_scale_linear_fp4.view(

350 - torch.float8_e4m3fn

351 - ).reshape(*hidden_states_fp4.shape[:-1], -1),

352 - gemm1_weights=layer.w13_weight.data,

353 - gemm1_weights_scale=layer.w13_weight_scale.data.view(torch.float8_e4m3fn),

354 - gemm1_bias=None,

355 - gemm1_alpha=None,

356 - gemm1_beta=None,

357 - gemm1_clamp_limit=None,

358 - gemm2_weights=layer.w2_weight.data,

359 - gemm2_weights_scale=layer.w2_weight_scale.data.view(torch.float8_e4m3fn),

360 - gemm2_bias=None,

361 - output1_scale_scalar=layer.g1_scale_c.data,

362 - output1_scale_gate_scalar=layer.g1_alphas.data,

363 - output2_scale_scalar=layer.g2_alphas.data,

364 - num_experts=global_num_experts,

365 - top_k=top_k,

366 - n_group=num_expert_group if num_expert_group is not None else 0,

367 - topk_group=topk_group if topk_group is not None else 0,

368 - intermediate_size=layer.intermediate_size_per_partition,

369 - local_expert_offset=layer.ep_rank * layer.local_num_experts,

370 - local_num_experts=layer.local_num_experts,

371 - routed_scaling_factor=None,

372 - routing_method_type=routing_method_type,

373 - do_finalize=True,

374 - activation_type=activation_type,

375 - )[0]

376 -

377 - return out

378 -

379 -

380 -def flashinfer_trtllm_fp4_routed_moe(

381 - layer: torch.nn.Module,

382 - x: torch.Tensor,

383 - topk_ids: torch.Tensor,

384 - topk_weights: torch.Tensor,

385 - top_k: int,

386 - activation: MoEActivation,

387 - global_num_experts: int,

388 -) -> torch.Tensor:

389 - """

390 - Apply FlashInfer TensorRT-LLM FP4 MoE kernel. Uses packed

391 - input top k expert indices and scores rather than computing

392 - top k expert indices from scores.

393 -

394 - Args:

395 - layer: The MoE layer with weights and scales

396 - x: Input tensor

397 - topk_ids: Ids of selected experts

398 - top_k: Number of experts to select per token

399 - activation: Activation function to use

400 - global_num_experts: Total number of experts across all ranks

401 -

402 - Returns:

403 - Output tensor from the MoE layer

404 - """

405 - import flashinfer

406 -

407 - # https://github.com/flashinfer-ai/flashinfer/blob/f0277fd1bff90e309e5c19cab36c5dae056d685d/flashinfer/fused_moe/core.py#L2535

408 - assert activation == MoEActivation.SILU, (

409 - "Only SiLU activation is supported for FlashInfer TRTLLM FP4 Routed MoE. "

410 - f"{activation} found instead."

411 - )

412 -

413 - # Pack top k ids and expert weights into a single int32 tensor, as

414 - # required by TRT-LLM

415 - packed_tensor = (topk_ids.to(torch.int32) << 16) | topk_weights.to(

416 - torch.bfloat16

417 - ).view(torch.int16)

418 -

419 - if isinstance(x, tuple):

420 - # Hidden_states is the already quantized

421 - hidden_states_fp4, hidden_states_scale_linear_fp4 = x

422 - else:

423 - # Quantize input to FP4

424 - (hidden_states_fp4, hidden_states_scale_linear_fp4) = ops.scaled_fp4_quant(

425 - x, layer.a1_gscale, is_sf_swizzled_layout=False

426 - )

427 -

428 - # Call TRT-LLM FP4 block-scale MoE kernel

429 - out = flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe(

430 - topk_ids=packed_tensor,

431 - routing_bias=None,

432 - hidden_states=hidden_states_fp4,

433 - hidden_states_scale=hidden_states_scale_linear_fp4.view(

434 - torch.float8_e4m3fn

435 - ).reshape(*hidden_states_fp4.shape[:-1], -1),

436 - gemm1_weights=layer.w13_weight.data,

437 - gemm1_weights_scale=layer.w13_weight_scale.data.view(torch.float8_e4m3fn),

438 - gemm1_bias=None,

439 - gemm1_alpha=None,

440 - gemm1_beta=None,

441 - gemm1_clamp_limit=None,

442 - gemm2_weights=layer.w2_weight.data,

443 - gemm2_weights_scale=layer.w2_weight_scale.data.view(torch.float8_e4m3fn),

444 - gemm2_bias=None,

445 - output1_scale_scalar=layer.g1_scale_c.data,

446 - output1_scale_gate_scalar=layer.g1_alphas.data,

447 - output2_scale_scalar=layer.g2_alphas.data,

448 - num_experts=global_num_experts,

449 - top_k=top_k,

450 - n_group=0,

451 - topk_group=0,

452 - intermediate_size=layer.intermediate_size_per_partition,

453 - local_expert_offset=layer.ep_rank * layer.local_num_experts,

454 - local_num_experts=layer.local_num_experts,

455 - routed_scaling_factor=None,

456 - routing_method_type=1,

457 - do_finalize=True,

458 - )[0]

459 -

460 - return out

461 -

462 -

192 def prepare_nvfp4_moe_layer_for_fi_or_cutlass(

193 backend: "NvFp4MoeBackend",

194 layer: "FusedMoE",

@@ -526,6 +255,7 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass(

255 )

256 )

257 layer.intermediate_size_per_partition = padded_intermediate

258 + layer.moe_config.intermediate_size_per_partition = padded_intermediate

259

260 w13, w13_scale, w2, w2_scale = prepare_static_weights_for_trtllm_fp4_moe(

261 w13,

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

+11 -108

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 from enum import Enum

4 +from typing import TYPE_CHECKING

6 import torch

@@ -10,6 +11,9 @@ from vllm.model_executor.layers.fused_moe.activation import MoEActivation

11 from vllm.platforms import current_platform

12 from vllm.utils.math_utils import round_up

14 +if TYPE_CHECKING:

15 + from flashinfer.fused_moe.core import ActivationType

16 +

17 logger = init_logger(__name__)

@@ -20,6 +24,10 @@ class FlashinferMoeBackend(Enum):

26 def activation_to_flashinfer_int(activation: MoEActivation) -> int:

27 + return activation_to_flashinfer_type(activation).value

28 +

29 +

30 +def activation_to_flashinfer_type(activation: MoEActivation) -> "ActivationType":

31 from flashinfer.fused_moe.core import ActivationType

33 # silu and gelu are mapped to their gated versions SwiGLU and GeGLU respectively

@@ -30,7 +38,7 @@ def activation_to_flashinfer_int(activation: MoEActivation) -> int:

38 MoEActivation.GELU: ActivationType.Geglu,

39 MoEActivation.RELU2_NO_MUL: ActivationType.Relu2,

40 }

33 - return ACTIVATION_TO_FI_ACTIVATION[activation].value

41 + return ACTIVATION_TO_FI_ACTIVATION[activation]

44 def swap_w13_to_w31(x: torch.Tensor) -> torch.Tensor:

@@ -87,104 +95,6 @@ def rotate_weights_for_fi_trtllm_fp8_per_tensor_moe(

95 )

90 -def register_scales_for_trtllm_fp8_per_tensor_moe(

91 - layer: torch.nn.Module,

92 - w13_scale: torch.Tensor,

93 - w13_input_scale: torch.Tensor,

94 - w2_scale: torch.Tensor,

95 - w2_input_scale: torch.Tensor,

96 -) -> None:

97 - """Register necessary scales for FlashInfer TRTLLM FP8 MoE kernel"""

98 - g1_alphas, g2_alphas = make_fp8_moe_alpha_scales_for_fi(

99 - w13_scale=w13_scale,

100 - w13_input_scale=w13_input_scale,

101 - w2_scale=w2_scale,

102 - w2_input_scale=w2_input_scale,

103 - )

104 - layer.w2_input_scale_inv = 1.0 / w2_input_scale

105 - layer.output1_scales_gate_scalar = g1_alphas

106 -

107 - if layer.activation.is_gated:

108 - layer.output1_scales_scalar = g1_alphas * layer.w2_input_scale_inv

109 - else:

110 - layer.output1_scales_scalar = (

111 - torch.ones_like(g1_alphas) * layer.w2_input_scale_inv

112 - )

113 - layer.output2_scales_scalar = g2_alphas

114 -

115 -

116 -def apply_fi_trtllm_fp8_per_tensor_moe(

117 - layer: torch.nn.Module,

118 - hidden_states: torch.Tensor,

119 - router_logits: torch.Tensor,

120 - routing_bias: torch.Tensor | None,

121 - top_k: int,

122 - num_expert_group: int | None,

123 - topk_group: int | None,

124 - global_num_experts: int,

125 - apply_router_weight_on_input: bool,

126 -) -> torch.Tensor:

127 - from flashinfer.fused_moe import RoutingMethodType

128 -

129 - import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401

130 - from vllm.model_executor.models.llama4 import Llama4MoE

131 -

132 - # Added to the layer by: register_scales_for_trtllm_fp8_per_tensor_moe

133 - assert (

134 - hasattr(layer, "output1_scales_scalar")

135 - and hasattr(layer, "output1_scales_gate_scalar")

136 - and hasattr(layer, "output2_scales_scalar")

137 - )

138 -

139 - if layer.routing_method_type == RoutingMethodType.Llama4:

140 - assert (

141 - not layer.renormalize

142 - and layer.custom_routing_function == Llama4MoE.custom_routing_function

143 - ), (

144 - "FusedMoE flashinfer kernels with Llama4 routing method are only "

145 - "supported for Llama4"

146 - )

147 - else:

148 - assert layer.custom_routing_function is None, (

149 - "Custom routing function is only supported for Llama4"

150 - )

151 - activation_type = activation_to_flashinfer_int(layer.activation)

152 -

153 - return torch.ops.vllm.fi_trtllm_fp8_per_tensor_moe(

154 - routing_logits=router_logits,

155 - routing_bias=routing_bias,

156 - hidden_states=hidden_states,

157 - input_scale=layer.w13_input_scale,

158 - gemm1_weights=layer.w13_weight,

159 - gemm2_weights=layer.w2_weight,

160 - output1_scales_scalar=layer.output1_scales_scalar,

161 - output1_scales_gate_scalar=layer.output1_scales_gate_scalar,

162 - output2_scales_scalar=layer.output2_scales_scalar,

163 - num_experts=global_num_experts,

164 - top_k=top_k,

165 - num_expert_group=num_expert_group,

166 - topk_group=topk_group,

167 - intermediate_size=layer.intermediate_size_per_partition,

168 - local_expert_offset=layer.ep_rank * layer.local_num_experts,

169 - local_num_experts=layer.local_num_experts,

170 - use_routing_scales_on_input=apply_router_weight_on_input,

171 - routing_method_type=layer.routing_method_type,

172 - activation_type=activation_type,

173 - )

174 -

175 -

176 -def make_fp8_moe_alpha_scales_for_fi(

177 - w13_scale: torch.Tensor,

178 - w13_input_scale: torch.Tensor,

179 - w2_scale: torch.Tensor,

180 - w2_input_scale: torch.Tensor,

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

182 - g1_alphas = (w13_scale * w13_input_scale).squeeze()

183 - g2_alphas = (w2_scale * w2_input_scale).squeeze()

184 -

185 - return g1_alphas, g2_alphas

186 -

187 -

98 def get_flashinfer_moe_backend() -> FlashinferMoeBackend:

99 backend_map = {

100 "throughput": FlashinferMoeBackend.CUTLASS,

@@ -432,6 +342,7 @@ def prepare_fp8_moe_layer_for_fi(

342 min_alignment,

343 )

344 layer.intermediate_size_per_partition = new_intermediate

345 + layer.moe_config.intermediate_size_per_partition = new_intermediate

346

347 # FI kernels require W31 layout rather than W13.

348 if layer.moe_config.is_act_and_mul:

@@ -440,20 +351,12 @@ def prepare_fp8_moe_layer_for_fi(

351 w13_scale = swap_w13_to_w31(w13_scale)

352

353 # FI TRT-LLM FP8 per-tensor MoE kernel requires weight shuffle

443 - # and registration of alpha scales. Note that we do not register

444 - # as nn.Parameters since they are not needed for weight-reloading.

354 + # and registration of alpha scales.

355 if is_trtllm and not block_quant:

356 assert w13_input_scale is not None

357 assert w2_input_scale is not None

358

359 rotate_weights_for_fi_trtllm_fp8_per_tensor_moe(w13, w2, is_gated)

450 - register_scales_for_trtllm_fp8_per_tensor_moe(

451 - layer,

452 - w13_scale=w13_scale,

453 - w13_input_scale=w13_input_scale,

454 - w2_scale=w2_scale,

455 - w2_input_scale=w2_input_scale,

456 - )

360

361 # Clamp block scales to avoid NaN from the FlashInfer CUTLASS kernel.

362 # Some FP8 models have near-zero block scales (~1e-23) for dead/unused

vllm/model_executor/warmup/deep_gemm_warmup.py

+1 -1

@@ -172,7 +172,7 @@ def _fused_moe_grouped_gemm_may_use_deep_gemm(module: torch.nn.Module) -> bool:

172

173 # Further check if the ModularKernel implementation uses the DeepGemmExperts

174 return isinstance(

175 - module.quant_method.moe_mk, (DeepGemmExperts, TritonOrDeepGemmExperts)

175 + module.quant_method.moe_kernel, (DeepGemmExperts, TritonOrDeepGemmExperts)

176 )

177

178

vllm/model_executor/warmup/kernel_warmup.py

+9 -2

@@ -88,9 +88,14 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:

88 Without autotuning, FlashInfer will rely on heuristics, which may

89 be significantly slower.

90 """

91 - from vllm.utils.flashinfer import autotune

91 + import vllm.utils.flashinfer as fi_utils

92 +

93 + with torch.inference_mode(), fi_utils.autotune():

94 + # Certain FlashInfer kernels (e.g. nvfp4 routed moe) are

95 + # incompatible with autotuning. This state is used to skip

96 + # those kernels during the autotuning process.

97 + fi_utils._is_fi_autotuning = True

93 - with torch.inference_mode(), autotune():

99 # We skip EPLB here since we don't want to record dummy metrics

100 # When autotuning with number of tokens m, flashinfer will autotune

101 # operations for all number of tokens up to m.

@@ -100,3 +105,5 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:

105 skip_eplb=True,

106 is_profile=True,

107 )

108 +

109 + fi_utils._is_fi_autotuning = False

vllm/utils/flashinfer.py

@@ -140,6 +140,7 @@ autotune = _lazy_import_wrapper(

140 "autotune",

141 fallback_fn=lambda *args, **kwargs: contextlib.nullcontext(),

142 )

143 +_is_fi_autotuning: bool = False

144

145

146 @functools.cache