@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9f6dcb71ae

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 10, 12:45 UTC

[MoE Refactor][16/N] Apply Refactor to NVFP4 (#31692)

Signed-off-by: Robert Shaw <robshaw@redhat.com> Signed-off-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Pavani Majety <pmajety@nvidia.com>

Robert Shaw committed Jan 7, 2026 at 22:46 UTC 9f6dcb71ae4fd410d144fa3d4bb135ea9ab74185

15 files changed +780 -684

benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py

+34 -19

@@ -11,14 +11,20 @@ import nvtx

11 import torch

12 import torch.utils.benchmark as benchmark

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

15 from vllm import _custom_ops as ops

16 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config

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

18 fp8_w8a8_moe_quant_config,

19 nvfp4_moe_quant_config,

20 )

20 -from vllm.model_executor.layers.fused_moe.cutlass_moe import cutlass_moe_fp4

21 +from vllm.model_executor.layers.fused_moe.cutlass_moe import (

22 + CutlassExpertsFp4,

23 +)

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

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

26 + MoEPrepareAndFinalizeNoEP,

27 +)

28 from vllm.scalar_type import scalar_types

29 from vllm.utils.argparse_utils import FlexibleArgumentParser

30 from vllm.v1.worker.workspace import init_workspace_manager

@@ -188,19 +194,24 @@ def bench_run(

194 g1_alphas=w1_gs,

195 g2_alphas=w2_gs,

196 )

197 +

198 + kernel = mk.FusedMoEModularKernel(

199 + MoEPrepareAndFinalizeNoEP(defer_input_quant=True),

200 + CutlassExpertsFp4(

201 + out_dtype=dtype,

202 + max_experts_per_worker=e,

203 + quant_config=quant_config,

204 + ),

205 + )

206 +

207 for _ in range(num_repeats):

208 with nvtx.annotate("cutlass_moe_fp4", color="green"):

193 - cutlass_moe_fp4(

194 - a=a,

195 - w1_fp4=w1_fp4,

196 - w2_fp4=w2_fp4,

209 + kernel(

210 + hidden_states=a,

211 + w1=w1_fp4,

212 + w2=w2_fp4,

213 topk_weights=topk_weights,

214 topk_ids=topk_ids,

199 - m=m,

200 - n=n,

201 - k=k,

202 - e=num_experts,

203 - quant_config=quant_config,

215 )

216

217 def run_cutlass_from_graph(

@@ -230,20 +241,24 @@ def bench_run(

241 g2_alphas=w2_gs,

242 )

243

244 + kernel = mk.FusedMoEModularKernel(

245 + MoEPrepareAndFinalizeNoEP(defer_input_quant=True),

246 + CutlassExpertsFp4(

247 + out_dtype=dtype,

248 + max_experts_per_worker=e,

249 + quant_config=quant_config,

250 + ),

251 + )

252 +

253 with set_current_vllm_config(

254 VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))

255 ):

236 - return cutlass_moe_fp4(

237 - a=a,

238 - w1_fp4=w1_fp4,

239 - w2_fp4=w2_fp4,

256 + return kernel(

257 + hidden_states=a,

258 + w1=w1_fp4,

259 + w2=w2_fp4,

260 topk_weights=topk_weights,

261 topk_ids=topk_ids,

242 - m=m,

243 - n=n,

244 - k=k,

245 - e=num_experts,

246 - quant_config=quant_config,

262 )

263

264 def run_triton_from_graph(

docs/design/moe_kernel_features.md

+1 -1

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

86 | triton | standard | all1 | G,A,T | silu, gelu,swigluoai,silu_no_mul,gelu_no_mul | Y | Y | [`fused_experts`][vllm.model_executor.layers.fused_moe.fused_moe.fused_experts],[`TritonExperts`][vllm.model_executor.layers.fused_moe.fused_moe.TritonExperts] |

87 | triton (batched) | batched | all1 | G,A,T | silu, gelu | 6 | Y | [`BatchedTritonExperts`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedTritonExperts] |

88 | deep gemm | standard,batched | fp8 | G(128),A,T | silu, gelu | 6 | Y | [`deep_gemm_moe_fp8`][vllm.model_executor.layers.fused_moe.deep_gemm_moe.deep_gemm_moe_fp8],[`DeepGemmExperts`][vllm.model_executor.layers.fused_moe.deep_gemm_moe.DeepGemmExperts],[`BatchedDeepGemmExperts`][vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe.BatchedDeepGemmExperts] |

89 -| cutlass_fp4 | standard,batched | nvfp4 | A,T | silu | Y | Y | [`cutlass_moe_fp4`][vllm.model_executor.layers.fused_moe.cutlass_moe.cutlass_moe_fp4],[`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp4] |

89 +| cutlass_fp4 | standard,batched | nvfp4 | A,T | silu | Y | Y | [`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp4] |

90 | cutlass_fp8 | standard,batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp8],[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassBatchedExpertsFp8] |

91 | flashinfer | standard | nvfp4,fp8 | T | 5 | N | Y | [`flashinfer_cutlass_moe_fp4`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.flashinfer_cutlass_moe_fp4],[`FlashInferExperts`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts] |

92 | gpt oss triton | standard | N/A | N/A | 5 | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.OAITritonExperts] |

tests/kernels/moe/test_nvfp4_moe.py

+20 -10

@@ -3,6 +3,7 @@

3 import pytest

4 import torch

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

7 from tests.kernels.moe.utils import make_test_weights

8 from tests.kernels.quantization.nvfp4_utils import (

9 FLOAT4_E2M1_MAX,

@@ -13,8 +14,13 @@ from tests.kernels.utils import torch_moe

14 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.config import nvfp4_moe_quant_config

16 -from vllm.model_executor.layers.fused_moe.cutlass_moe import cutlass_moe_fp4

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

18 + CutlassExpertsFp4,

19 +)

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

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

22 + MoEPrepareAndFinalizeNoEP,

23 +)

24 from vllm.platforms import current_platform

25 from vllm.utils.torch_utils import set_random_seed

@@ -83,17 +89,21 @@ def test_cutlass_fp4_moe_no_graph(

89 w2_scale=w2_blockscale,

90 )

86 - cutlass_output = cutlass_moe_fp4(

87 - a=a,

88 - w1_fp4=w1_q,

89 - w2_fp4=w2_q,

92 + kernel = mk.FusedMoEModularKernel(

93 + MoEPrepareAndFinalizeNoEP(defer_input_quant=True),

94 + CutlassExpertsFp4(

95 + out_dtype=dtype,

96 + max_experts_per_worker=e,

97 + quant_config=quant_config,

98 + ),

99 + )

100 +

101 + cutlass_output = kernel(

102 + hidden_states=a,

103 + w1=w1_q,

104 + w2=w2_q,

105 topk_weights=topk_weights,

106 topk_ids=topk_ids,

92 - quant_config=quant_config,

93 - m=m,

94 - n=n,

95 - k=k,

96 - e=e,

107 )

108

109 # Reference check:

vllm/model_executor/layers/fused_moe/__init__.py

-2

@@ -72,7 +72,6 @@ if HAS_TRITON:

72 CutlassBatchedExpertsFp8,

73 CutlassExpertsFp8,

74 CutlassExpertsW4A8Fp8,

75 - cutlass_moe_fp4,

75 cutlass_moe_w4a8_fp8,

76 )

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

@@ -95,7 +94,6 @@ if HAS_TRITON:

94 "fused_experts",

95 "get_config_file_name",

96 "GroupedTopk",

98 - "cutlass_moe_fp4",

97 "cutlass_moe_w4a8_fp8",

98 "CutlassExpertsFp8",

99 "CutlassBatchedExpertsFp8",

vllm/model_executor/layers/fused_moe/config.py

+23

@@ -336,6 +336,10 @@ class FusedMoEQuantConfig:

336 def use_int4_w4a16(self) -> bool:

337 return self._a1.dtype is None and self._w1.dtype == "int4"

338

339 + @property

340 + def use_nvfp4_w4a16(self) -> bool:

341 + return self._a1.dtype is None and self._w1.dtype == "nvfp4"

342 +

343 @property

344 def ocp_mx_scheme(self) -> str | None:

345 if not hasattr(self, "_ocp_mx_scheme"):

@@ -690,6 +694,25 @@ def nvfp4_moe_quant_config(

694 )

695

696

697 +def nvfp4_w4a16_moe_quant_config(

698 + g1_alphas: torch.Tensor,

699 + g2_alphas: torch.Tensor,

700 + w1_scale: torch.Tensor,

701 + w2_scale: torch.Tensor,

702 +) -> FusedMoEQuantConfig:

703 + """

704 + Construct a quant config for 16-but activations and nvp4 weights.

705 + """

706 + return FusedMoEQuantConfig.make(

707 + quant_dtype=None,

708 + w1_scale=w1_scale,

709 + w2_scale=w2_scale,

710 + g1_alphas=g1_alphas,

711 + g2_alphas=g2_alphas,

712 + weight_dtype="nvfp4",

713 + )

714 +

715 +

716 def int4_w4a16_moe_quant_config(

717 w1_scale: torch.Tensor,

718 w2_scale: torch.Tensor,

vllm/model_executor/layers/fused_moe/cutlass_moe.py

-62

@@ -706,68 +706,6 @@ class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):

706 )

707

708

709 -def cutlass_moe_fp4(

710 - a: torch.Tensor,

711 - w1_fp4: torch.Tensor,

712 - w2_fp4: torch.Tensor,

713 - topk_weights: torch.Tensor,

714 - topk_ids: torch.Tensor,

715 - quant_config: FusedMoEQuantConfig,

716 - m: int,

717 - n: int,

718 - k: int,

719 - e: int,

720 - expert_map: torch.Tensor | None = None,

721 - apply_router_weight_on_input: bool = False,

722 -) -> torch.Tensor:

723 - assert expert_map is None, (

724 - "Expert Parallelism / expert_map "

725 - "is currently not supported for "

726 - "ModelOptNvFp4FusedMoE's cutlass_moe_fp4."

727 - )

728 -

729 - # TODO(bnell): this feels a bit hacky

730 - # NVFP4 requires two levels of quantization, which involves

731 - # computing some scaling factors dynamically. This makes it

732 - # incompatible with the typical prepare -> MoE -> finalize

733 - # pipeline. Move the quantization logic into the MoE body.

734 - quant_config = FusedMoEQuantConfig.make(

735 - quant_dtype=None, # skip quantization in prepare/finalize

736 - per_act_token_quant=quant_config.per_act_token_quant,

737 - per_out_ch_quant=quant_config.per_out_ch_quant,

738 - block_shape=quant_config.block_shape,

739 - g1_alphas=quant_config.g1_alphas,

740 - g2_alphas=quant_config.g2_alphas,

741 - a1_gscale=quant_config.a1_gscale,

742 - a2_gscale=quant_config.a2_gscale,

743 - w1_scale=quant_config.w1_scale,

744 - w2_scale=quant_config.w2_scale,

745 - )

746 -

747 - fn = mk.FusedMoEModularKernel(

748 - MoEPrepareAndFinalizeNoEP(),

749 - CutlassExpertsFp4(

750 - max_experts_per_worker=e,

751 - out_dtype=a.dtype,

752 - quant_config=quant_config,

753 - use_batched_format=False,

754 - ),

755 - )

756 -

757 - return fn(

758 - hidden_states=a,

759 - w1=w1_fp4,

760 - w2=w2_fp4,

761 - topk_weights=topk_weights,

762 - topk_ids=topk_ids,

763 - inplace=False,

764 - activation="silu",

765 - global_num_experts=e,

766 - expert_map=None,

767 - apply_router_weight_on_input=apply_router_weight_on_input,

768 - )

769 -

770 -

709 # W4A8

710 def run_cutlass_moe_w4a8_fp8(

711 output: torch.Tensor,

vllm/model_executor/layers/fused_moe/flashinfer_cutedsl_moe.py

-39

@@ -335,42 +335,3 @@ def flashinfer_cutedsl_moe_masked(

335 alpha_dtype=get_cute_dtype(w2_alpha),

336 ) # in logical [m, k, l]

337 out = out.permute(2, 0, 1)

338 -

339 -

340 -def flashinfer_cutedsl_moe_fp4(

341 - hidden_states: torch.Tensor,

342 - w1: torch.Tensor,

343 - w2: torch.Tensor,

344 - topk_weights: torch.Tensor,

345 - topk_ids: torch.Tensor,

346 - quant_config: FusedMoEQuantConfig,

347 - inplace: bool = False,

348 - activation: str = "silu",

349 - global_num_experts: int = -1,

350 - expert_map: torch.Tensor | None = None,

351 - apply_router_weight_on_input: bool = False,

352 -) -> torch.Tensor:

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

354 - create_flashinfer_prepare_finalize,

355 - )

356 -

357 - fused_experts = mk.FusedMoEModularKernel(

358 - create_flashinfer_prepare_finalize(use_dp=False), # could be swapped later

359 - FlashInferCuteDSLExperts(

360 - out_dtype=hidden_states.dtype,

361 - quant_config=quant_config,

362 - ),

363 - )

364 -

365 - return fused_experts(

366 - hidden_states=hidden_states,

367 - w1=w1,

368 - w2=w2,

369 - topk_weights=topk_weights,

370 - topk_ids=topk_ids,

371 - inplace=inplace,

372 - activation=activation,

373 - global_num_experts=global_num_experts,

374 - expert_map=expert_map,

375 - apply_router_weight_on_input=apply_router_weight_on_input,

376 - )

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py

+7 -11

@@ -355,21 +355,17 @@ def create_flashinfer_prepare_finalize(

355 use_deepseek_fp8_block_scale: bool = False,

356 ) -> FlashInferCutlassMoEPrepareAndFinalize | MoEPrepareAndFinalizeNoEP:

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

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

359 - # once we complete the FP8 refactor.

360 - if use_nvfp4:

361 - if enable_alltoallv:

362 - return FlashInferAllToAllMoEPrepareAndFinalize(use_dp)

363 - else:

364 - return FlashInferAllGatherMoEPrepareAndFinalize(use_dp)

358

366 - # FP8 DP path currently supported via AllGather.

359 if use_dp:

360 + if enable_alltoallv:

361 + assert use_nvfp4

362 + return FlashInferAllToAllMoEPrepareAndFinalize(use_dp)

363 return FlashInferAllGatherMoEPrepareAndFinalize(

364 use_dp=True,

365 use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale,

366 )

367 else:

373 - # NOTE(rob): CUTLASS FP8 block quant executes the input

374 - # quantzation and grouped gemm in a single kernel.

375 - return MoEPrepareAndFinalizeNoEP(defer_input_quant=use_deepseek_fp8_block_scale)

368 + # CUTLASS FP8 BLOCK and CUTLASS NVFP4 apply input quantization

369 + # in a single call with the MoE experts kernel.

370 + defer_input_quant = use_deepseek_fp8_block_scale or use_nvfp4

371 + return MoEPrepareAndFinalizeNoEP(defer_input_quant=defer_input_quant)

vllm/model_executor/layers/fused_moe/fused_marlin_moe.py

+5 -2

@@ -540,9 +540,10 @@ class MarlinExpertsBase(mk.FusedMoEPermuteExpertsUnpermute):

540 # TODO (varun) : Enable activation quantization

541 assert (

542 quant_config.use_mxfp4_w4a16

543 + or quant_config.use_nvfp4_w4a16

544 or quant_config.use_int4_w4a16

545 or quant_config.use_fp8_w8a16

545 - ), "Supports only mxfp4_w4a16, int4_w4a16 or fp8_w8a16"

546 + ), "Supports only {mxfp,nvfp,int}4_w4a16 or fp8_w8a16"

547 self.w13_g_idx = w13_g_idx

548 self.w2_g_idx = w2_g_idx

549 self.w13_g_idx_sort_indices = w13_g_idx_sort_indices

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

556 # uint4b8 will be set for int4 weight and float4_e2m1f will be used for mxfp4

557 if self.quant_config.use_int4_w4a16:

558 return scalar_types.uint4b8.id

558 - elif self.quant_config.use_mxfp4_w4a16:

559 + elif self.quant_config.use_mxfp4_w4a16 or self.quant_config.use_nvfp4_w4a16:

560 return scalar_types.float4_e2m1f.id

561 elif (

562 self.quant_config.use_fp8_w8a16

@@ -692,6 +693,8 @@ class MarlinExperts(MarlinExpertsBase):

693 gating_output=None,

694 topk_weights=topk_weights,

695 topk_ids=topk_ids,

696 + global_scale1=self.g1_alphas,

697 + global_scale2=self.g2_alphas,

698 quant_type_id=self.quant_type_id,

699 apply_router_weight_on_input=apply_router_weight_on_input,

700 global_num_experts=global_num_experts,

vllm/model_executor/layers/fused_moe/layer.py

+1 -9

@@ -38,9 +38,6 @@ from vllm.model_executor.layers.fused_moe.routing_simulator import RoutingSimula

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

39 QuantizationConfig,

40 )

41 -from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

42 - is_flashinfer_supporting_global_sf,

43 -)

41 from vllm.platforms import current_platform

42 from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe

43 from vllm.utils.math_utils import cdiv, round_up

@@ -1125,14 +1122,9 @@ class FusedMoE(CustomOp):

1122 global_expert_id = expert_id

1123 expert_id = self._map_global_expert_id_to_local_expert_id(global_expert_id)

1124

1128 - allow_flashinfer = getattr(self.quant_method, "allow_flashinfer", False)

1129 - moe_backend = getattr(self.quant_method, "flashinfer_moe_backend", None)

1130 -

1125 use_global_sf = (

1132 - allow_flashinfer

1133 - and is_flashinfer_supporting_global_sf(moe_backend)

1126 + getattr(self.quant_method, "use_global_sf", False)

1127 and "input_scale" in weight_name

1135 - and quant_method_name == "ModelOptNvFp4FusedMoE"

1128 )

1129

1130 if expert_id == -1 and not use_global_sf:

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

+280

@@ -0,0 +1,280 @@

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

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

3 +from enum import Enum

4 +

5 +import torch

6 +

7 +import vllm.envs as envs

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

9 +from vllm.logger import init_logger

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

11 + FusedMoEConfig,

12 + FusedMoEQuantConfig,

13 + nvfp4_moe_quant_config,

14 + nvfp4_w4a16_moe_quant_config,

15 +)

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

17 + CutlassExpertsFp4,

18 +)

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

20 + FlashInferExperts,

21 +)

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

23 + MarlinExperts,

24 +)

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

26 + MoEPrepareAndFinalizeNoEP,

27 +)

28 +from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (

29 + is_flashinfer_fp4_cutedsl_moe_available,

30 + is_flashinfer_fp4_cutlass_moe_available,

31 + prepare_nvfp4_moe_layer_for_fi_or_cutlass,

32 +)

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

34 + FlashinferMoeBackend,

35 + get_flashinfer_moe_backend,

36 +)

37 +from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (

38 + is_fp4_marlin_supported,

39 + prepare_nvfp4_moe_layer_for_marlin,

40 +)

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

42 + cutlass_fp4_supported,

43 +)

44 +

45 +logger = init_logger(__name__)

46 +

47 +

48 +class NvFp4MoeBackend(Enum):

49 + FLASHINFER_CUTLASS = "FlashInfer CUTLASS"

50 + FLASHINFER_TRTLLM = "FlashInfer TRTLLM"

51 + FLASHINFER_CUTEDSL = "FlashInfer CUTEDSL"

52 + VLLM_CUTLASS = "vLLM CUTASS"

53 + MARLIN = "vLLM MARLIN"

54 +

55 +

56 +FLASHINFER_NVFP4_MOE_BACKENDS = [

57 + NvFp4MoeBackend.FLASHINFER_CUTLASS,

58 + NvFp4MoeBackend.FLASHINFER_TRTLLM,

59 + NvFp4MoeBackend.FLASHINFER_CUTEDSL,

60 +]

61 +

62 +fi_2_vllm_backend_map: dict[FlashinferMoeBackend, NvFp4MoeBackend] = {

63 + FlashinferMoeBackend.CUTLASS: NvFp4MoeBackend.FLASHINFER_CUTLASS,

64 + FlashinferMoeBackend.TENSORRT_LLM: NvFp4MoeBackend.FLASHINFER_TRTLLM,

65 + FlashinferMoeBackend.CUTEDSL: NvFp4MoeBackend.FLASHINFER_CUTEDSL,

66 +}

67 +

68 +

69 +def is_global_sf_supported_for_nvfp4_backend(backend: NvFp4MoeBackend) -> bool:

70 + # Checks whether `backend` supports quantizing with scaling factors

71 + # of all experts in Expert Parallel Mode when all experts are not

72 + # on the same rank.

73 +

74 + return backend in [

75 + NvFp4MoeBackend.FLASHINFER_CUTLASS,

76 + NvFp4MoeBackend.FLASHINFER_TRTLLM,

77 + ]

78 +

79 +

80 +def select_nvfp4_moe_backend() -> NvFp4MoeBackend:

81 + def _make_log_backend(backend: NvFp4MoeBackend):

82 + return f"Using {backend.value} backend for NvFp4 MoE"

83 +

84 + if cutlass_fp4_supported() and not envs.VLLM_TEST_FORCE_FP8_MARLIN:

85 + allow_flashinfer = (

86 + is_flashinfer_fp4_cutlass_moe_available()

87 + or is_flashinfer_fp4_cutedsl_moe_available()

88 + )

89 + if allow_flashinfer and envs.VLLM_USE_FLASHINFER_MOE_FP4:

90 + backend = fi_2_vllm_backend_map[get_flashinfer_moe_backend()]

91 + else:

92 + backend = NvFp4MoeBackend.VLLM_CUTLASS

93 + elif is_fp4_marlin_supported():

94 + backend = NvFp4MoeBackend.MARLIN

95 + else:

96 + raise ValueError("No NvFp4 kernel backend available for NvFp4 MoE.")

97 +

98 + # Log warning if FI backend requested but not available.

99 + if (

100 + backend not in FLASHINFER_NVFP4_MOE_BACKENDS

101 + and envs.VLLM_USE_FLASHINFER_MOE_FP4

102 + ):

103 + logger.warning_once(

104 + "Requested FlashInfer backend for NvFp4 MoE, but it's not available. "

105 + "Falling back to %s for NvFp4 MoE",

106 + backend.value,

107 + scope="local",

108 + )

109 + else:

110 + logger.info_once(_make_log_backend(backend), scope="local")

111 + return backend

112 +

113 +

114 +def convert_to_nvfp4_moe_kernel_format(

115 + nvfp4_backend: NvFp4MoeBackend,

116 + layer: torch.nn.Module,

117 + w13: torch.Tensor,

118 + w13_scale: torch.Tensor,

119 + w13_scale_2: torch.Tensor,

120 + a13_scale: torch.Tensor | None,

121 + w2: torch.Tensor,

122 + w2_scale: torch.Tensor,

123 + w2_scale_2: torch.Tensor,

124 + a2_scale: torch.Tensor | None,

125 + is_act_and_mul: bool,

126 +) -> tuple[

127 + torch.Tensor,

128 + torch.Tensor,

129 + torch.Tensor,

130 + torch.Tensor,

131 + torch.Tensor,

132 + torch.Tensor,

133 + torch.Tensor,

134 + torch.Tensor,

135 +]:

136 + if (

137 + nvfp4_backend in FLASHINFER_NVFP4_MOE_BACKENDS

138 + or nvfp4_backend == NvFp4MoeBackend.VLLM_CUTLASS

139 + ):

140 + (

141 + w13,

142 + w13_scale,

143 + w13_scale_2,

144 + a13_scale,

145 + w2,

146 + w2_scale,

147 + w2_scale_2,

148 + a2_scale,

149 + ) = prepare_nvfp4_moe_layer_for_fi_or_cutlass(

150 + backend=nvfp4_backend,

151 + layer=layer,

152 + w13=w13,

153 + w13_scale=w13_scale,

154 + w13_scale_2=w13_scale_2,

155 + a13_scale=a13_scale,

156 + w2=w2,

157 + w2_scale=w2_scale,

158 + w2_scale_2=w2_scale_2,

159 + a2_scale=a2_scale,

160 + is_act_and_mul=is_act_and_mul,

161 + )

162 + elif nvfp4_backend == NvFp4MoeBackend.MARLIN:

163 + a13_scale = None

164 + a2_scale = None

165 + (

166 + w13,

167 + w13_scale,

168 + w13_scale_2,

169 + w2,

170 + w2_scale,

171 + w2_scale_2,

172 + ) = prepare_nvfp4_moe_layer_for_marlin(

173 + layer=layer,

174 + w13=w13,

175 + w13_scale=w13_scale,

176 + w13_scale_2=w13_scale_2,

177 + w2=w2,

178 + w2_scale=w2_scale,

179 + w2_scale_2=w2_scale_2,

180 + )

181 + else:

182 + raise ValueError(f"Unknown NvFp4 backend for MoE: {nvfp4_backend}")

183 +

184 + return (

185 + w13,

186 + w13_scale,

187 + w13_scale_2,

188 + a13_scale,

189 + w2,

190 + w2_scale,

191 + w2_scale_2,

192 + a2_scale,

193 + )

194 +

195 +

196 +def make_nvfp4_moe_quant_config(

197 + backend: NvFp4MoeBackend,

198 + w13_scale: torch.Tensor,

199 + w2_scale: torch.Tensor,

200 + w13_scale_2: torch.Tensor,

201 + w2_scale_2: torch.Tensor,

202 + a13_scale: torch.Tensor,

203 + a2_scale: torch.Tensor,

204 +) -> FusedMoEQuantConfig | None:

205 + UNSUPPORTED = [NvFp4MoeBackend.FLASHINFER_TRTLLM]

206 + if backend in UNSUPPORTED:

207 + return None

208 +

209 + elif backend == NvFp4MoeBackend.MARLIN:

210 + return nvfp4_w4a16_moe_quant_config(

211 + g1_alphas=w13_scale_2,

212 + g2_alphas=w2_scale_2,

213 + w1_scale=w13_scale,

214 + w2_scale=w2_scale,

215 + )

216 +

217 + g1_alphas = a13_scale * w13_scale_2

218 + g2_alphas = a2_scale * w2_scale_2

219 + return nvfp4_moe_quant_config(

220 + g1_alphas=g1_alphas,

221 + g2_alphas=g2_alphas,

222 + a1_gscale=(1.0 / a13_scale),

223 + a2_gscale=(1.0 / a2_scale),

224 + w1_scale=w13_scale,

225 + w2_scale=w2_scale,

226 + )

227 +

228 +

229 +def make_nvfp4_moe_kernel(

230 + backend: NvFp4MoeBackend,

231 + quant_config: FusedMoEQuantConfig,

232 + moe_config: FusedMoEConfig,

233 +) -> mk.FusedMoEModularKernel | None:

234 + assert moe_config.dp_size == 1

235 +

236 + UNSUPPORTED_BACKENDS = [

237 + # TRTLLM does not use the modular kernl abstraction.

238 + NvFp4MoeBackend.FLASHINFER_TRTLLM,

239 + # CUTEDSL is used with BATCHED (masked) format only.

240 + # TODO: add here once we support dp/ep via the oracle.

241 + NvFp4MoeBackend.FLASHINFER_CUTEDSL,

242 + ]

243 +

244 + if backend in UNSUPPORTED_BACKENDS:

245 + return None

246 +

247 + elif backend == NvFp4MoeBackend.FLASHINFER_CUTLASS:

248 + return mk.FusedMoEModularKernel(

249 + MoEPrepareAndFinalizeNoEP(defer_input_quant=True),

250 + FlashInferExperts(

251 + out_dtype=moe_config.in_dtype,

252 + quant_config=quant_config,

253 + ep_rank=moe_config.ep_rank,

254 + ep_size=moe_config.ep_size,

255 + tp_rank=moe_config.tp_rank,

256 + tp_size=moe_config.tp_size,

257 + use_dp=False,

258 + use_deepseek_fp8_block_scale=False,

259 + ),

260 + )

261 +

262 + elif backend == NvFp4MoeBackend.VLLM_CUTLASS:

263 + return mk.FusedMoEModularKernel(

264 + MoEPrepareAndFinalizeNoEP(defer_input_quant=True),

265 + CutlassExpertsFp4(

266 + out_dtype=moe_config.in_dtype,

267 + # TODO(rob): see what impact this has on expert map?

268 + max_experts_per_worker=moe_config.num_experts,

269 + quant_config=quant_config,

270 + ),

271 + )

272 +

273 + elif backend == NvFp4MoeBackend.MARLIN:

274 + return mk.FusedMoEModularKernel(

275 + MoEPrepareAndFinalizeNoEP(),

276 + MarlinExperts(quant_config=quant_config),

277 + )

278 +

279 + else:

280 + raise ValueError(f"Unknown NvFp4 MoE backend: {backend}")

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

+115 -240

@@ -11,7 +11,6 @@ from compressed_tensors.quantization import (

11 QuantizationArgs,

12 QuantizationStrategy,

13 )

14 -from torch.nn.parameter import Parameter

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

16 from vllm import _custom_ops as ops

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

33 int4_w4afp8_moe_quant_config,

34 int8_w8a8_moe_quant_config,

35 int8_w8a16_moe_quant_config,

37 - nvfp4_moe_quant_config,

36 )

37 from vllm.model_executor.layers.fused_moe.cpu_fused_moe import select_experts

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

41 - is_valid_flashinfer_cutlass_fused_moe,

42 -)

38 from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (

39 BatchedMarlinExperts,

40 MarlinExperts,

@@ -51,6 +46,15 @@ from vllm.model_executor.layers.fused_moe.oracle.fp8 import (

46 make_fp8_moe_kernel,

47 select_fp8_moe_backend,

48 )

49 +from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (

50 + FLASHINFER_NVFP4_MOE_BACKENDS,

51 + NvFp4MoeBackend,

52 + convert_to_nvfp4_moe_kernel_format,

53 + is_global_sf_supported_for_nvfp4_backend,

54 + make_nvfp4_moe_kernel,

55 + make_nvfp4_moe_quant_config,

56 + select_nvfp4_moe_backend,

57 +)

58 from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa

59 WNA16_SUPPORTED_BITS,

60 WNA16_SUPPORTED_TYPES_MAP,

@@ -58,14 +62,9 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compress

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

63 build_flashinfer_fp4_cutlass_moe_prepare_finalize,

64 flashinfer_trtllm_fp4_moe,

61 - prepare_static_weights_for_trtllm_fp4_moe,

62 - reorder_w1w3_to_w3w1,

65 + flashinfer_trtllm_fp4_routed_moe,

66 select_nvfp4_gemm_impl,

67 )

65 -from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

66 - FlashinferMoeBackend,

67 - get_flashinfer_moe_backend,

68 -)

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

69 process_fp8_input_tensor_strategy_moe,

70 process_fp8_weight_tensor_strategy_moe,

@@ -77,20 +76,15 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (

76 marlin_make_workspace_new,

77 marlin_moe_permute_scales,

78 )

80 -from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (

81 - prepare_moe_fp4_layer_for_marlin,

82 -)

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

80 convert_bf16_scales_to_fp8,

81 convert_packed_uint4b8_to_signed_int4_inplace,

86 - swizzle_blockscale,

82 )

83 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (

84 normalize_e4m3fn_to_e4m3fnuz,

85 )

86 from vllm.model_executor.utils import replace_parameter, set_weight_attrs

87 from vllm.platforms import CpuArchEnum, current_platform

93 -from vllm.scalar_type import scalar_types

89 logger = init_logger(__name__)

@@ -218,31 +212,19 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):

212

213 class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

214 def __init__(self, moe: FusedMoEConfig, layer_name: str | None = None):

221 - from vllm.model_executor.layers.quantization.utils.nvfp4_moe_support import ( # noqa: E501

222 - detect_nvfp4_moe_support,

223 - )

215 + if not moe.is_act_and_mul:

216 + raise ValueError(

217 + "CompressedTensorsW4A4Nvfp4MoEMethod does not yet "

218 + "support non gated MoE models."

219 + )

220

221 super().__init__(moe)

226 - _nvfp4 = detect_nvfp4_moe_support(self.__class__.__name__)

227 - self.cutlass_nvfp4_supported = _nvfp4.cutlass_supported

228 - self.allow_flashinfer = _nvfp4.allow_flashinfer

229 - self.use_marlin = _nvfp4.use_marlin

222 self.group_size = 16

231 - self.layer_name = layer_name

232 - self.marlin_input_dtype = (

233 - get_marlin_input_dtype(layer_name) if self.use_marlin else None

234 - )

235 - self.flashinfer_moe_backend = None

236 - if self.allow_flashinfer:

237 - self.flashinfer_moe_backend = get_flashinfer_moe_backend()

238 - logger.info_once(

239 - f"Using FlashInfer {self.flashinfer_moe_backend.value} kernels"

240 - " for CompressedTensorsW4A4Nvfp4MoEMethod."

241 - )

242 - elif self.use_marlin:

243 - logger.info_once("Using Marlin for CompressedTensorsW4A4Nvfp4MoEMethod.")

244 - else:

245 - logger.info_once("Using Cutlass for CompressedTensorsW4A4Nvfp4MoEMethod.")

223 + self.nvfp4_backend = select_nvfp4_moe_backend()

224 + self.use_global_sf = is_global_sf_supported_for_nvfp4_backend(

225 + self.nvfp4_backend

226 + )

227 + self.kernel: mk.FusedMoEModularKernel | None = None

228

229 def create_weights(

230 self,

@@ -355,7 +337,13 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

337 set_weight_attrs(w2_input_scale, extra_weight_attrs)

338

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

358 - # From packed to weight

340 + """

341 + Convert NVFP4 MoE weights into kernel format and setup the kernel.

342 + """

343 + # NOTE(rob): wN_weight_packed -> wN_weight is because ModularKernelMethod

344 + # requires this naming convention. However, the name change breaks

345 + # reloading because the state dict no longer matches disk. Once we

346 + # remove MKM, we should revert this change to ensure compatibility.

347 layer.w13_weight = torch.nn.Parameter(

348 layer.w13_weight_packed.data, requires_grad=False

349 )

@@ -366,144 +354,79 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

354 )

355 delattr(layer, "w2_weight_packed")

356

369 - # reorder GEMM1 weights and block scales for FlashInfer CUTLASS kernel.

370 - if self.allow_flashinfer:

371 - w, s = reorder_w1w3_to_w3w1(

372 - layer.w13_weight.data, layer.w13_weight_scale.data, dim=-2

373 - )

374 - layer.w13_weight = torch.nn.Parameter(w, requires_grad=False)

375 - layer.w13_weight_scale = torch.nn.Parameter(s, requires_grad=False)

376 -

377 - if not torch.allclose(

357 + # Use a single gscale for w13.

358 + if self.moe.is_act_and_mul and not torch.allclose(

359 layer.w13_weight_global_scale[:, 0], layer.w13_weight_global_scale[:, 1]

360 ):

361 logger.warning_once(

362 "w1_weight_global_scale must match w3_weight_global_scale. "

382 - "Accuracy may be affected."

383 - )

384 -

385 - # Take inverse of global scale saved to disk

386 - layer.w13_weight_scale_2 = torch.nn.Parameter(

387 - 1 / layer.w13_weight_global_scale[:, 0], requires_grad=False

388 - )

389 -

390 - layer.w2_weight_scale_2 = torch.nn.Parameter(

391 - 1 / layer.w2_weight_global_scale.data, requires_grad=False

392 - )

393 -

394 - if self.use_marlin:

395 - prepare_moe_fp4_layer_for_marlin(layer, input_dtype=self.marlin_input_dtype)

396 - return

397 - # w13

398 - if (

399 - self.allow_flashinfer

400 - and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

401 - ):

402 - w13_input_global_scale = (

403 - layer.w13_input_global_scale.min()

404 - .to(torch.float32)

405 - .expand(layer.num_experts)

406 - )

407 - else:

408 - w13_input_global_scale = layer.w13_input_global_scale.min(dim=1).values.to(

409 - torch.float32

410 - )

411 - layer.g1_alphas = torch.nn.Parameter(

412 - ((1 / w13_input_global_scale) * layer.w13_weight_scale_2),

413 - requires_grad=False,

414 - )

415 -

416 - layer.w13_input_scale_quant = torch.nn.Parameter(

417 - (w13_input_global_scale), requires_grad=False

418 - )

419 -

420 - # w2

421 - if (

422 - self.allow_flashinfer

423 - and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

424 - ):

425 - w2_input_global_scale = (

426 - layer.w2_input_global_scale.min()

427 - .to(torch.float32)

428 - .expand(layer.num_experts)

429 - )

430 - else:

431 - w2_input_global_scale = layer.w2_input_global_scale

432 -

433 - layer.g2_alphas = torch.nn.Parameter(

434 - ((1 / w2_input_global_scale) * layer.w2_weight_scale_2).to(torch.float32),

435 - requires_grad=False,

436 - )

437 -

438 - layer.w2_input_scale_quant = torch.nn.Parameter(

439 - (w2_input_global_scale), requires_grad=False

363 + "Accuracy may be affected.",

364 + )

365 + w13_weight_global_scale = layer.w13_weight_global_scale[:, 0].contiguous()

366 +

367 + # Shuffle weights into the NvFp4 kernel format.

368 + (

369 + w13,

370 + w13_scale,

371 + w13_scale_2,

372 + a13_scale,

373 + w2,

374 + w2_scale,

375 + w2_scale_2,

376 + a2_scale,

377 + ) = convert_to_nvfp4_moe_kernel_format(

378 + nvfp4_backend=self.nvfp4_backend,

379 + layer=layer,

380 + w13=layer.w13_weight,

381 + w13_scale=layer.w13_weight_scale,

382 + w13_scale_2=(1.0 / w13_weight_global_scale),

383 + a13_scale=(1.0 / layer.w13_input_global_scale),

384 + w2=layer.w2_weight,

385 + w2_scale=layer.w2_weight_scale,

386 + w2_scale_2=(1.0 / layer.w2_weight_global_scale),

387 + a2_scale=(1.0 / layer.w2_input_global_scale),

388 + is_act_and_mul=self.moe.is_act_and_mul,

389 )

390

442 - # TensorRT-LLM specific processing

443 - if (

444 - self.allow_flashinfer

445 - and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

446 - ):

447 - # Prepare static weights for TRT-LLM kernel

448 - # alternate: prepare_static_weight_layouts_for_trtllm_moe

449 - (

450 - gemm1_weights_fp4_shuffled,

451 - gemm1_scales_fp4_shuffled,

452 - gemm2_weights_fp4_shuffled,

453 - gemm2_scales_fp4_shuffled,

454 - ) = prepare_static_weights_for_trtllm_fp4_moe(

455 - layer.w13_weight,

456 - layer.w2_weight,

457 - layer.w13_weight_scale,

458 - layer.w2_weight_scale,

459 - layer.w2_weight.size(-2), # hidden_size

460 - layer.w13_weight.size(-2) // 2, # intermediate_size

461 - layer.w13_weight.size(0), # num_experts

462 - )

463 - logger.debug_once("Finished shuffling weights for TRT-LLM MOE")

464 -

465 - layer.w13_weight = Parameter(

466 - gemm1_weights_fp4_shuffled, requires_grad=False

467 - )

468 - layer.w2_weight = Parameter(gemm2_weights_fp4_shuffled, requires_grad=False)

469 - layer.w13_weight_scale = Parameter(

470 - gemm1_scales_fp4_shuffled, requires_grad=False

471 - )

472 - layer.w2_weight_scale = Parameter(

473 - gemm2_scales_fp4_shuffled, requires_grad=False

474 - )

475 -

476 - # Additional parameter needed for TRT-LLM

477 - layer.g1_scale_c = Parameter(

478 - (layer.w2_input_scale_quant * layer.g1_alphas).to(torch.float32),

479 - requires_grad=False,

480 - )

481 - else:

482 - # swizzle weight scales

483 - layer.w13_weight_scale = torch.nn.Parameter(

484 - swizzle_blockscale(layer.w13_weight_scale), requires_grad=False

485 - )

391 + replace_parameter(layer, "w13_weight", w13)

392 + replace_parameter(layer, "w13_weight_scale", w13_scale)

393 + replace_parameter(layer, "w2_weight", w2)

394 + replace_parameter(layer, "w2_weight_scale", w2_scale)

395 + layer.w13_weight_scale_2 = w13_scale_2

396 + layer.w2_weight_scale_2 = w2_scale_2

397 + layer.w13_input_scale = a13_scale

398 + layer.w2_input_scale = a2_scale

399

487 - layer.w2_weight_scale = torch.nn.Parameter(

488 - swizzle_blockscale(layer.w2_weight_scale), requires_grad=False

400 + # Initialize the kernel that will be called in apply().

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

402 + use_dp = self.moe.dp_size > 1

403 + if self.moe_quant_config is not None and not use_dp:

404 + self.kernel = make_nvfp4_moe_kernel(

405 + backend=self.nvfp4_backend,

406 + quant_config=self.moe_quant_config,

407 + moe_config=self.moe,

408 )

409

410 def maybe_make_prepare_finalize(

411 self,

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

413 ) -> mk.FusedMoEPrepareAndFinalize | None:

495 - if self.use_marlin or (

496 - self.allow_flashinfer

497 - and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

498 - ):

414 + UNSUPPORTED = [NvFp4MoeBackend.MARLIN, NvFp4MoeBackend.FLASHINFER_TRTLLM]

415 + if self.nvfp4_backend in UNSUPPORTED:

416 return None

500 - elif not self.allow_flashinfer:

417 + elif self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_CUTLASS:

418 + # TP case: avoid convert to ModularKernelMethod - to be refactored.

419 + if self.moe.dp_size == 1:

420 + return None

421 + # For now, fp4 moe only works with the flashinfer dispatcher.

422 + prepare_finalize = build_flashinfer_fp4_cutlass_moe_prepare_finalize(

423 + self.moe

424 + )

425 + logger.debug_once("%s", prepare_finalize.__class__.__name__)

426 + return prepare_finalize

427 + else:

428 return super().maybe_make_prepare_finalize(routing_tables)

429

503 - prepare_finalize = build_flashinfer_fp4_cutlass_moe_prepare_finalize(self.moe)

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

505 - return prepare_finalize

506 -

430 def select_gemm_impl(

431 self,

432 prepare_finalize: mk.FusedMoEPrepareAndFinalize,

@@ -514,7 +437,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

437 experts = select_nvfp4_gemm_impl(

438 self.moe,

439 self.moe_quant_config,

517 - allow_flashinfer=self.allow_flashinfer,

440 + allow_flashinfer=(self.nvfp4_backend in FLASHINFER_NVFP4_MOE_BACKENDS),

441 )

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

443 return experts

@@ -522,19 +445,14 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

445 def get_fused_moe_quant_config(

446 self, layer: torch.nn.Module

447 ) -> FusedMoEQuantConfig | None:

525 - if (

526 - self.use_marlin

527 - or self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

528 - ):

529 - return None

530 -

531 - return nvfp4_moe_quant_config(

532 - g1_alphas=layer.g1_alphas,

533 - g2_alphas=layer.g2_alphas,

534 - a1_gscale=layer.w13_input_scale_quant,

535 - a2_gscale=layer.w2_input_scale_quant,

536 - w1_scale=layer.w13_weight_scale,

448 + return make_nvfp4_moe_quant_config(

449 + backend=self.nvfp4_backend,

450 + w13_scale=layer.w13_weight_scale,

451 w2_scale=layer.w2_weight_scale,

452 + w13_scale_2=layer.w13_weight_scale_2,

453 + w2_scale_2=layer.w2_weight_scale_2,

454 + a13_scale=layer.w13_input_scale,

455 + a2_scale=layer.w2_input_scale,

456 )

457

458 def apply(

@@ -546,14 +464,9 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

464 assert layer.activation == "silu", "Only SiLU activation is supported."

465

466 if (

549 - self.allow_flashinfer

550 - and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

467 + self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM

468 + and not layer.enable_eplb

469 ):

552 - if layer.enable_eplb:

553 - raise NotImplementedError(

554 - "EPLB not supported for `CompressedTensorsW4A4MoEMethod` yet."

555 - )

556 -

470 return flashinfer_trtllm_fp4_moe(

471 layer=layer,

472 x=x,

@@ -566,79 +479,41 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

479 e_score_correction_bias=layer.e_score_correction_bias,

480 )

481

482 + # Hidden_states in select_experts is only used to extract metadata

483 + if isinstance(x, tuple):

484 + x_routing, _ = x

485 + else:

486 + x_routing = x

487 topk_weights, topk_ids = layer.select_experts(

570 - hidden_states=x,

488 + hidden_states=x_routing,

489 router_logits=router_logits,

490 )

491

574 - if self.use_marlin:

575 - return fused_marlin_moe(

492 + # EPLB path

493 + if self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM:

494 + assert layer.enable_eplb

495 + return flashinfer_trtllm_fp4_routed_moe(

496 + layer=layer,

497 + x=x,

498 + topk_ids=topk_ids,

499 + topk_weights=topk_weights,

500 + top_k=layer.top_k,

501 + global_num_experts=layer.global_num_experts,

502 + )

503 + else:

504 + assert self.kernel is not None

505 + return self.kernel(

506 x,

507 layer.w13_weight,

508 layer.w2_weight,

579 - None,

580 - None,

581 - layer.w13_weight_scale,

582 - layer.w2_weight_scale,

583 - router_logits,

509 topk_weights,

510 topk_ids,

586 - global_scale1=layer.w13_weight_scale_2,

587 - global_scale2=layer.w2_weight_scale_2,

588 - quant_type_id=scalar_types.float4_e2m1f.id,

589 - apply_router_weight_on_input=layer.apply_router_weight_on_input,

590 - global_num_experts=layer.global_num_experts,

591 - expert_map=layer.expert_map,

592 - input_dtype=self.marlin_input_dtype,

593 - workspace=layer.workspace,

594 - )

595 -

596 - # FlashInfer fused experts path

597 - elif self.allow_flashinfer:

598 - from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( # noqa: E501

599 - flashinfer_cutlass_moe_fp4,

600 - )

601 -

602 - assert is_valid_flashinfer_cutlass_fused_moe(

603 - x, layer.w13_weight, layer.w2_weight

604 - ), "Flashinfer CUTLASS Fused MoE not applicable!"

605 -

606 - assert self.moe_quant_config is not None

607 -

608 - return flashinfer_cutlass_moe_fp4(

609 - hidden_states=x,

610 - w1=layer.w13_weight,

611 - w2=layer.w2_weight,

612 - topk_weights=topk_weights,

613 - topk_ids=topk_ids,

614 - quant_config=self.moe_quant_config,

615 - inplace=False, # TODO(shuw): fix later, now output is high prec

511 + inplace=False,

512 activation=layer.activation,

513 global_num_experts=layer.global_num_experts,

514 expert_map=layer.expert_map,

515 apply_router_weight_on_input=layer.apply_router_weight_on_input,

516 )

621 - else:

622 - # If no modular kernel is provided, use cutlass_moe_fp4 for TP case

623 - # only (no EP).

624 - from vllm.model_executor.layers.fused_moe.cutlass_moe import cutlass_moe_fp4

625 -

626 - assert self.moe_quant_config is not None

627 - return cutlass_moe_fp4(

628 - a=x,

629 - w1_fp4=layer.w13_weight,

630 - w2_fp4=layer.w2_weight,

631 - topk_weights=topk_weights,

632 - topk_ids=topk_ids,

633 - quant_config=self.moe_quant_config,

634 - expert_map=layer.expert_map,

635 - apply_router_weight_on_input=layer.apply_router_weight_on_input,

636 - # TODO(bnell): derive these from arguments

637 - m=x.shape[0],

638 - n=layer.w2_weight.shape[2] * 2,

639 - k=x.shape[1],

640 - e=layer.w13_weight.shape[0],

641 - ).to(x.dtype)

517

518

519 class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

vllm/model_executor/layers/quantization/modelopt.py

+87 -285

@@ -15,9 +15,7 @@ from vllm.attention.layer import Attention

15 from vllm.logger import init_logger

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

17 FusedMoEQuantConfig,

18 - nvfp4_moe_quant_config,

18 )

20 -from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe

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

20 FusedMoE,

21 FusedMoEMethodBase,

@@ -30,6 +28,15 @@ from vllm.model_executor.layers.fused_moe.oracle.fp8 import (

28 make_fp8_moe_quant_config,

29 select_fp8_moe_backend,

30 )

31 +from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (

32 + FLASHINFER_NVFP4_MOE_BACKENDS,

33 + NvFp4MoeBackend,

34 + convert_to_nvfp4_moe_kernel_format,

35 + is_global_sf_supported_for_nvfp4_backend,

36 + make_nvfp4_moe_kernel,

37 + make_nvfp4_moe_quant_config,

38 + select_nvfp4_moe_backend,

39 +)

40 from vllm.model_executor.layers.linear import (

41 LinearBase,

42 LinearMethodBase,

@@ -45,16 +52,11 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (

52 build_flashinfer_fp4_cutlass_moe_prepare_finalize,

53 flashinfer_trtllm_fp4_moe,

54 flashinfer_trtllm_fp4_routed_moe,

48 - prepare_static_weights_for_trtllm_fp4_moe,

49 - reorder_w1w3_to_w3w1,

55 select_nvfp4_gemm_impl,

56 )

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

53 - FlashinferMoeBackend,

58 apply_fi_trtllm_fp8_per_tensor_moe,

59 build_flashinfer_fp8_cutlass_moe_prepare_finalize,

56 - get_flashinfer_moe_backend,

57 - is_flashinfer_supporting_global_sf,

60 select_cutlass_fp8_gemm_impl,

61 )

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

@@ -69,7 +71,6 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (

71 apply_fp4_marlin_linear,

72 is_fp4_marlin_supported,

73 prepare_fp4_layer_for_marlin,

72 - prepare_moe_fp4_layer_for_marlin,

74 )

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

76 GroupShape,

@@ -89,7 +90,6 @@ from vllm.model_executor.parameter import (

90 PerTensorScaleParameter,

91 )

92 from vllm.model_executor.utils import replace_parameter

92 -from vllm.scalar_type import scalar_types

93 from vllm.utils.flashinfer import (

94 flashinfer_scaled_fp4_mm,

95 has_flashinfer,

@@ -1327,43 +1327,32 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1327 quant_config: ModelOptNvFp4Config,

1328 layer: FusedMoE,

1329 ) -> None:

1330 - from vllm.model_executor.layers.quantization.utils.nvfp4_moe_support import (

1331 - detect_nvfp4_moe_support, # noqa: E501

1332 - )

1333 -

1330 super().__init__(layer.moe_config)

1331 self.quant_config = quant_config

1336 - self.layer = layer

1337 - _nvfp4 = detect_nvfp4_moe_support(self.__class__.__name__)

1338 - self.cutlass_nvfp4_supported = _nvfp4.cutlass_supported

1339 - self.allow_flashinfer = _nvfp4.allow_flashinfer

1340 - self.use_marlin = _nvfp4.use_marlin

1341 - self.marlin_input_dtype = None

1342 - self.flashinfer_moe_backend = None

1343 - if self.allow_flashinfer:

1344 - self.flashinfer_moe_backend = get_flashinfer_moe_backend()

1345 - logger.info_once(

1346 - f"Using FlashInfer {self.flashinfer_moe_backend.value} kernels"

1347 - " for ModelOptNvFp4FusedMoE."

1332 + self.nvfp4_backend = select_nvfp4_moe_backend()

1333 + # TODO: move this type of check into the oracle.

1334 + if (

1335 + not self.moe.is_act_and_mul

1336 + and not self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_CUTLASS

1337 + ):

1338 + raise NotImplementedError(

1339 + "Non-gated activations are only supported by FlashInfer "

1340 + "CUTLASS NvFP4 MoE backend."

1341 )

1349 - elif self.use_marlin:

1350 - logger.info_once("Using Marlin for ModelOptNvFp4FusedMoE.")

1351 - else:

1352 - logger.info_once("Using Cutlass for ModelOptNvFp4FusedMoE.")

1342 +

1343 + self.use_global_sf = is_global_sf_supported_for_nvfp4_backend(

1344 + self.nvfp4_backend

1345 + )

1346 + self.kernel: mk.FusedMoEModularKernel | None = None

1347

1348 def maybe_make_prepare_finalize(

1349 self,

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

1351 ) -> mk.FusedMoEPrepareAndFinalize | None:

1358 - if self.use_marlin or (

1359 - self.allow_flashinfer

1360 - and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

1361 - ):

1352 + UNSUPPORTED = [NvFp4MoeBackend.MARLIN, NvFp4MoeBackend.FLASHINFER_TRTLLM]

1353 + if self.nvfp4_backend in UNSUPPORTED:

1354 return None

1363 - elif (

1364 - self.allow_flashinfer

1365 - and self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS

1366 - ):

1355 + elif self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_CUTLASS:

1356 # TP case: avoid convert to ModularKernelMethod - to be refactored.

1357 if self.moe.dp_size == 1:

1358 return None

@@ -1385,7 +1374,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1374 experts = select_nvfp4_gemm_impl(

1375 self.moe,

1376 self.moe_quant_config,

1388 - allow_flashinfer=self.allow_flashinfer,

1377 + allow_flashinfer=self.nvfp4_backend in FLASHINFER_NVFP4_MOE_BACKENDS,

1378 )

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

1380 return experts

@@ -1405,11 +1394,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1394 params_dtype: torch.dtype,

1395 **extra_weight_attrs,

1396 ):

1408 - if not self.quant_config.is_checkpoint_nvfp4_serialized:

1409 - raise ValueError(

1410 - "NVFP4 quantization was selected, "

1411 - " dynamic quantization is not supported."

1412 - )

1397 + assert self.quant_config.is_checkpoint_nvfp4_serialized

1398

1399 layer.num_experts = num_experts

1400 layer.params_dtype = params_dtype

@@ -1498,14 +1483,12 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1483 {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}

1484 )

1485

1501 - use_global_sf = self.allow_flashinfer and is_flashinfer_supporting_global_sf(

1502 - self.flashinfer_moe_backend

1486 + global_sf_num_experts = (

1487 + global_num_experts if self.use_global_sf else num_experts

1488 )

1504 - global_scale_num_experts = global_num_experts if use_global_sf else num_experts

1505 -

1489 w13_input_scale = PerTensorScaleParameter(

1490 data=torch.empty(

1508 - global_scale_num_experts,

1491 + global_sf_num_experts,

1492 2 if self.moe.is_act_and_mul else 1,

1493 dtype=torch.float32,

1494 ),

@@ -1514,32 +1497,17 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1497 layer.register_parameter("w13_input_scale", w13_input_scale)

1498

1499 w2_input_scale = PerTensorScaleParameter(

1517 - data=torch.empty(global_scale_num_experts, dtype=torch.float32),

1500 + data=torch.empty(global_sf_num_experts, dtype=torch.float32),

1501 weight_loader=weight_loader,

1502 )

1503 layer.register_parameter("w2_input_scale", w2_input_scale)

1504

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

1523 - # GEMM 1 processing

1524 - gemm1_weight = layer.w13_weight.data

1525 - gemm1_weight_scale = layer.w13_weight_scale.data

1526 -

1527 - if (

1528 - self.allow_flashinfer

1529 - and (

1530 - self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS

1531 - or self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

1532 - )

1533 - and self.moe.is_act_and_mul

1534 - ):

1535 - gemm1_weight, gemm1_weight_scale = reorder_w1w3_to_w3w1(

1536 - gemm1_weight, gemm1_weight_scale, dim=-2

1537 - )

1538 -

1539 - layer.w13_weight = Parameter(gemm1_weight, requires_grad=False)

1540 - layer.w13_weight_scale = Parameter(gemm1_weight_scale, requires_grad=False)

1506 + """

1507 + Convert NVFP4 MoE weights into kernel format and setup the kernel.

1508 + """

1509

1542 - # Common processing for w13_weight_scale_2

1510 + # Use a single gscale for w13.

1511 if self.moe.is_act_and_mul and not torch.allclose(

1512 layer.w13_weight_scale_2[:, 0], layer.w13_weight_scale_2[:, 1]

1513 ):

@@ -1547,136 +1515,47 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1515 "w1_weight_scale_2 must match w3_weight_scale_2. "

1516 "Accuracy may be affected."

1517 )

1550 -

1518 w13_weight_scale_2 = layer.w13_weight_scale_2[:, 0].contiguous()

1552 - layer.w13_weight_scale_2 = Parameter(w13_weight_scale_2, requires_grad=False)

1519

1554 - # Common processing for input scales and alphas

1555 - use_global_sf = self.allow_flashinfer and is_flashinfer_supporting_global_sf(

1556 - self.flashinfer_moe_backend

1557 - )

1558 - if use_global_sf:

1559 - # For backends provide by Flashinfer, the input global scales are

1560 - # shared across all experts.

1561 - w13_input_scale = (

1562 - layer.w13_input_scale.max().to(torch.float32).expand(layer.num_experts)

1563 - )

1564 - else:

1565 - w13_input_scale = layer.w13_input_scale.max(dim=1).values.to(torch.float32)

1566 - layer.g1_alphas = Parameter(

1567 - (w13_input_scale * w13_weight_scale_2).to(torch.float32),

1568 - requires_grad=False,

1569 - )

1570 -

1571 - # This is for quantization, so we need to invert it.

1572 - layer.w13_input_scale_quant = Parameter(

1573 - (1 / w13_input_scale).to(torch.float32), requires_grad=False

1574 - )

1575 -

1576 - # GEMM 2 processing

1577 - if use_global_sf:

1578 - # For backends provide by Flashinfer, the input global scales are

1579 - # shared across all experts.

1580 - w2_input_scale = (

1581 - layer.w2_input_scale.max().to(torch.float32).expand(layer.num_experts)

1582 - )

1583 - else:

1584 - w2_input_scale = layer.w2_input_scale

1585 - layer.g2_alphas = Parameter(

1586 - (w2_input_scale * layer.w2_weight_scale_2).to(torch.float32),

1587 - requires_grad=False,

1588 - )

1589 -

1590 - # This is for quantization, so we need to invert it.

1591 - layer.w2_input_scale_quant = Parameter(

1592 - (1 / w2_input_scale).to(torch.float32), requires_grad=False

1520 + (

1521 + w13,

1522 + w13_scale,

1523 + w13_scale_2,

1524 + a13_scale,

1525 + w2,

1526 + w2_scale,

1527 + w2_scale_2,

1528 + a2_scale,

1529 + ) = convert_to_nvfp4_moe_kernel_format(

1530 + nvfp4_backend=self.nvfp4_backend,

1531 + layer=layer,

1532 + w13=layer.w13_weight,

1533 + w13_scale=layer.w13_weight_scale,

1534 + w13_scale_2=w13_weight_scale_2,

1535 + a13_scale=layer.w13_input_scale,

1536 + w2=layer.w2_weight,

1537 + w2_scale=layer.w2_weight_scale,

1538 + w2_scale_2=layer.w2_weight_scale_2,

1539 + a2_scale=layer.w2_input_scale,

1540 + is_act_and_mul=self.moe.is_act_and_mul,

1541 )

1542

1595 - # TensorRT-LLM specific processing

1596 - if (

1597 - self.allow_flashinfer

1598 - and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

1599 - ):

1600 - # Prepare static weights for TRT-LLM kernel

1601 - # alternate: prepare_static_weight_layouts_for_trtllm_moe

1602 - (

1603 - gemm1_weights_fp4_shuffled,

1604 - gemm1_scales_fp4_shuffled,

1605 - gemm2_weights_fp4_shuffled,

1606 - gemm2_scales_fp4_shuffled,

1607 - ) = prepare_static_weights_for_trtllm_fp4_moe(

1608 - layer.w13_weight,

1609 - layer.w2_weight,

1610 - layer.w13_weight_scale,

1611 - layer.w2_weight_scale,

1612 - layer.w2_weight.size(-2), # hidden_size

1613 - layer.w13_weight.size(-2) // 2, # intermediate_size

1614 - layer.w13_weight.size(0), # num_experts

1615 - )

1616 - logger.debug_once("Finished shuffling weights for TRT-LLM MOE")

1617 -

1618 - layer.w13_weight = Parameter(

1619 - gemm1_weights_fp4_shuffled, requires_grad=False

1620 - )

1621 - layer.w2_weight = Parameter(gemm2_weights_fp4_shuffled, requires_grad=False)

1622 - layer.w13_weight_scale = Parameter(

1623 - gemm1_scales_fp4_shuffled, requires_grad=False

1624 - )

1625 - layer.w2_weight_scale = Parameter(

1626 - gemm2_scales_fp4_shuffled, requires_grad=False

1627 - )

1628 -

1629 - # Additional parameter needed for TRT-LLM

1630 - layer.g1_scale_c = Parameter(

1631 - (layer.w2_input_scale_quant * layer.g1_alphas).to(torch.float32),

1632 - requires_grad=False,

1633 - )

1634 - elif self.use_marlin:

1635 - # Marlin processing

1636 - prepare_moe_fp4_layer_for_marlin(layer)

1637 - del layer.g1_alphas

1638 - del layer.g2_alphas

1639 - del layer.w13_input_scale_quant

1640 - del layer.w2_input_scale_quant

1641 - else:

1642 - # Non-TRT-LLM processing (Cutlass or non-flashinfer)

1643 - w13_blockscale_swizzled = swizzle_blockscale(layer.w13_weight_scale)

1644 - layer.w13_weight_scale = Parameter(

1645 - w13_blockscale_swizzled, requires_grad=False

1646 - )

1647 -

1648 - w13_weight = layer.w13_weight

1649 - intermediate_size_pad = w13_blockscale_swizzled.size(1) - w13_weight.size(1)

1650 - if intermediate_size_pad:

1651 - # padding gated activations will require to split w1 and w3

1652 - # and pad them individually

1653 - assert not self.moe.is_act_and_mul, (

1654 - "The intermediate size required padding, "

1655 - "but padding is not implemented for gated activations"

1656 - )

1657 -

1658 - layer.w13_weight = Parameter(

1659 - torch.nn.functional.pad(

1660 - w13_weight, (0, 0, 0, intermediate_size_pad)

1661 - ),

1662 - requires_grad=False,

1663 - )

1664 - layer.w2_weight = Parameter(

1665 - torch.nn.functional.pad(

1666 - layer.w2_weight, (0, intermediate_size_pad // 2, 0, 0)

1667 - ),

1668 - requires_grad=False,

1669 - )

1670 - layer.w2_weight_scale = Parameter(

1671 - torch.nn.functional.pad(

1672 - layer.w2_weight_scale, (0, intermediate_size_pad // 16)

1673 - ),

1674 - requires_grad=False,

1675 - )

1543 + replace_parameter(layer, "w13_weight", w13)

1544 + replace_parameter(layer, "w13_weight_scale", w13_scale)

1545 + replace_parameter(layer, "w13_weight_scale_2", w13_scale_2)

1546 + replace_parameter(layer, "w13_input_scale", a13_scale)

1547 + replace_parameter(layer, "w2_weight", w2)

1548 + replace_parameter(layer, "w2_weight_scale", w2_scale)

1549 + replace_parameter(layer, "w2_weight_scale_2", w2_scale_2)

1550 + replace_parameter(layer, "w2_input_scale", a2_scale)

1551

1677 - w2_blockscale_swizzled = swizzle_blockscale(layer.w2_weight_scale)

1678 - layer.w2_weight_scale = Parameter(

1679 - w2_blockscale_swizzled, requires_grad=False

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

1553 + use_dp = self.moe.dp_size > 1

1554 + if self.moe_quant_config is not None and not use_dp:

1555 + self.kernel = make_nvfp4_moe_kernel(

1556 + backend=self.nvfp4_backend,

1557 + quant_config=self.moe_quant_config,

1558 + moe_config=self.moe,

1559 )

1560

1561 def prepare_dp_allgather_tensor(

@@ -1688,7 +1567,8 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1567 """Optionally prepare extra tensors to carry through DP allgather/EP."""

1568 import flashinfer

1569

1691 - a1_gscale = layer.w13_input_scale_quant

1570 + assert self.moe_quant_config is not None

1571 + a1_gscale = self.moe_quant_config.a1_gscale

1572 hidden_states_fp4, hidden_states_sf = flashinfer.fp4_quantize(

1573 hidden_states,

1574 a1_gscale,

@@ -1700,19 +1580,14 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1580 def get_fused_moe_quant_config(

1581 self, layer: torch.nn.Module

1582 ) -> FusedMoEQuantConfig | None:

1703 - if (

1704 - self.use_marlin

1705 - or self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

1706 - ):

1707 - return None

1708 -

1709 - return nvfp4_moe_quant_config(

1710 - w1_scale=layer.w13_weight_scale,

1583 + return make_nvfp4_moe_quant_config(

1584 + backend=self.nvfp4_backend,

1585 + w13_scale=layer.w13_weight_scale,

1586 w2_scale=layer.w2_weight_scale,

1712 - g1_alphas=layer.g1_alphas,

1713 - g2_alphas=layer.g2_alphas,

1714 - a1_gscale=layer.w13_input_scale_quant,

1715 - a2_gscale=layer.w2_input_scale_quant,

1587 + w13_scale_2=layer.w13_weight_scale_2,

1588 + w2_scale_2=layer.w2_weight_scale_2,

1589 + a13_scale=layer.w13_input_scale,

1590 + a2_scale=layer.w2_input_scale,

1591 )

1592

1593 @property

@@ -1725,18 +1600,8 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1600 x: torch.Tensor,

1601 router_logits: torch.Tensor,

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

1728 - if not self.moe.is_act_and_mul:

1729 - assert (

1730 - self.allow_flashinfer

1731 - and self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS

1732 - ), (

1733 - "Non-gated activations are only supported by the"

1734 - " flashinfer CUTLASS backend for modelopt checkpoints"

1735 - )

1736 -

1603 if (

1738 - self.allow_flashinfer

1739 - and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

1604 + self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM

1605 and not layer.enable_eplb

1606 ):

1607 return flashinfer_trtllm_fp4_moe(

@@ -1762,10 +1627,8 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1627 )

1628

1629 # EPLB path

1765 - if (

1766 - self.allow_flashinfer

1767 - and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM

1768 - ):

1630 + if self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM:

1631 + assert layer.enable_eplb

1632 return flashinfer_trtllm_fp4_routed_moe(

1633 layer=layer,

1634 x=x,

@@ -1774,81 +1637,20 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1637 top_k=layer.top_k,

1638 global_num_experts=layer.global_num_experts,

1639 )

1777 -

1778 - if self.use_marlin:

1779 - return fused_marlin_moe(

1640 + else:

1641 + assert self.kernel is not None

1642 + return self.kernel(

1643 x,

1644 layer.w13_weight,

1645 layer.w2_weight,

1783 - None,

1784 - None,

1785 - layer.w13_weight_scale,

1786 - layer.w2_weight_scale,

1787 - router_logits,

1646 topk_weights,

1647 topk_ids,

1790 - global_scale1=layer.w13_weight_scale_2,

1791 - global_scale2=layer.w2_weight_scale_2,

1792 - quant_type_id=scalar_types.float4_e2m1f.id,

1793 - apply_router_weight_on_input=layer.apply_router_weight_on_input,

1794 - global_num_experts=layer.global_num_experts,

1795 - expert_map=layer.expert_map,

1796 - input_dtype=self.marlin_input_dtype,

1797 - )

1798 -

1799 - elif self.allow_flashinfer:

1800 - assert self.flashinfer_moe_backend in (

1801 - FlashinferMoeBackend.CUTLASS,

1802 - FlashinferMoeBackend.CUTEDSL,

1803 - )

1804 - if self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:

1805 - from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( # noqa: E501

1806 - flashinfer_cutlass_moe_fp4,

1807 - )

1808 -

1809 - flashinfer_fn_moe_fp4 = flashinfer_cutlass_moe_fp4

1810 - else:

1811 - from vllm.model_executor.layers.fused_moe.flashinfer_cutedsl_moe import ( # noqa: E501

1812 - flashinfer_cutedsl_moe_fp4,

1813 - )

1814 -

1815 - flashinfer_fn_moe_fp4 = flashinfer_cutedsl_moe_fp4

1816 -

1817 - assert self.moe_quant_config is not None

1818 - return flashinfer_fn_moe_fp4(

1819 - hidden_states=x,

1820 - w1=layer.w13_weight,

1821 - w2=layer.w2_weight,

1822 - topk_weights=topk_weights,

1823 - topk_ids=topk_ids,

1824 - quant_config=self.moe_quant_config,

1648 inplace=False,

1649 activation=layer.activation,

1650 global_num_experts=layer.global_num_experts,

1651 expert_map=layer.expert_map,

1652 apply_router_weight_on_input=layer.apply_router_weight_on_input,

1653 )

1831 - else:

1832 - # If no modular kernel is provided, use cutlass_moe_fp4 for TP case

1833 - # only (no EP).

1834 - from vllm.model_executor.layers.fused_moe.cutlass_moe import cutlass_moe_fp4

1835 -

1836 - assert self.moe_quant_config is not None

1837 - return cutlass_moe_fp4(

1838 - a=x,

1839 - w1_fp4=layer.w13_weight,

1840 - w2_fp4=layer.w2_weight,

1841 - topk_weights=topk_weights,

1842 - topk_ids=topk_ids,

1843 - quant_config=self.moe_quant_config,

1844 - expert_map=layer.expert_map,

1845 - apply_router_weight_on_input=layer.apply_router_weight_on_input,

1846 - # TODO: derive from arguments

1847 - m=x.shape[0],

1848 - n=layer.w2_weight.shape[2] * 2,

1849 - k=x.shape[1],

1850 - e=layer.w13_weight.shape[0],

1851 - )

1654

1655

1656 ModelOptNvFp4Config.LinearMethodCls = ModelOptNvFp4LinearMethod

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

+106 -4

@@ -2,10 +2,13 @@

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

3 """Utility helpers for NVFP4 + FlashInfer fused-MoE path"""

5 +from typing import TYPE_CHECKING

6 +

7 import torch

9 import vllm.envs as envs

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

11 +from vllm.logger import init_logger

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

13 FusedMoEConfig,

14 FusedMoEQuantConfig,

@@ -20,12 +23,23 @@ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (

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

24 create_flashinfer_prepare_finalize,

25 )

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

27 + swizzle_blockscale,

28 +)

29 from vllm.platforms import current_platform

30 from vllm.utils.flashinfer import (

31 has_flashinfer_cutedsl_grouped_gemm_nt_masked,

32 has_flashinfer_cutlass_fused_moe,

33 )

35 +if TYPE_CHECKING:

36 + from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (

37 + NvFp4MoeBackend,

38 + )

39 +

40 +logger = init_logger(__name__)

41 +

42 +

43 __all__ = [

44 "is_flashinfer_fp4_cutlass_moe_available",

45 "is_flashinfer_fp4_cutedsl_moe_available",

@@ -273,10 +287,9 @@ def flashinfer_trtllm_fp4_moe(

287 hidden_states_fp4, hidden_states_scale_linear_fp4 = x

288 else:

289 # hidden_states is the already quantized

276 - a1_gscale = layer.w13_input_scale_quant

290 (hidden_states_fp4, hidden_states_scale_linear_fp4) = flashinfer.fp4_quantize(

291 x,

279 - a1_gscale,

292 + layer.a1_gscale,

293 is_sf_swizzled_layout=False,

294 )

295

@@ -369,10 +382,9 @@ def flashinfer_trtllm_fp4_routed_moe(

382 hidden_states_fp4, hidden_states_scale_linear_fp4 = x

383 else:

384 # Quantize input to FP4

372 - a1_gscale = layer.w13_input_scale_quant

385 (hidden_states_fp4, hidden_states_scale_linear_fp4) = flashinfer.fp4_quantize(

386 x,

375 - a1_gscale,

387 + layer.a1_gscale,

388 is_sf_swizzled_layout=False,

389 )

390

@@ -410,3 +422,93 @@ def flashinfer_trtllm_fp4_routed_moe(

422 )[0]

423

424 return out

425 +

426 +

427 +def prepare_nvfp4_moe_layer_for_fi_or_cutlass(

428 + backend: "NvFp4MoeBackend",

429 + layer: torch.nn.Module,

430 + w13: torch.Tensor,

431 + w13_scale: torch.Tensor,

432 + w13_scale_2: torch.Tensor,

433 + a13_scale: torch.Tensor,

434 + w2: torch.Tensor,

435 + w2_scale: torch.Tensor,

436 + w2_scale_2: torch.Tensor,

437 + a2_scale: torch.Tensor,

438 + is_act_and_mul: bool,

439 +) -> tuple[

440 + torch.Tensor,

441 + torch.Tensor,

442 + torch.Tensor,

443 + torch.Tensor,

444 + torch.Tensor,

445 + torch.Tensor,

446 + torch.Tensor,

447 + torch.Tensor,

448 +]:

449 + # Delayed import for circular dependency avoidance.

450 + from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (

451 + NvFp4MoeBackend,

452 + is_global_sf_supported_for_nvfp4_backend,

453 + )

454 +

455 + assert backend in [

456 + NvFp4MoeBackend.VLLM_CUTLASS,

457 + NvFp4MoeBackend.FLASHINFER_CUTLASS,

458 + NvFp4MoeBackend.FLASHINFER_TRTLLM,

459 + NvFp4MoeBackend.FLASHINFER_TRTLLM,

460 + ]

461 +

462 + # Reorder [w1, w3] to [w3, w1] for FI NVFP4 MoE kernels.

463 + if is_act_and_mul and backend in [

464 + NvFp4MoeBackend.FLASHINFER_CUTLASS,

465 + NvFp4MoeBackend.FLASHINFER_TRTLLM,

466 + ]:

467 + w13, w13_scale = reorder_w1w3_to_w3w1(w13, w13_scale)

468 +

469 + # For some FI kernels, the input scales are shared by all experts.

470 + if is_global_sf_supported_for_nvfp4_backend(backend):

471 + num_experts = w13.shape[0]

472 + a13_scale = a13_scale.max().to(torch.float32).expand(num_experts)

473 + a2_scale = a2_scale.max().to(torch.float32).expand(num_experts)

474 + else:

475 + a13_scale = a13_scale.max(dim=1).values.to(torch.float32)

476 +

477 + # Shuffle weights and scales for FI TRTLLM NVFP4 MoE kernels.

478 + if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM:

479 + w13, w13_scale, w2, w2_scale = prepare_static_weights_for_trtllm_fp4_moe(

480 + w13,

481 + w2,

482 + w13_scale,

483 + w2_scale,

484 + w2.size(-2), # hidden_size

485 + w13.size(-2) // 2, # intermediate_size

486 + w13.size(0), # num_experts

487 + )

488 +

489 + # We do not need to make this a parameter, because

490 + # it is not used during the weight (re)-loading process.

491 + layer.g1_scale_c = a13_scale * w13_scale_2 / a2_scale

492 + layer.a1_gscale = 1.0 / a13_scale

493 + layer.g1_alphas = a13_scale * w13_scale_2

494 + layer.g2_alphas = a2_scale * w2_scale_2

495 + else:

496 + # Swizzle the block scales for other FI NVFP4 MoE kernels.

497 + w13_scale = swizzle_blockscale(w13_scale)

498 +

499 + # Apply padding if needed.

500 + pad_size = w13_scale.size(1) - w13.size(1)

501 + if pad_size > 0:

502 + if is_act_and_mul:

503 + raise NotImplementedError(

504 + "Intermediate size padding for w1 and w3, for %s "

505 + "NvFp4 backend, but this is not currently supported",

506 + backend.value,

507 + )

508 + w13 = torch.nn.functional.pad(w13, (0, 0, 0, pad_size))

509 + w2 = torch.nn.functional.pad(w2, (0, pad_size // 2, 0, 0))

510 + w2_scale = torch.nn.functional.pad(w2_scale, (0, pad_size // 16))

511 +

512 + w2_scale = swizzle_blockscale(w2_scale)

513 +

514 + return w13, w13_scale, w13_scale_2, a13_scale, w2, w2_scale, w2_scale_2, a2_scale

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

+101

@@ -8,6 +8,7 @@ import vllm._custom_ops as ops

8 from vllm.logger import init_logger

9 from vllm.model_executor.layers.quantization.utils.marlin_utils import (

10 USE_FP32_REDUCE_DEFAULT,

11 + get_marlin_input_dtype,

12 marlin_make_workspace_new,

13 marlin_permute_bias,

14 marlin_permute_scales,

@@ -226,6 +227,106 @@ def prepare_fp4_layer_for_marlin(

227 return

228

229

230 +def prepare_nvfp4_moe_layer_for_marlin(

231 + layer: torch.nn.Module,

232 + w13: torch.Tensor,

233 + w13_scale: torch.Tensor,

234 + w13_scale_2: torch.Tensor,

235 + w2: torch.Tensor,

236 + w2_scale: torch.Tensor,

237 + w2_scale_2: torch.Tensor,

238 +) -> tuple[

239 + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor

240 +]:

241 + logger.warning_once(

242 + "Your GPU does not have native support for FP4 computation but "

243 + "FP4 quantization is being used. Weight-only FP4 compression will "

244 + "be used leveraging the Marlin kernel. This may degrade "

245 + "performance for compute-heavy workloads."

246 + )

247 +

248 + input_dtype = get_marlin_input_dtype(prefix="")

249 + if input_dtype is not None and input_dtype.itemsize == 1:

250 + raise RuntimeError("NVFP4 weight + INT8/FP8 activation is not supported.")

251 +

252 + GROUP_SIZE = 16

253 + E = layer.num_experts

254 + K = layer.hidden_size

255 + N = layer.intermediate_size_per_partition

256 +

257 + device = w13.device

258 + param_dtype = layer.params_dtype

259 + is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1

260 +

261 + # WORKSPACE

262 + layer.workspace = marlin_make_workspace_new(device, 4)

263 + perm = torch.empty(0, dtype=torch.int, device=device)

264 +

265 + # WEIGHT

266 + # Repack weights to marlin format

267 + def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor:

268 + tensor_list = []

269 + if "w13" in name:

270 + size_n, size_k = N * 2, K

271 + else:

272 + size_n, size_k = K, N

273 +

274 + assert weight.shape == (E, size_n, size_k // 2)

275 +

276 + for i in range(E):

277 + qweight = weight[i].view(torch.int32).T.contiguous()

278 +

279 + marlin_qweight = ops.gptq_marlin_repack(

280 + b_q_weight=qweight,

281 + perm=perm,

282 + size_k=size_k,

283 + size_n=size_n,

284 + num_bits=4,

285 + is_a_8bit=is_a_8bit,

286 + )

287 + tensor_list.append(marlin_qweight)

288 +

289 + return torch.cat([x.unsqueeze(0) for x in tensor_list], 0)

290 +

291 + w13 = repack_weight(w13, "w13")

292 + w2 = repack_weight(w2, "w2")

293 +

294 + # WEIGHT SCALES

295 + # Permute scales

296 + def premute_scales(

297 + scales: torch.Tensor, g_scales: torch.Tensor, name: str

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

299 + scales = scales.to(param_dtype)

300 + g_scales = g_scales.to(param_dtype)

301 +

302 + tensor_list = []

303 + if "w13" in name:

304 + size_n, size_k = N * 2, K

305 + else:

306 + size_n, size_k = K, N

307 +

308 + for i in range(E):

309 + scale = scales[i].T

310 + marlin_scales = marlin_permute_scales(

311 + s=scale,

312 + size_k=size_k,

313 + size_n=size_n,

314 + group_size=GROUP_SIZE,

315 + is_a_8bit=is_a_8bit,

316 + )

317 + marlin_scales = nvfp4_marlin_process_scales(marlin_scales)

318 + tensor_list.append(marlin_scales)

319 +

320 + scales = torch.cat([x.unsqueeze(0) for x in tensor_list], 0)

321 + g_scales = nvfp4_marlin_process_global_scale(g_scales)

322 + return scales, g_scales

323 +

324 + w13_scale, w13_scale_2 = premute_scales(w13_scale, w13_scale_2, "w13")

325 + w2_scale, w2_scale_2 = premute_scales(w2_scale, w2_scale_2, "w2")

326 +

327 + return w13, w13_scale, w13_scale_2, w2, w2_scale, w2_scale_2

328 +

329 +

330 def prepare_moe_fp4_layer_for_marlin(

331 layer: torch.nn.Module, input_dtype: torch.dtype | None = None

332 ) -> None: