@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 42135d6898

Fork 0

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

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

Synced Nov 9, 21:48 UTC

[MoE Refactor] Oracle Select FP8+NVFP4 Kernels In Priority (#32414)

Robert Shaw committed Jan 21, 2026 at 08:22 UTC 42135d689830c0e764d925b6454bc68ba6c6cab4

82 files changed +2699 -1552

.buildkite/test-pipeline.yaml

+40

@@ -634,6 +634,46 @@ steps:

634 - pip install helion

635 - pytest -v -s kernels/helion/

636

637 +

638 +- label: Kernels FP8 MoE Test (1 H100)

639 + timeout_in_minutes: 90

640 + gpu: h100

641 + num_gpus: 1

642 + optional: true

643 + commands:

644 + - pytest -v -s kernels/moe/test_cutlass_moe.py

645 + - pytest -v -s kernels/moe/test_flashinfer.py

646 + - pytest -v -s kernels/moe/test_gpt_oss_triton_kernels.py

647 + - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py

648 + - pytest -v -s kernels/moe/test_moe.py

649 + # - pytest -v -s kernels/moe/test_block_fp8.py - failing on main

650 + - pytest -v -s kernels/moe/test_block_int8.py

651 + - pytest -v -s kernels/moe/test_triton_moe_no_act_mul.py

652 + - pytest -v -s kernels/moe/test_triton_moe_ptpc_fp8.py

653 +

654 +- label: Kernels FP8 MoE Test (2 H100s)

655 + timeout_in_minutes: 90

656 + gpu: h100

657 + num_gpus: 2

658 + optional: true

659 + commands:

660 + - pytest -v -s kernels/moe/test_deepep_deepgemm_moe.py

661 + - pytest -v -s kernels/moe/test_deepep_moe.py

662 + - pytest -v -s kernels/moe/test_pplx_cutlass_moe.py

663 + # - pytest -v -s kernels/moe/test_pplx_moe.py - failing on main

664 +

665 +- label: Kernels Fp4 MoE Test (B200)

666 + timeout_in_minutes: 60

667 + gpu: b200

668 + num_gpus: 1

669 + optional: true

670 + commands:

671 + - pytest -v -s kernels/moe/test_cutedsl_moe.py

672 + - pytest -v -s kernels/moe/test_flashinfer_moe.py

673 + - pytest -v -s kernels/moe/test_nvfp4_moe.py

674 + - pytest -v -s kernels/moe/test_ocp_mx_moe.py

675 +

676 +

677 - label: Model Executor Test # 23min

678 timeout_in_minutes: 35

679 torch_nightly: true

benchmarks/kernels/benchmark_cutlass_moe_fp8.py

+7 -5

@@ -9,6 +9,7 @@ but use different quantization strategies and backends.

9 import torch

11 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.config import fp8_w8a8_moe_quant_config

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

@@ -138,12 +139,13 @@ def bench_run(

139 fn = mk.FusedMoEModularKernel(

140 MoEPrepareAndFinalizeNoEP(),

141 CutlassExpertsFp8(

141 - out_dtype=a.dtype,

142 - e=num_experts,

143 - n=n,

144 - k=k,

142 + moe_config=make_dummy_moe_config(

143 + num_experts=num_experts,

144 + hidden_dim=k,

145 + intermediate_size_per_partition=n,

146 + in_dtype=a.dtype,

147 + ),

148 quant_config=quant_config,

146 - device=w1.device,

149 ),

150 )

151

benchmarks/kernels/benchmark_cutlass_moe_nvfp4.py

+3 -4

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

12 import torch.utils.benchmark as benchmark

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

@@ -198,8 +199,7 @@ def bench_run(

199 kernel = mk.FusedMoEModularKernel(

200 MoEPrepareAndFinalizeNoEP(defer_input_quant=True),

201 CutlassExpertsFp4(

201 - out_dtype=dtype,

202 - max_experts_per_worker=e,

202 + make_dummy_moe_config(),

203 quant_config=quant_config,

204 ),

205 )

@@ -244,8 +244,7 @@ def bench_run(

244 kernel = mk.FusedMoEModularKernel(

245 MoEPrepareAndFinalizeNoEP(defer_input_quant=True),

246 CutlassExpertsFp4(

247 - out_dtype=dtype,

248 - max_experts_per_worker=e,

247 + make_dummy_moe_config(),

248 quant_config=quant_config,

249 ),

250 )

benchmarks/kernels/benchmark_grouped_gemm_cutlass.py

+13 -12

@@ -6,6 +6,7 @@ import torch.utils.benchmark as benchmark

6 from benchmark_shapes import WEIGHT_SHAPES_MOE

8 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.config import fp8_w8a8_moe_quant_config

@@ -134,13 +135,13 @@ def bench_run(

135 fn = mk.FusedMoEModularKernel(

136 MoEPrepareAndFinalizeNoEP(),

137 CutlassExpertsFp8(

137 - out_dtype=a.dtype,

138 - # NOTE(rob): w2 is shaped as [E, hidden, intermediate]

139 - e=w2.shape[0],

140 - n=w2.shape[2],

141 - k=w2.shape[1],

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

144 quant_config=quant_config,

143 - device=w1.device,

145 ),

146 )

147

@@ -166,13 +167,13 @@ def bench_run(

167 fn = mk.FusedMoEModularKernel(

168 MoEPrepareAndFinalizeNoEP(),

169 CutlassExpertsFp8(

169 - out_dtype=a.dtype,

170 - # NOTE(rob): w2 is shaped as [E, hidden, intermediate]

171 - e=w2.shape[0],

172 - n=w2.shape[2],

173 - k=w2.shape[1],

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

176 quant_config=quant_config,

175 - device=w1.device,

177 ),

178 )

179

benchmarks/kernels/benchmark_moe.py

+29 -1

@@ -16,10 +16,16 @@ import torch

16 from ray.experimental.tqdm_ray import tqdm

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

19 + FusedMoEConfig,

20 + FusedMoEParallelConfig,

21 FusedMoEQuantConfig,

22 + RoutingMethodType,

23 _get_config_dtype_str,

24 )

25 from vllm.model_executor.layers.fused_moe.fused_moe import *

26 +from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (

27 + TritonOrDeepGemmExperts,

28 +)

29 from vllm.platforms import current_platform

30 from vllm.transformers_utils.config import get_config

31 from vllm.triton_utils import triton

@@ -194,10 +200,33 @@ def benchmark_config(

200 block_shape=block_quant_shape,

201 )

202

203 + deep_gemm_experts = mk.FusedMoEModularKernel(

204 + prepare_finalize=MoEPrepareAndFinalizeNoEP(),

205 + fused_experts=TritonOrDeepGemmExperts(

206 + moe_config=FusedMoEConfig(

207 + num_experts=num_experts,

208 + experts_per_token=topk,

209 + hidden_dim=hidden_size,

210 + intermediate_size_per_partition=shard_intermediate_size,

211 + num_local_experts=num_experts,

212 + activation="silu",

213 + parallel_config=FusedMoEParallelConfig.make_no_parallel(),

214 + in_dtype=init_dtype,

215 + routing_method=RoutingMethodType.TopK,

216 + ),

217 + quant_config=quant_config,

218 + ),

219 + )

220 +

221 with override_config(config):

222 topk_weights, topk_ids, token_expert_indices = fused_topk(

223 x, input_gating, topk, renormalize=not use_deep_gemm

224 )

225 +

226 + if use_deep_gemm:

227 + return deep_gemm_experts(

228 + x, w1, w2, topk_weights, topk_ids, inplace=True

229 + )

230 return fused_experts(

231 x,

232 w1,

@@ -206,7 +235,6 @@ def benchmark_config(

235 topk_ids,

236 inplace=True,

237 quant_config=quant_config,

209 - allow_deep_gemm=use_deep_gemm,

238 )

239

240 # JIT compilation & warmup

docs/design/moe_kernel_features.md

+2 -2

@@ -85,10 +85,10 @@ To be used with a particular `FusedMoEPrepareAndFinalize` subclass, MoE kernels

85 |--------|-------------------|--------------|---------------|---------------------|-----------------------|---------|--------|

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

88 +| deep gemm | standard,batched | fp8 | G(128),A,T | silu, gelu | 6 | Y | [`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 | [`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] |

91 +| flashinfer | standard | nvfp4,fp8 | T | 5 | N | Y | [`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] |

93 | marlin | standard,batched | 3 / N/A | 3 / N/A | silu,swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe],[`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] |

94 | trtllm | standard | mxfp4,nvfp4 | G(16),G(32) | 5 | N | Y | [`TrtLlmGenExperts`][vllm.model_executor.layers.fused_moe.trtllm_moe.TrtLlmGenExperts] |

tests/compile/test_fusion_attn.py

+3 -3

@@ -43,7 +43,7 @@ from vllm.forward_context import get_forward_context, set_forward_context

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

44 QuantKey,

45 kFp8StaticTensorSym,

46 - kNvfp4Quant,

46 + kNvfp4Dynamic,

47 )

48 from vllm.platforms import current_platform

49 from vllm.utils.flashinfer import has_flashinfer

@@ -215,7 +215,7 @@ class TestAttentionFp8StaticQuantPatternModel(AttentionQuantPatternModel):

215 class TestAttentionNvfp4QuantPatternModel(AttentionQuantPatternModel):

216 """Test model for AttentionNvfp4QuantPattern fusion."""

217

218 - quant_key = kNvfp4Quant

218 + quant_key = kNvfp4Dynamic

219

220 def __init__(self, *args, **kwargs):

221 super().__init__(*args, **kwargs)

@@ -468,7 +468,7 @@ def test_attention_quant_pattern(

468

469 # Note: for fp8, fully_replaced=False because query quant ops remain in graph.

470 # Only output quant ops are fused into attention.

471 - test_backend.check_before_ops([quant_op], fully_replaced=quant_key is kNvfp4Quant)

471 + test_backend.check_before_ops([quant_op], fully_replaced=quant_key is kNvfp4Dynamic)

472

473 # access the underlying `AttnFusionPass` on the `LazyInitPass`

474 assert attn_pass.pass_.matched_count == sum(attn_fusion_supported)

tests/compile/test_silu_mul_quant_fusion.py

+3 -3

@@ -44,7 +44,7 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import W8A8BlockFp8

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

45 GroupShape,

46 kFp8StaticTensorSym,

47 - kNvfp4Quant,

47 + kNvfp4Dynamic,

48 )

49 from vllm.platforms import current_platform

@@ -134,11 +134,11 @@ class TestSiluMulNvfp4QuantModel(torch.nn.Module):

134 def ops_in_model_before(self):

135 return [

136 SILU_MUL_OP if self.enable_silu_mul_custom_op else torch.ops.aten.mul,

137 - QUANT_OPS[kNvfp4Quant],

137 + QUANT_OPS[kNvfp4Dynamic],

138 ]

139

140 def ops_in_model_after(self):

141 - return [FUSED_OPS[kNvfp4Quant]]

141 + return [FUSED_OPS[kNvfp4Dynamic]]

142

143

144 class TestSiluMulGroupFp8QuantModel(torch.nn.Module):

tests/evals/gsm8k/configs/moe-refactor-dp-ep/Llama-4-Scout-Fp8-ModelOpt-triton.yaml

@@ -3,3 +3,6 @@ accuracy_threshold: 0.92

3 num_questions: 1319

4 num_fewshot: 5

5 server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel"

6 +env:

7 + VLLM_USE_FLASHINFER_MOE_FP8: "0"

8 + VLLM_USE_DEEP_GEMM: "0"

tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-Fp8-AutoFp8-deepgemm-deepep-ll.yaml

-1

@@ -6,4 +6,3 @@ server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enab

6 env:

7 VLLM_USE_DEEP_GEMM: "1"

8 VLLM_USE_DEEP_GEMM_MOE: "1"

9 - VLLM_USE_DEEP_GEMM_E8M0: "0"

tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-Fp8-CT-Block-deepgemm-deepep-ll.yaml

-1

@@ -6,4 +6,3 @@ server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enab

6 env:

7 VLLM_USE_DEEP_GEMM: "1"

8 VLLM_USE_DEEP_GEMM_MOE: "1"

9 - VLLM_USE_DEEP_GEMM_E8M0: "0"

tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-Fp8-ModelOpt-triton.yaml

@@ -3,3 +3,5 @@ accuracy_threshold: 0.92

3 num_questions: 1319

4 num_fewshot: 5

5 server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"

6 +env:

7 + VLLM_USE_FLASHINFER_MOE_FP8: "0"

tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-AutoFp8-fi-cutlass.yaml

-2

@@ -4,7 +4,5 @@ num_questions: 1319

4 num_fewshot: 5

5 server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"

6 env:

7 - VLLM_USE_DEEP_GEMM: "0"

8 - VLLM_USE_DEEP_GEMM_MOE: "0"

7 VLLM_USE_FLASHINFER_MOE_FP8: "1"

8 VLLM_FLASHINFER_MOE_BACKEND: "throughput"

tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-AutoFp8-fi-trtllm.yaml

-2

@@ -4,7 +4,5 @@ num_questions: 1319

4 num_fewshot: 5

5 server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"

6 env:

7 - VLLM_USE_DEEP_GEMM: "0"

8 - VLLM_USE_DEEP_GEMM_MOE: "0"

7 VLLM_USE_FLASHINFER_MOE_FP8: "1"

8 VLLM_FLASHINFER_MOE_BACKEND: "latency"

tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-AutoFp8-marlin.yaml

-2

@@ -4,6 +4,4 @@ num_questions: 1319

4 num_fewshot: 5

5 server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"

6 env:

7 - VLLM_USE_DEEP_GEMM: "0"

8 - VLLM_USE_DEEP_GEMM_MOE: "0"

7 VLLM_TEST_FORCE_FP8_MARLIN: "1"

tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-AutoFp8-triton.yaml

+1 -1

@@ -4,5 +4,5 @@ num_questions: 1319

4 num_fewshot: 5

5 server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"

6 env:

7 + VLLM_USE_FLASHINFER_MOE_FP8: "0"

8 VLLM_USE_DEEP_GEMM: "0"

8 - VLLM_USE_DEEP_GEMM_MOE: "0"

tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-CT-Block-fi-cutlass.yaml

-2

@@ -4,7 +4,5 @@ num_questions: 1319

4 num_fewshot: 5

5 server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"

6 env:

7 - VLLM_USE_DEEP_GEMM: "0"

8 - VLLM_USE_DEEP_GEMM_MOE: "0"

7 VLLM_USE_FLASHINFER_MOE_FP8: "1"

8 VLLM_FLASHINFER_MOE_BACKEND: "throughput"

tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-CT-Block-marlin.yaml

-2

@@ -4,6 +4,4 @@ num_questions: 1319

4 num_fewshot: 5

5 server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"

6 env:

7 - VLLM_USE_DEEP_GEMM: "0"

8 - VLLM_USE_DEEP_GEMM_MOE: "0"

7 VLLM_TEST_FORCE_FP8_MARLIN: "1"

tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-CT-Block-triton.yaml

+1 -1

@@ -4,5 +4,5 @@ num_questions: 1319

4 num_fewshot: 5

5 server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"

6 env:

7 + VLLM_USE_FLASHINFER_MOE_FP8: "0"

8 VLLM_USE_DEEP_GEMM: "0"

8 - VLLM_USE_DEEP_GEMM_MOE: "0"

tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-CT-vllm-cutlass.yaml

@@ -3,3 +3,5 @@ accuracy_threshold: 0.88

3 num_questions: 1319

4 num_fewshot: 5

5 server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"

6 +env:

7 + VLLM_USE_FLASHINFER_MOE_FP4: "0"

tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-ModelOpt-vllm-cutlass.yaml

@@ -3,3 +3,5 @@ accuracy_threshold: 0.88

3 num_questions: 1319

4 num_fewshot: 5

5 server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"

6 +env:

7 + VLLM_USE_FLASHINFER_MOE_FP4: "0"

tests/kernels/moe/modular_kernel_tools/common.py

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

26 FusedMoEConfig,

27 FusedMoEParallelConfig,

28 FusedMoEQuantConfig,

29 + RoutingMethodType,

30 )

31 from vllm.utils.import_utils import has_deep_ep, has_deep_gemm, has_pplx

@@ -574,10 +575,14 @@ def make_modular_kernel(

575 num_experts=config.E,

576 experts_per_token=config.topk,

577 hidden_dim=config.K,

578 + intermediate_size_per_partition=config.N,

579 num_local_experts=config.num_local_experts,

580 moe_parallel_config=moe_parallel_config,

581 in_dtype=config.dtype,

582 max_num_tokens=next_power_of_2(config.M),

583 + activation="silu",

584 + device=vllm_config.device_config.device,

585 + routing_method=RoutingMethodType.DeepSeekV3,

586 )

587

588 # make modular kernel

tests/kernels/moe/modular_kernel_tools/mk_objects.py

+17 -75

@@ -425,84 +425,26 @@ def make_fused_experts(

425 num_dispatchers: int,

426 N: int,

427 ) -> mk.FusedMoEPermuteExpertsUnpermute:

428 - batch_kwargs = {

429 - "max_num_tokens": moe.max_num_tokens,

430 - "num_dispatchers": num_dispatchers,

431 - }

432 - quant_kwargs = {

433 - "quant_config": quant_config,

434 - }

435 - deepgemm_kwargs = {"allow_deep_gemm": has_deep_gemm()}

436 -

437 - torch.set_printoptions(threshold=0, edgeitems=0, linewidth=10000)

438 -

439 - if fused_experts_type == BatchedDeepGemmExperts:

440 - kwargs = batch_kwargs | quant_kwargs

441 - print(f"Making BatchedDeepGemmExperts {kwargs} ...")

442 - experts = BatchedDeepGemmExperts(**kwargs)

443 - elif fused_experts_type == BatchedTritonExperts:

444 - kwargs = batch_kwargs | quant_kwargs

445 - print(f"Making BatchedTritonExperts {kwargs} ...")

446 - experts = BatchedTritonExperts(**kwargs)

447 - elif fused_experts_type == DeepGemmExperts:

448 - print(f"Making DeepGemmExperts {quant_config} ...")

449 - experts = DeepGemmExperts(quant_config)

450 - elif fused_experts_type == TritonExperts:

451 - kwargs = quant_kwargs

452 - print(f"Making TritonExperts {kwargs} ...")

453 - experts = TritonExperts(**kwargs)

454 - elif fused_experts_type == TritonOrDeepGemmExperts:

455 - kwargs = quant_kwargs | deepgemm_kwargs

456 - print(f"Making TritonOrDeepGemmExperts {kwargs} ...")

457 - experts = TritonOrDeepGemmExperts(**kwargs)

458 - elif fused_experts_type == NaiveBatchedExperts:

459 - kwargs = batch_kwargs | quant_kwargs

460 - print(f"Making NaiveBatchedExperts {kwargs} ...")

461 - experts = NaiveBatchedExperts(**kwargs)

462 - elif fused_experts_type == CutlassExpertsFp8:

463 - strides = make_cutlass_strides(moe.num_experts, N, moe.hidden_dim)

464 - kwargs = {

465 - "out_dtype": moe.in_dtype,

466 - "ab_strides1": strides[0],

467 - "ab_strides2": strides[1],

468 - "c_strides1": strides[2],

469 - "c_strides2": strides[3],

470 - } | quant_kwargs

471 - print(f"Making CutlassExpertsFp8 {kwargs} ...")

472 - experts = CutlassExpertsFp8(**kwargs)

473 - elif fused_experts_type == CutlassBatchedExpertsFp8:

474 - strides = make_cutlass_strides(moe.num_experts, N, moe.hidden_dim)

475 - kwargs = {

476 - "max_experts_per_worker": moe.num_local_experts,

477 - "num_dispatchers": num_dispatchers,

478 - "out_dtype": moe.in_dtype,

479 - "ab_strides1": strides[0],

480 - "ab_strides2": strides[1],

481 - "c_strides1": strides[2],

482 - "c_strides2": strides[3],

483 - } | quant_kwargs

484 - print(f"Making CutlassBatchedExpertsFp8 {kwargs} ...")

485 - experts = CutlassBatchedExpertsFp8(**kwargs)

486 - elif fused_experts_type == CutlassExpertsFp4:

428 + if (

429 + fused_experts_type.activation_format()

430 + == mk.FusedMoEActivationFormat.BatchedExperts

431 + ):

432 kwargs = {

488 - "max_experts_per_worker": moe.num_local_experts,

433 + "moe_config": moe,

434 + "quant_config": quant_config,

435 + "max_num_tokens": moe.max_num_tokens,

436 "num_dispatchers": num_dispatchers,

490 - "out_dtype": moe.in_dtype,

491 - } | quant_kwargs

492 - print(f"Making CutlassExpertsFp4 {kwargs} ...")

493 - experts = CutlassExpertsFp4(**kwargs)

494 - elif fused_experts_type == FlashInferExperts:

495 - kwargs = {

496 - "out_dtype": moe.in_dtype,

497 - "ep_rank": moe.ep_rank,

498 - "ep_size": moe.ep_size,

499 - "tp_rank": moe.tp_rank,

500 - "tp_size": moe.tp_size,

501 - } | quant_kwargs

502 - print(f"Making FlashInferExperts {kwargs} ...")

503 - experts = FlashInferExperts(**kwargs)

437 + }

438 else:

505 - raise RuntimeError(f"Unknown fused experts type: {fused_experts_type}")

439 + kwargs = {

440 + "moe_config": moe,

441 + "quant_config": quant_config,

442 + }

443 +

444 + torch.set_printoptions(threshold=0, edgeitems=0, linewidth=10000)

445 +

446 + print(f"Making {fused_experts_type.__class__.__name__} {kwargs} ...")

447 + experts = fused_experts_type(**kwargs)

448

449 torch.set_printoptions(threshold=1000, edgeitems=5, linewidth=80)

450

tests/kernels/moe/test_batched_deepgemm.py

@@ -16,6 +16,7 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularK

16 from vllm.utils.deep_gemm import calc_diff, is_deep_gemm_supported

18 from .test_deepgemm import make_block_quant_fp8_weights

19 +from .utils import make_dummy_moe_config

21 BLOCK_SIZE = [128, 128]

@@ -71,6 +72,7 @@ def test_batched_deepgemm_vs_triton(

72 max_num_tokens=max_num_tokens,

73 num_dispatchers=1,

74 quant_config=quant_config,

75 + moe_config=make_dummy_moe_config(),

76 )

77 mk_triton = FusedMoEModularKernel(prep_finalize, triton_experts)

@@ -89,6 +91,7 @@ def test_batched_deepgemm_vs_triton(

91 max_num_tokens=max_num_tokens,

92 num_dispatchers=1,

93 quant_config=quant_config,

94 + moe_config=make_dummy_moe_config(),

95 )

96 mk_deepgemm = FusedMoEModularKernel(prep_finalize, deepgemm_experts)

tests/kernels/moe/test_block_fp8.py

+39 -3

@@ -4,7 +4,12 @@

4 import pytest

5 import torch

7 -from tests.kernels.moe.utils import make_test_quant_config, make_test_weights

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

8 +from tests.kernels.moe.utils import (

9 + make_dummy_moe_config,

10 + make_test_quant_config,

11 + make_test_weights,

12 +)

13 from tests.kernels.quant_utils import (

14 native_per_token_group_quant_fp8,

15 native_w8a8_block_matmul,

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

20 fused_experts,

21 fused_topk,

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.deep_gemm_moe import (

27 _valid_deep_gemm_shape,

20 - deep_gemm_moe_fp8,

28 )

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

30 modular_triton_fused_moe,

31 )

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

33 + MoEPrepareAndFinalizeNoEP,

34 +)

35 +from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (

36 + TritonOrDeepGemmExperts,

37 +)

38 from vllm.platforms import current_platform

39 from vllm.utils.deep_gemm import (

40 get_mk_alignment_for_contiguous_layout,

@@ -161,7 +174,7 @@ def test_w8a8_block_fp8_fused_moe(

174 block_shape=block_size,

175 )

176

164 - m_fused_moe = modular_triton_fused_moe(quant_config)

177 + m_fused_moe = modular_triton_fused_moe(make_dummy_moe_config(), quant_config)

178

179 topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False)

180

@@ -236,6 +249,29 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed, monkeypatch)

249

250 topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False)

251

252 + quant_config = fp8_w8a8_moe_quant_config(

253 + w1_scale=w1_s,

254 + w2_scale=w2_s,

255 + block_shape=block_size,

256 + )

257 +

258 + deep_gemm_experts = mk.FusedMoEModularKernel(

259 + prepare_finalize=MoEPrepareAndFinalizeNoEP(),

260 + fused_experts=TritonOrDeepGemmExperts(

261 + moe_config=make_dummy_moe_config(),

262 + quant_config=quant_config,

263 + ),

264 + )

265 +

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

267 + return deep_gemm_experts(

268 + hidden_states=a,

269 + w1=w1,

270 + w2=w2,

271 + topk_weights=topk_weights,

272 + topk_ids=topk_ids,

273 + )

274 +

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

276 with set_current_vllm_config(vllm_config):

277 ref_out = torch_w8a8_block_fp8_moe(

tests/kernels/moe/test_cutlass_moe.py

+16 -13

@@ -8,6 +8,7 @@ import pytest

8 import torch

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

11 +from tests.kernels.moe.utils import make_dummy_moe_config

12 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

@@ -193,16 +194,18 @@ def run_with_expert_maps(

194

195 out_tensor = torch.zeros_like(cutlass_moe_kwargs["hidden_states"])

196 for kwargs, new_quant_config in slice_experts():

197 + w2 = kwargs["w2"]

198 + a = kwargs["hidden_states"]

199 kernel = mk.FusedMoEModularKernel(

200 MoEPrepareAndFinalizeNoEP(),

201 CutlassExpertsFp8(

199 - out_dtype=kwargs["hidden_states"].dtype,

200 - # NOTE(rob): w2 is shaped as [E, hidden, intermediate]

201 - e=kwargs["w2"].shape[0], # type: ignore[union-attr]

202 - n=kwargs["w2"].shape[2], # type: ignore[union-attr]

203 - k=kwargs["w2"].shape[1], # type: ignore[union-attr]

202 + moe_config=make_dummy_moe_config(

203 + num_experts=w2.shape[0],

204 + hidden_dim=w2.shape[1],

205 + intermediate_size_per_partition=w2.shape[2],

206 + in_dtype=a.dtype,

207 + ),

208 quant_config=new_quant_config,

205 - device="cuda",

209 ),

210 )

211 out_tensor = out_tensor + kernel(**kwargs)

@@ -249,19 +252,19 @@ def run_8_bit(

252 "topk_ids": topk_ids,

253 }

254

252 - num_experts = moe_tensors.w1.size(0)

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

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

257 if not with_ep:

258 kernel = mk.FusedMoEModularKernel(

259 MoEPrepareAndFinalizeNoEP(),

260 CutlassExpertsFp8(

258 - out_dtype=moe_tensors.a.dtype,

259 - # NOTE(rob): w2 is shaped as [E, hidden, intermediate]

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

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

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

261 + moe_config=make_dummy_moe_config(

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

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

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

265 + in_dtype=moe_tensors.a.dtype,

266 + ),

267 quant_config=quant_config,

264 - device="cuda",

268 ),

269 )

270 return kernel(**kwargs)

tests/kernels/moe/test_deepep_deepgemm_moe.py

+6 -5

@@ -33,7 +33,7 @@ from vllm.v1.worker.workspace import init_workspace_manager

34 from ...utils import multi_gpu_test

35 from .parallel_utils import ProcessGroupInfo, parallel_launch

36 -from .utils import make_test_weights

36 +from .utils import make_dummy_moe_config, make_test_weights

38 if has_deep_ep():

39 from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import (

@@ -192,6 +192,7 @@ def make_ll_modular_kernel(

192 max_num_tokens=max_tokens_per_rank,

193 num_dispatchers=pgi.world_size // dp_size,

194 quant_config=quant_config,

195 + moe_config=make_dummy_moe_config(),

196 )

197 mk = FusedMoEModularKernel(prepare_finalize=a2a, fused_experts=fused_experts)

198 return mk

@@ -219,7 +220,10 @@ def make_ht_modular_kernel(

220 block_shape=test_config.block_size,

221 )

222

222 - fused_experts = DeepGemmExperts(quant_config)

223 + fused_experts = DeepGemmExperts(

224 + moe_config=make_dummy_moe_config(),

225 + quant_config=quant_config,

226 + )

227 mk = FusedMoEModularKernel(prepare_finalize=a2a, fused_experts=fused_experts)

228 return mk

229

@@ -349,9 +353,6 @@ def triton_impl(

353 topk_ids=topk_ids,

354 inplace=False,

355 quant_config=quant_config,

352 - # Make sure this is set to False so we

353 - # don't end up comparing the same implementation.

354 - allow_deep_gemm=False,

356 )

357

358

tests/kernels/moe/test_deepep_moe.py

+11 -2

@@ -10,11 +10,14 @@ import pytest

10 import torch.distributed

11 from torch.distributed import ProcessGroup

13 +from tests.kernels.moe.utils import make_dummy_moe_config

14 from vllm import _custom_ops as ops

15 from vllm.config import VllmConfig, set_current_vllm_config

16 from vllm.model_executor.layers.activation import SiluAndMul

17 from vllm.model_executor.layers.fused_moe import TritonExperts

17 -from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

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

19 + FusedMoEQuantConfig,

20 +)

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

22 from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel

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

@@ -160,15 +163,21 @@ def make_modular_kernel(

163

164 num_dispatchers = pgi.world_size // dp_size

165

166 + moe_config = make_dummy_moe_config()

167 +

168 if low_latency_mode:

169 assert not quant_config.per_act_token_quant, "not supported in ll mode"

170 fused_experts = BatchedTritonExperts(

171 max_num_tokens=MAX_TOKENS_PER_RANK,

172 num_dispatchers=num_dispatchers,

173 + moe_config=moe_config,

174 quant_config=quant_config,

175 )

176 else:

171 - fused_experts = TritonExperts(quant_config=quant_config)

177 + fused_experts = TritonExperts(

178 + moe_config=moe_config,

179 + quant_config=quant_config,

180 + )

181

182 mk = FusedMoEModularKernel(prepare_finalize=a2a, fused_experts=fused_experts)

183 return mk

tests/kernels/moe/test_deepgemm.py

+26 -13

@@ -11,10 +11,19 @@ import math

11 import pytest

12 import torch

14 -from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config

15 -

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.config import (

18 + fp8_w8a8_moe_quant_config,

19 +)

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

24 +from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (

25 + TritonOrDeepGemmExperts,

26 +)

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

28 per_token_group_quant_fp8,

29 )

@@ -100,6 +109,14 @@ def run_single_case(m, n, k, topk, num_experts, block_size):

109 block_shape=block_size,

110 )

111

112 + deep_gemm_experts = mk.FusedMoEModularKernel(

113 + prepare_finalize=MoEPrepareAndFinalizeNoEP(),

114 + fused_experts=TritonOrDeepGemmExperts(

115 + moe_config=make_dummy_moe_config(),

116 + quant_config=quant_config,

117 + ),

118 + )

119 +

120 # triton reference

121 out_triton = fused_experts(

122 hidden_states=tokens_bf16,

@@ -109,19 +126,16 @@ def run_single_case(m, n, k, topk, num_experts, block_size):

126 topk_ids=topk_ids,

127 inplace=False,

128 quant_config=quant_config,

112 - allow_deep_gemm=False,

129 )

130

131 # DeepGemm

116 - out_deepgemm = fused_experts(

132 + out_deepgemm = deep_gemm_experts(

133 hidden_states=tokens_bf16,

134 w1=w1,

135 w2=w2,

136 topk_weights=topk_weights,

137 topk_ids=topk_ids,

138 inplace=False,

123 - quant_config=quant_config,

124 - allow_deep_gemm=True,

139 )

140 diff = calc_diff(out_deepgemm, out_triton)

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

@@ -147,20 +161,19 @@ def test_deepgemm_vs_triton(m, n, k, topk, num_experts, monkeypatch, workspace_i

161 with monkeypatch.context() as mp:

162 mp.setenv("VLLM_USE_DEEP_GEMM", "1")

163

150 - _fused_moe_mod = importlib.import_module(

151 - "vllm.model_executor.layers.fused_moe.fused_moe"

152 - )

164 + _DeepGemmExperts = importlib.import_module(

165 + "vllm.model_executor.layers.fused_moe.deep_gemm_moe"

166 + ).DeepGemmExperts

167

168 call_counter = {"cnt": 0}

169

156 - orig_fn = _fused_moe_mod.deep_gemm_moe_fp8

170 + orig_fn = _DeepGemmExperts.apply

171

158 - def _spy_deep_gemm_moe_fp8(*args, **kwargs):

172 + def _spy_apply(*args, **kwargs):

173 call_counter["cnt"] += 1

174 return orig_fn(*args, **kwargs)

175

162 - monkeypatch.setattr(_fused_moe_mod, "deep_gemm_moe_fp8", _spy_deep_gemm_moe_fp8)

163 -

176 + monkeypatch.setattr(_DeepGemmExperts, "apply", _spy_apply)

177 if topk > num_experts:

178 pytest.skip(f"topk={topk} > num_experts={num_experts}")

179

tests/kernels/moe/test_flashinfer.py

+25 -20

@@ -8,7 +8,10 @@ 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.config import (

11 + FusedMoEConfig,

12 + FusedMoEParallelConfig,

13 FusedMoEQuantConfig,

14 + RoutingMethodType,

15 fp8_w8a8_moe_quant_config,

16 )

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

@@ -116,18 +119,7 @@ class TestData:

119 layer.w13_weight_scale = w13_weight_scale

120 layer.w2_weight_scale = w2_weight_scale

121 # Setup dummy config.

119 - layer.moe_parallel_config = mk.FusedMoEParallelConfig(

120 - tp_size=1,

121 - pcp_size=1,

122 - dp_size=1,

123 - ep_size=1,

124 - tp_rank=0,

125 - pcp_rank=0,

126 - dp_rank=0,

127 - ep_rank=0,

128 - use_ep=False,

129 - all2all_backend="naive",

130 - )

122 + layer.moe_parallel_config = mk.FusedMoEParallelConfig.make_no_parallel()

123

124 # flashinfer expects swapped rows for w13

125 layer.w13_weight.data = swap_w13_to_w31(layer.w13_weight.data)

@@ -238,6 +230,8 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(

230 ):

231 set_random_seed(7)

232 monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")

233 + assert activation in ["silu", "relu2_no_mul"]

234 + is_act_and_mul = activation == "silu_and_mul"

235 with set_current_vllm_config(vllm_config):

236 td = TestData.make_moe_tensors_8bit(

237 m, k, n, e, is_trtllm=False, activation=activation

@@ -285,19 +279,30 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(

279 td.layer.get_fused_moe_quant_config = get_fused_moe_quant_config

280 td.layer.quant_method = td.layer

281

282 + moe_config = FusedMoEConfig(

283 + num_experts=e,

284 + experts_per_token=topk,

285 + hidden_dim=k,

286 + intermediate_size_per_partition=n,

287 + num_local_experts=e,

288 + activation=activation,

289 + device="cuda",

290 + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),

291 + in_dtype=torch.bfloat16,

292 + is_act_and_mul=is_act_and_mul,

293 + routing_method=RoutingMethodType.TopK,

294 + )

295 +

296 kernel = mk.FusedMoEModularKernel(

297 MoEPrepareAndFinalizeNoEP(

290 - defer_input_quant=quant_config.is_block_quantized

298 + defer_input_quant=FlashInferExperts.expects_unquantized_inputs(

299 + moe_config=moe_config,

300 + quant_config=quant_config,

301 + )

302 ),

303 FlashInferExperts(

293 - out_dtype=td.layer.orig_dtype,

304 + moe_config=moe_config,

305 quant_config=quant_config,

295 - ep_rank=td.layer.moe_parallel_config.ep_rank,

296 - ep_size=td.layer.moe_parallel_config.ep_size,

297 - tp_rank=td.layer.moe_parallel_config.tp_rank,

298 - tp_size=td.layer.moe_parallel_config.tp_size,

299 - use_dp=False,

300 - use_deepseek_fp8_block_scale=False,

306 ),

307 )

308

tests/kernels/moe/test_flashinfer_moe.py

+29 -5

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

13 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.config import (

17 + FusedMoEConfig,

18 + FusedMoEParallelConfig,

19 + RoutingMethodType,

20 +)

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

22 FlashInferExperts,

23 is_valid_flashinfer_cutlass_fused_moe,

24 )

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

21 - create_flashinfer_prepare_finalize,

22 -)

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

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

27 + MoEPrepareAndFinalizeNoEP,

28 +)

29 from vllm.platforms import current_platform

30 from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe

31 from vllm.utils.torch_utils import set_random_seed

@@ -86,9 +91,28 @@ def test_flashinfer_fp4_moe_no_graph(

92 assert is_valid_flashinfer_cutlass_fused_moe(a, w1_q, w2_q)

94 + moe_config = FusedMoEConfig(

95 + num_experts=e,

96 + experts_per_token=topk,

97 + hidden_dim=k,

98 + intermediate_size_per_partition=n,

99 + num_local_experts=e,

100 + activation=activation,

101 + device="cuda",

102 + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),

103 + in_dtype=dtype,

104 + is_act_and_mul=is_gated_act,

105 + routing_method=RoutingMethodType.TopK,

106 + )

107 +

108 flashinfer_experts = FusedMoEModularKernel(

90 - create_flashinfer_prepare_finalize(use_dp=False, use_nvfp4=True),

91 - FlashInferExperts(out_dtype=dtype, quant_config=quant_config),

109 + MoEPrepareAndFinalizeNoEP(

110 + defer_input_quant=FlashInferExperts.expects_unquantized_inputs(

111 + moe_config=moe_config,

112 + quant_config=quant_config,

113 + )

114 + ),

115 + FlashInferExperts(moe_config=moe_config, quant_config=quant_config),

116 )

117

118 fi_activation = {"silu_and_mul": "silu", "relu2": "relu2_no_mul"}[activation]

tests/kernels/moe/test_modular_oai_triton_moe.py

+4 -2

@@ -36,6 +36,8 @@ from vllm.model_executor.layers.utils import shuffle_weight

36 from vllm.platforms import current_platform

37 from vllm.utils.torch_utils import set_random_seed

39 +from .utils import make_dummy_moe_config

40 +

41 MNK = [

42 (1, 512, 384),

43 (1, 2880, 2880),

@@ -174,9 +176,9 @@ def oai_triton_moe_impl(

176 )

177

178 if unfused:

177 - fused_experts = UnfusedOAITritonExperts(quant_config)

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

180 else:

179 - fused_experts = OAITritonExperts(quant_config)

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

182

183 mk = FusedMoEModularKernel(MoEPrepareAndFinalizeNoEP(), fused_experts)

184

tests/kernels/moe/test_moe.py

+3 -3

@@ -18,7 +18,7 @@ from transformers import MixtralConfig

18 from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock

20 import vllm.model_executor.layers.fused_moe # noqa

21 -from tests.kernels.moe.utils import fused_moe

21 +from tests.kernels.moe.utils import fused_moe, make_dummy_moe_config

22 from tests.kernels.utils import opcheck, stack_and_dev, torch_experts, torch_moe

23 from vllm._aiter_ops import rocm_aiter_ops

24 from vllm.config import VllmConfig, set_current_vllm_config

@@ -332,7 +332,7 @@ def test_fused_moe(

332 #

333 quant_config = FUSED_MOE_UNQUANTIZED_CONFIG

334

335 - m_fused_moe_fn = modular_triton_fused_moe(quant_config)

335 + m_fused_moe_fn = modular_triton_fused_moe(make_dummy_moe_config(), quant_config)

336

337 def m_fused_moe(

338 a: torch.Tensor,

@@ -437,7 +437,7 @@ def test_naive_block_assignment_moe(

437 #

438 quant_config = FUSED_MOE_UNQUANTIZED_CONFIG

439

440 - m_fused_moe_fn = modular_triton_fused_moe(quant_config)

440 + m_fused_moe_fn = modular_triton_fused_moe(make_dummy_moe_config(), quant_config)

441

442 def m_fused_moe(

443 a: torch.Tensor,

tests/kernels/moe/test_nvfp4_moe.py

+2 -3

@@ -4,7 +4,7 @@ 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

7 +from tests.kernels.moe.utils import make_dummy_moe_config, make_test_weights

8 from tests.kernels.quantization.nvfp4_utils import (

9 FLOAT4_E2M1_MAX,

10 FLOAT8_E4M3_MAX,

@@ -92,8 +92,7 @@ def test_cutlass_fp4_moe_no_graph(

92 kernel = mk.FusedMoEModularKernel(

93 MoEPrepareAndFinalizeNoEP(defer_input_quant=True),

94 CutlassExpertsFp4(

95 - out_dtype=dtype,

96 - max_experts_per_worker=e,

95 + moe_config=make_dummy_moe_config(),

96 quant_config=quant_config,

97 ),

98 )

tests/kernels/moe/test_pplx_cutlass_moe.py

+26 -21

@@ -9,12 +9,18 @@ from tests.kernels.utils import torch_experts

9 from vllm import _custom_ops as ops

10 from vllm.config import VllmConfig, set_current_vllm_config

11 from vllm.model_executor.layers.fused_moe import fused_topk

12 -from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config

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

13 + FusedMoEConfig,

14 + FusedMoEParallelConfig,

15 + RoutingMethodType,

16 + fp8_w8a8_moe_quant_config,

17 +)

18 from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassBatchedExpertsFp8

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

20 from vllm.platforms import current_platform

21 from vllm.utils.math_utils import cdiv

22 from vllm.utils.torch_utils import set_random_seed

23 +from vllm.v1.worker.workspace import init_workspace_manager

25 from ...utils import multi_gpu_test

26 from .parallel_utils import ProcessGroupInfo, parallel_launch

@@ -79,6 +85,8 @@ def pplx_cutlass_moe(

85 PplxPrepareAndFinalize,

86 )

88 + init_workspace_manager(torch.cuda.current_device())

89 +

90 assert torch.cuda.current_device() == pgi.local_rank

92 num_tokens, hidden_dim = a.shape

@@ -132,28 +140,23 @@ def pplx_cutlass_moe(

140 num_dispatchers=num_dispatchers,

141 )

142

135 - ab_strides1 = torch.full(

136 - (num_local_experts,), hidden_dim, device="cuda", dtype=torch.int64

137 - )

138 - ab_strides2 = torch.full(

139 - (num_local_experts,), intermediate_dim, device="cuda", dtype=torch.int64

140 - )

141 - c_strides1 = torch.full(

142 - (num_local_experts,), 2 * intermediate_dim, device="cuda", dtype=torch.int64

143 - )

144 - c_strides2 = torch.full(

145 - (num_local_experts,), hidden_dim, device="cuda", dtype=torch.int64

146 - )

143 + def make_moe_config() -> FusedMoEConfig:

144 + return FusedMoEConfig(

145 + num_experts=num_experts,

146 + experts_per_token=topk,

147 + hidden_dim=hidden_dim,

148 + intermediate_size_per_partition=intermediate_dim,

149 + num_local_experts=num_local_experts,

150 + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),

151 + activation="silu",

152 + in_dtype=torch.bfloat16,

153 + device="cuda",

154 + routing_method=RoutingMethodType.Llama4,

155 + )

156

157 experts = CutlassBatchedExpertsFp8(

149 - num_local_experts,

150 - num_dispatchers,

151 - out_dtype,

152 - ab_strides1,

153 - ab_strides2,

154 - c_strides1,

155 - c_strides2,

156 - fp8_w8a8_moe_quant_config(

158 + moe_config=make_moe_config(),

159 + quant_config=fp8_w8a8_moe_quant_config(

160 per_act_token_quant=per_act_token,

161 per_out_ch_quant=per_out_ch,

162 w1_scale=chunk_by_rank(w1_scale, rank, world_size),

@@ -162,6 +165,8 @@ def pplx_cutlass_moe(

165 if per_act_token

166 else a1_scale[rank],

167 ),

168 + max_num_tokens=max_num_tokens,

169 + num_dispatchers=num_dispatchers,

170 )

171

172 fused_cutlass_experts = FusedMoEModularKernel(

tests/kernels/moe/test_pplx_moe.py

@@ -29,6 +29,7 @@ except ImportError:

30 from tests.kernels.moe.modular_kernel_tools.parallel_utils import _set_vllm_config

31 from tests.kernels.moe.utils import (

32 + make_dummy_moe_config,

33 make_shared_experts,

34 make_test_weights,

35 naive_batched_moe,

@@ -584,6 +585,7 @@ def pplx_moe(

585 max_num_tokens=max_num_tokens,

586 num_dispatchers=prepare_finalize.num_dispatchers(),

587 quant_config=quant_config,

588 + moe_config=make_dummy_moe_config(),

589 )

590

591 fused_experts = FusedMoEModularKernel(

tests/kernels/moe/test_routing.py

-7

@@ -6,7 +6,6 @@ import pytest

6 import torch

8 from vllm.distributed.eplb.eplb_state import EplbLayerState

9 -from vllm.model_executor.layers.fused_moe.config import RoutingMethodType

9 from vllm.model_executor.layers.fused_moe.router.router_factory import (

10 create_fused_moe_router,

11 )

@@ -385,17 +384,11 @@ def test_grouped_topk(

384 global_num_experts,

385 )

386

388 - routing_method_type = None

389 - if scoring_func == "llama4":

390 - routing_method_type = RoutingMethodType.Llama4

391 - scoring_func = "sigmoid"

392 -

387 router = create_fused_moe_router(

388 use_grouped_topk=True,

389 num_expert_group=num_expert_group,

390 topk_group=topk_group,

391 scoring_func=scoring_func,

398 - routing_method_type=routing_method_type,

392 e_score_correction_bias=e_score_correction_bias,

393 routed_scaling_factor=routed_scaling_factor,

394 top_k=top_k,

tests/kernels/moe/test_triton_moe_no_act_mul.py

+13 -3

@@ -10,6 +10,7 @@ equals N (not N // 2 like gated activations).

10 import pytest

11 import torch

13 +from tests.kernels.moe.utils import make_dummy_moe_config

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

15 FUSED_MOE_UNQUANTIZED_CONFIG,

16 )

@@ -78,7 +79,10 @@ def test_triton_experts_no_mul_activation(

79 m, n, k, NUM_EXPERTS, topk

80 )

81 - experts = TritonExperts(FUSED_MOE_UNQUANTIZED_CONFIG)

82 + experts = TritonExperts(

83 + moe_config=make_dummy_moe_config(),

84 + quant_config=FUSED_MOE_UNQUANTIZED_CONFIG,

85 + )

87 ws1_shape, ws2_shape, out_shape = experts.workspace_shapes(

88 M=m,

@@ -151,7 +155,10 @@ def test_workspace_shapes_no_mul_vs_gated():

155

156 M, N, K, topk = 64, 256, 128, 2

157

154 - experts = TritonExperts(FUSED_MOE_UNQUANTIZED_CONFIG)

158 + experts = TritonExperts(

159 + moe_config=make_dummy_moe_config(),

160 + quant_config=FUSED_MOE_UNQUANTIZED_CONFIG,

161 + )

162

163 ws1_no_mul, _, out_no_mul = experts.workspace_shapes(

164 M, N, K, topk, 8, 8, None, SILU_NO_MUL

@@ -187,7 +194,10 @@ def test_adjust_n_for_activation():

194 """Test the adjust_N_for_activation method."""

195 from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts

196

190 - experts = TritonExperts(FUSED_MOE_UNQUANTIZED_CONFIG)

197 + experts = TritonExperts(

198 + moe_config=make_dummy_moe_config(),

199 + quant_config=FUSED_MOE_UNQUANTIZED_CONFIG,

200 + )

201

202 N = 256

203

tests/kernels/moe/utils.py

+36 -1

@@ -8,7 +8,12 @@ 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 import fused_experts, fused_topk

11 -from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

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

12 + FusedMoEConfig,

13 + FusedMoEParallelConfig,

14 + FusedMoEQuantConfig,

15 + RoutingMethodType,

16 +)

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

18 BatchedPrepareAndFinalize,

19 BatchedTritonExperts,

@@ -20,6 +25,34 @@ from vllm.utils.deep_gemm import per_block_cast_to_fp8

25 from vllm.utils.math_utils import round_up

28 +def make_dummy_moe_config(

29 + num_experts: int = 1,

30 + experts_per_token: int = 1,

31 + hidden_dim: int = 1,

32 + intermediate_size_per_partition: int = 1,

33 + in_dtype: torch.dtype = torch.bfloat16,

34 +) -> FusedMoEConfig:

35 + """

36 + This is a dummy config for the mk constructor interface

37 + as most kernels like DeepGEMM, CUTLASSFp4, Triton, MARLIN

38 + do not actually use this config.

39 +

40 + CUTLASSFp8 needs to set some params for workshapes.

41 + """

42 + return FusedMoEConfig(

43 + num_experts=num_experts,

44 + experts_per_token=experts_per_token,

45 + hidden_dim=hidden_dim,

46 + intermediate_size_per_partition=intermediate_size_per_partition,

47 + num_local_experts=num_experts,

48 + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),

49 + activation="silu",

50 + in_dtype=in_dtype,

51 + device="cuda",

52 + routing_method=RoutingMethodType.TopK,

53 + )

54 +

55 +

56 def triton_moe(

57 a: torch.Tensor,

58 w1: torch.Tensor,

@@ -81,6 +114,7 @@ def batched_moe(

114 max_num_tokens=max_num_tokens,

115 num_dispatchers=1,

116 quant_config=quant_config,

117 + moe_config=make_dummy_moe_config(),

118 ),

119 )

120

@@ -121,6 +155,7 @@ def naive_batched_moe(

155 max_num_tokens=max_num_tokens,

156 num_dispatchers=1,

157 quant_config=quant_config,

158 + moe_config=make_dummy_moe_config(),

159 ),

160 )

161

tools/vllm-rocm/pin_rocm_dependencies.py

+2 -1

@@ -11,10 +11,11 @@ This ensures that 'pip install vllm' automatically installs the correct custom w

11 instead of allowing pip to download different versions from PyPI.

12 """

14 -import re

14 import sys

15 from pathlib import Path

17 +import regex as re

18 +

20 def extract_version_from_wheel(wheel_name: str) -> str:

21 """

vllm/compilation/activation_quant_fusion.py

+3 -3

@@ -18,7 +18,7 @@ from vllm.logger import init_logger

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

19 QuantKey,

20 kFp8StaticTensorSym,

21 - kNvfp4Quant,

21 + kNvfp4Dynamic,

22 )

23 from vllm.platforms import current_platform

@@ -41,7 +41,7 @@ silu_and_mul_nvfp4_quant_supported = current_platform.is_cuda() and hasattr(

41 torch.ops._C, "silu_and_mul_nvfp4_quant"

42 )

43 if silu_and_mul_nvfp4_quant_supported:

44 - FUSED_OPS[kNvfp4Quant] = torch.ops._C.silu_and_mul_nvfp4_quant.default # noqa: E501

44 + FUSED_OPS[kNvfp4Dynamic] = torch.ops._C.silu_and_mul_nvfp4_quant.default # noqa: E501

47 class ActivationQuantPattern(ABC):

@@ -129,7 +129,7 @@ class SiluMulNvfp4QuantPattern(ActivationQuantPattern):

129 """

130

131 def __init__(self) -> None:

132 - super().__init__(kNvfp4Quant)

132 + super().__init__(kNvfp4Dynamic)

133

134 def get_inputs(self) -> list[torch.Tensor]:

135 result = self.empty_quant(5, 32)

vllm/compilation/fusion.py

+2 -2

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

20 kFp8DynamicTensorSym,

21 kFp8DynamicTokenSym,

22 kFp8StaticTensorSym,

23 - kNvfp4Quant,

23 + kNvfp4Dynamic,

24 kStaticTensorScale,

25 )

26 from vllm.platforms import current_platform

@@ -63,7 +63,7 @@ QUANT_OPS: dict[QuantKey, OpOverload] = {

63 kFp8DynamicTokenSym: torch.ops._C.dynamic_per_token_scaled_fp8_quant.default, # noqa: E501

64 }

65 if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"):

66 - QUANT_OPS[kNvfp4Quant] = torch.ops._C.scaled_fp4_quant.default

66 + QUANT_OPS[kNvfp4Dynamic] = torch.ops._C.scaled_fp4_quant.default

67 if current_platform.is_cuda():

68 QUANT_OPS[kFp8Dynamic128Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501

69 QUANT_OPS[kFp8Dynamic64Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501

vllm/compilation/fusion_attn.py

+2 -2

@@ -16,7 +16,7 @@ from vllm.config import VllmConfig, get_layers_from_vllm_config

16 from vllm.logger import init_logger

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

18 QuantKey,

19 - kNvfp4Quant,

19 + kNvfp4Dynamic,

20 kStaticTensorScale,

21 )

22 from vllm.platforms import current_platform

@@ -217,7 +217,7 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):

217 """

218

219 def __init__(self, layer: Attention, dtype: torch.dtype) -> None:

220 - super().__init__(layer, kNvfp4Quant, dtype)

220 + super().__init__(layer, kNvfp4Dynamic, dtype)

221

222 def _register(self, pm_pass: PatternMatcherPass) -> None:

223 def pattern(

vllm/compilation/matcher_utils.py

+2 -2

@@ -21,7 +21,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

21 kFp8DynamicTensorSym,

22 kFp8DynamicTokenSym,

23 kFp8StaticTensorSym,

24 - kNvfp4Quant,

24 + kNvfp4Dynamic,

25 )

26 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding

27 from vllm.platforms import current_platform

@@ -38,7 +38,7 @@ QUANT_OPS: dict[QuantKey, OpOverload] = {

38 }

40 if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"):

41 - QUANT_OPS[kNvfp4Quant] = torch.ops._C.scaled_fp4_quant.default # noqa: E501

41 + QUANT_OPS[kNvfp4Dynamic] = torch.ops._C.scaled_fp4_quant.default # noqa: E501

43 if current_platform.is_cuda():

44 QUANT_OPS[kFp8Dynamic128Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501

vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py

+48 -13

@@ -7,11 +7,20 @@ import torch

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

8 from vllm.forward_context import get_forward_context, is_forward_context_available

9 from vllm.logger import init_logger

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

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

11 + FusedMoEConfig,

12 + FusedMoEParallelConfig,

13 + FusedMoEQuantConfig,

14 +)

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

16 TopKWeightAndReduceDelegate,

17 )

18 from vllm.model_executor.layers.fused_moe.utils import _resize_cache

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

20 + QuantKey,

21 + kFp8Dynamic128Sym,

22 + kFp8Static128BlockSym,

23 +)

24 from vllm.platforms import current_platform

25 from vllm.triton_utils import tl, triton

26 from vllm.utils.deep_gemm import (

@@ -19,6 +28,7 @@ from vllm.utils.deep_gemm import (

28 fp8_m_grouped_gemm_nt_masked,

29 get_mk_alignment_for_contiguous_layout,

30 is_deep_gemm_e8m0_used,

31 + is_deep_gemm_supported,

32 )

33 from vllm.utils.math_utils import cdiv, round_up

@@ -253,29 +263,52 @@ def persistent_masked_m_silu_mul_quant(

263 class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

264 def __init__(

265 self,

266 + moe_config: FusedMoEConfig,

267 + quant_config: FusedMoEQuantConfig,

268 max_num_tokens: int,

269 num_dispatchers: int,

258 - quant_config: FusedMoEQuantConfig,

270 ):

271 """

272 max_num_tokens: Maximum number of tokens from a DP Rank

273 num_dispatchers: The number of DP dispatchers.

274 quant_config: Quantization configuration

275 """

265 - super().__init__(quant_config)

276 + super().__init__(

277 + moe_config=moe_config,

278 + quant_config=quant_config,

279 + max_num_tokens=max_num_tokens,

280 + num_dispatchers=num_dispatchers,

281 + )

282 assert self.block_shape == get_mk_alignment_for_contiguous_layout()

283 assert self.quant_config.use_fp8_w8a8

268 - self.max_num_tokens = max_num_tokens

269 - self.num_dispatchers = num_dispatchers

284

271 - @property

272 - def activation_formats(

273 - self,

274 - ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:

275 - return (

276 - mk.FusedMoEActivationFormat.BatchedExperts,

277 - mk.FusedMoEActivationFormat.BatchedExperts,

278 - )

285 + @staticmethod

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

287 + return mk.FusedMoEActivationFormat.BatchedExperts

288 +

289 + @staticmethod

290 + def _supports_current_device() -> bool:

291 + return is_deep_gemm_supported()

292 +

293 + @staticmethod

294 + def _supports_no_act_and_mul() -> bool:

295 + return False

296 +

297 + @staticmethod

298 + def _supports_quant_scheme(

299 + weight_key: QuantKey | None,

300 + activation_key: QuantKey | None,

301 + ) -> bool:

302 + SUPPORTED_W_A = [(kFp8Static128BlockSym, kFp8Dynamic128Sym)]

303 + return (weight_key, activation_key) in SUPPORTED_W_A

304 +

305 + @staticmethod

306 + def _supports_activation(activation: str) -> bool:

307 + return activation in ["silu"]

308 +

309 + @staticmethod

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

311 + return True

312

313 def supports_chunking(self) -> bool:

314 return False

@@ -310,6 +343,8 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

343 # FIXME (varun): We should be able to dispatch only from the leader

344 # DP ranks in the case of TP > 1. At the moment, all the Ranks

345 # end up sending their tokens. This needs to be fixed.

346 + assert self.num_dispatchers is not None

347 + assert self.max_num_tokens is not None

348 num_dispatchers = self.num_dispatchers

349 num_experts = local_num_experts

350 max_num_tokens = M if self.max_num_tokens is None else self.max_num_tokens

vllm/model_executor/layers/fused_moe/config.py

+34 -1

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

862

863 use_ep: bool # whether to use EP or not

864 all2all_backend: str # all2all backend for MoE communication

865 + enable_eplb: bool # whether to enable expert load balancing

866

867 @property

868 def use_all2all_kernels(self):

@@ -882,6 +883,16 @@ class FusedMoEParallelConfig:

883 def use_deepep_ll_kernels(self):

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

885

886 + @property

887 + def use_batched_activation_format(self):

888 + return self.use_deepep_ll_kernels or self.use_pplx_kernels

889 +

890 + @property

891 + def use_naive_all2all_kernels(self):

892 + return self.use_all2all_kernels and (

893 + self.all2all_backend in ["naive", "allgather_reducescatter"]

894 + )

895 +

896 @staticmethod

897 def flatten_tp_across_dp_and_pcp(

898 tp_size: int, dp_size: int, dp_rank: int, pcp_size: int, pcp_rank: int

@@ -999,6 +1010,7 @@ class FusedMoEParallelConfig:

1010 ep_rank=0,

1011 use_ep=False,

1012 all2all_backend=vllm_parallel_config.all2all_backend,

1013 + enable_eplb=vllm_parallel_config.enable_eplb,

1014 )

1015 # DP + EP / TP + EP / DP + TP + EP

1016 assert use_ep

@@ -1017,6 +1029,24 @@ class FusedMoEParallelConfig:

1029 ep_rank=ep_rank,

1030 use_ep=True,

1031 all2all_backend=vllm_parallel_config.all2all_backend,

1032 + enable_eplb=vllm_parallel_config.enable_eplb,

1033 + )

1034 +

1035 + @classmethod

1036 + def make_no_parallel(cls) -> "FusedMoEParallelConfig":

1037 + """For usage in CI/CD and testing."""

1038 + return FusedMoEParallelConfig(

1039 + tp_size=1,

1040 + tp_rank=0,

1041 + pcp_size=1,

1042 + pcp_rank=0,

1043 + dp_size=1,

1044 + dp_rank=0,

1045 + ep_size=1,

1046 + ep_rank=0,

1047 + use_ep=False,

1048 + all2all_backend="naive",

1049 + enable_eplb=False,

1050 )

1051

1052

@@ -1026,8 +1056,11 @@ class FusedMoEConfig:

1056 num_experts: int

1057 experts_per_token: int

1058 hidden_dim: int

1029 -

1059 + intermediate_size_per_partition: int

1060 num_local_experts: int

1061 + activation: str

1062 + device: torch.device | str

1063 + routing_method: RoutingMethodType

1064 moe_parallel_config: FusedMoEParallelConfig

1065

1066 # The activation type.

vllm/model_executor/layers/fused_moe/cutlass_moe.py

+164 -91

@@ -7,7 +7,11 @@ import torch

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

8 from vllm import _custom_ops as ops

9 from vllm.logger import init_logger

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

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

11 + FusedMoEConfig,

12 + FusedMoEParallelConfig,

13 + FusedMoEQuantConfig,

14 +)

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

16 moe_permute,

17 moe_unpermute,

@@ -23,6 +27,19 @@ from vllm.model_executor.layers.fused_moe.utils import (

27 _resize_cache,

28 apply_moe_activation,

29 )

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

31 + QuantKey,

32 + kFp8DynamicTensorSym,

33 + kFp8DynamicTokenSym,

34 + kFp8StaticChannelSym,

35 + kFp8StaticTensorSym,

36 + kNvfp4Dynamic,

37 + kNvfp4Static,

38 +)

39 +from vllm.model_executor.layers.quantization.utils.w8a8_utils import (

40 + cutlass_group_gemm_supported,

41 +)

42 +from vllm.platforms import current_platform

43 from vllm.scalar_type import scalar_types

45 logger = init_logger(__name__)

@@ -238,29 +255,57 @@ def run_cutlass_moe_fp8(

255 class CutlassExpertsFp8Base(mk.FusedMoEPermuteExpertsUnpermute):

256 def __init__(

257 self,

241 - e: int,

242 - n: int,

243 - k: int,

244 - out_dtype: torch.dtype | None,

258 + moe_config: FusedMoEConfig,

259 quant_config: FusedMoEQuantConfig,

246 - device: torch.dtype,

260 + max_num_tokens: int | None = None,

261 + num_dispatchers: int | None = None,

262 ):

263 + super().__init__(

264 + moe_config=moe_config,

265 + quant_config=quant_config,

266 + max_num_tokens=max_num_tokens,

267 + num_dispatchers=num_dispatchers,

268 + )

269 assert quant_config.use_fp8_w8a8

249 - super().__init__(quant_config)

270

251 - # E: num_experts

252 - # N: intermediate size per partition

253 - # K: hidden dim

271 + e = moe_config.num_local_experts

272 + n = moe_config.intermediate_size_per_partition

273 + k = moe_config.hidden_dim

274 + device = moe_config.device

275 ab_strides1_c_strides2 = torch.full((e,), k, device=device, dtype=torch.int64)

276 ab_strides2 = torch.full((e,), n, device=device, dtype=torch.int64)

277 c_strides1 = torch.full((e,), 2 * n, device=device, dtype=torch.int64)

278

258 - self.out_dtype = out_dtype

279 + self.out_dtype = moe_config.in_dtype

280 self.ab_strides1 = ab_strides1_c_strides2

281 self.ab_strides2 = ab_strides2

282 self.c_strides1 = c_strides1

283 self.c_strides2 = ab_strides1_c_strides2

284

285 + @staticmethod

286 + def _supports_current_device() -> bool:

287 + return cutlass_group_gemm_supported()

288 +

289 + @staticmethod

290 + def _supports_no_act_and_mul() -> bool:

291 + return False

292 +

293 + @staticmethod

294 + def _supports_quant_scheme(

295 + weight_key: QuantKey | None,

296 + activation_key: QuantKey | None,

297 + ) -> bool:

298 + SUPPORTED_W_A = [

299 + (kFp8StaticChannelSym, kFp8DynamicTokenSym),

300 + (kFp8StaticTensorSym, kFp8DynamicTensorSym),

301 + (kFp8StaticTensorSym, kFp8StaticTensorSym),

302 + ]

303 + return (weight_key, activation_key) in SUPPORTED_W_A

304 +

305 + @staticmethod

306 + def _supports_activation(activation: str) -> bool:

307 + return activation in ["silu", "gelu", "swigluoai"]

308 +

309 def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:

310 # Let PrepareAndFinalize::finalize() decide the impl.

311 return TopKWeightAndReduceDelegate()

@@ -291,7 +336,7 @@ class CutlassExpertsFp8Base(mk.FusedMoEPermuteExpertsUnpermute):

336 expert_num_tokens = expert_tokens_meta.expert_num_tokens

337

338 use_batched_format = (

294 - self.activation_formats[0] == mk.FusedMoEActivationFormat.BatchedExperts

339 + self.activation_format() == mk.FusedMoEActivationFormat.BatchedExperts

340 )

341

342 in_dtype = hidden_states.dtype

@@ -324,20 +369,23 @@ class CutlassExpertsFp8Base(mk.FusedMoEPermuteExpertsUnpermute):

369

370

371 class CutlassExpertsFp8(CutlassExpertsFp8Base):

327 - @property

328 - def activation_formats(

329 - self,

330 - ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:

331 - return (

332 - mk.FusedMoEActivationFormat.Standard,

333 - mk.FusedMoEActivationFormat.Standard,

334 - )

372 + @staticmethod

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

374 + return mk.FusedMoEActivationFormat.Standard

375 +

376 + @staticmethod

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

378 + # CutlassExpertsFp8 does not support expert map, which is

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

380 + # Note that the BATCHED activation format does not use

381 + # the expert map for identifying experts.

382 + return not moe_parallel_config.use_all2all_kernels

383

384 def supports_chunking(self) -> bool:

385 return True

386

387 def supports_expert_map(self) -> bool:

340 - return True

388 + return False

389

390 def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:

391 # topk weights and reduction are fused in moe_unpermute cuda kernel

@@ -365,26 +413,16 @@ class CutlassExpertsFp8(CutlassExpertsFp8Base):

413

414

415 class CutlassBatchedExpertsFp8(CutlassExpertsFp8Base):

368 - def __init__(

369 - self,

370 - max_experts_per_worker: int,

371 - num_dispatchers: int,

372 - *args,

373 - **kwargs,

374 - ):

375 - super().__init__(*args, **kwargs)

376 - assert max_experts_per_worker > 0

377 - self.max_experts_per_worker = max_experts_per_worker

378 - self.num_dispatchers = num_dispatchers

416 + @staticmethod

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

418 + # BATCHED activation format works with EP because

419 + # expert_map is not used to identify experts (the

420 + # info is encoded/managed by the P/F logic).

421 + return True

422

380 - @property

381 - def activation_formats(

382 - self,

383 - ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:

384 - return (

385 - mk.FusedMoEActivationFormat.BatchedExperts,

386 - mk.FusedMoEActivationFormat.BatchedExperts,

387 - )

423 + @staticmethod

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

425 + return mk.FusedMoEActivationFormat.BatchedExperts

426

427 def supports_chunking(self) -> bool:

428 return False

@@ -408,14 +446,15 @@ class CutlassBatchedExpertsFp8(CutlassExpertsFp8Base):

446 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

447 num_dp = self.num_dispatchers

448 assert num_dp is not None

449 + experts_per_worker = self.moe_config.num_local_experts

450 activation_out_dim = self.adjust_N_for_activation(N, activation)

412 - workspace1 = (self.max_experts_per_worker, M * num_dp, max(N, K))

451 + workspace1 = (experts_per_worker, M * num_dp, max(N, K))

452 workspace2 = (

414 - self.max_experts_per_worker,

453 + experts_per_worker,

454 M * num_dp,

455 max(activation_out_dim, K),

456 )

418 - output = (self.max_experts_per_worker, M, K)

457 + output = (experts_per_worker, M, K)

458 return (workspace1, workspace2, output)

459

460

@@ -601,34 +640,41 @@ def run_cutlass_moe_fp4(

640 return

641

642

604 -# Split into batched and non-batched

643 class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):

606 - def __init__(

607 - self,

608 - max_experts_per_worker: int,

609 - out_dtype: torch.dtype,

610 - quant_config: FusedMoEQuantConfig,

611 - use_batched_format: bool = False,

612 - ):

613 - super().__init__(quant_config)

614 - self.max_experts_per_worker = max_experts_per_worker

615 - self.out_dtype = out_dtype

616 - self.use_batched_format = use_batched_format

644 + @staticmethod

645 + def expects_unquantized_inputs(

646 + moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig

647 + ) -> bool:

648 + return True

649

618 - @property

619 - def activation_formats(

620 - self,

621 - ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:

622 - if self.use_batched_format:

623 - return (

624 - mk.FusedMoEActivationFormat.BatchedExperts,

625 - mk.FusedMoEActivationFormat.BatchedExperts,

626 - )

627 - else:

628 - return (

629 - mk.FusedMoEActivationFormat.Standard,

630 - mk.FusedMoEActivationFormat.Standard,

631 - )

650 + @staticmethod

651 + def _supports_current_device() -> bool:

652 + return current_platform.has_device_capability((10, 0))

653 +

654 + @staticmethod

655 + def _supports_no_act_and_mul() -> bool:

656 + return False

657 +

658 + @staticmethod

659 + def _supports_quant_scheme(

660 + weight_key: QuantKey | None,

661 + activation_key: QuantKey | None,

662 + ) -> bool:

663 + return (weight_key, activation_key) == (kNvfp4Static, kNvfp4Dynamic)

664 +

665 + @staticmethod

666 + def _supports_activation(activation: str) -> bool:

667 + return activation in ["silu", "gelu", "swigluoai"]

668 +

669 + @staticmethod

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

671 + # CutlassExpertsFp4 does not support expert map, which is

672 + # needed for STANDARD activation format kernels in EP mode.

673 + return moe_parallel_config.ep_size == 1

674 +

675 + @staticmethod

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

677 + return mk.FusedMoEActivationFormat.Standard

678

679 def supports_expert_map(self) -> bool:

680 return False

@@ -640,7 +686,7 @@ class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):

686 return TopKWeightAndReduceNoOP()

687

688 def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype:

643 - return self.out_dtype if self.out_dtype is not None else act_dtype

689 + return act_dtype

690

691 def workspace_shapes(

692 self,

@@ -653,18 +699,9 @@ class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):

699 expert_tokens_meta: mk.ExpertTokensMetadata | None,

700 activation: str,

701 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

656 - activation_out_dim = self.adjust_N_for_activation(N, activation)

657 - workspace1: tuple[int, ...] = ()

658 - workspace2: tuple[int, ...] = ()

659 - output: tuple[int, ...] = ()

660 - if self.use_batched_format:

661 - workspace1 = (self.max_experts_per_worker, M, max(N, K))

662 - workspace2 = (self.max_experts_per_worker, M, activation_out_dim)

663 - output = (self.max_experts_per_worker, M, K)

664 - else:

665 - workspace1 = (M * topk, max(2 * N, K))

666 - workspace2 = (M * topk, N)

667 - output = (M, K)

702 + workspace1 = (M * topk, max(2 * N, K))

703 + workspace2 = (M * topk, N)

704 + output = (M, K)

705 return (workspace1, workspace2, output)

706

707 def apply(

@@ -869,10 +906,11 @@ class CutlassExpertsW4A8Fp8(mk.FusedMoEPermuteExpertsUnpermute):

906 c_strides2: torch.Tensor,

907 s_strides1: torch.Tensor,

908 s_strides2: torch.Tensor,

909 + moe_config: FusedMoEConfig,

910 quant_config: FusedMoEQuantConfig,

911 group_size: int,

912 ):

875 - super().__init__(quant_config)

913 + super().__init__(moe_config=moe_config, quant_config=quant_config)

914 self.out_dtype = out_dtype

915 self.a_strides1 = a_strides1

916 self.a_strides2 = a_strides2

@@ -884,13 +922,46 @@ class CutlassExpertsW4A8Fp8(mk.FusedMoEPermuteExpertsUnpermute):

922 self.s_strides2 = s_strides2

923 self.group_size = group_size

924

887 - @property

888 - def activation_formats(

889 - self,

890 - ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:

891 - return (

892 - mk.FusedMoEActivationFormat.Standard,

893 - mk.FusedMoEActivationFormat.Standard,

925 + @staticmethod

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

927 + return mk.FusedMoEActivationFormat.Standard

928 +

929 + @staticmethod

930 + def _supports_current_device() -> bool:

931 + raise NotImplementedError(

932 + "CutlassExpertsW4A8Fp8 is not yet used by an Oracle. "

933 + "This method should not be called."

934 + )

935 +

936 + @staticmethod

937 + def _supports_no_act_and_mul() -> bool:

938 + raise NotImplementedError(

939 + "CutlassExpertsW4A8Fp8 is not yet used by an Oracle. "

940 + "This method should not be called."

941 + )

942 +

943 + @staticmethod

944 + def _supports_quant_scheme(

945 + weight_key: QuantKey | None,

946 + activation_key: QuantKey | None,

947 + ) -> bool:

948 + raise NotImplementedError(

949 + "CutlassExpertsW4A8Fp8 is not yet used by an Oracle. "

950 + "This method should not be called."

951 + )

952 +

953 + @staticmethod

954 + def _supports_activation(activation: str) -> bool:

955 + raise NotImplementedError(

956 + "CutlassExpertsW4A8Fp8 is not yet used by an Oracle. "

957 + "This method should not be called."

958 + )

959 +

960 + @staticmethod

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

962 + raise NotImplementedError(

963 + "CutlassExpertsW4A8Fp8 is not yet used by an Oracle. "

964 + "This method should not be called."

965 )

966

967 def supports_chunking(self) -> bool:

@@ -947,7 +1018,7 @@ class CutlassExpertsW4A8Fp8(mk.FusedMoEPermuteExpertsUnpermute):

1018 expert_num_tokens = None

1019

1020 use_batched_format = (

950 - self.activation_formats[0] == mk.FusedMoEActivationFormat.BatchedExperts

1021 + self.activation_format() == mk.FusedMoEActivationFormat.BatchedExperts

1022 )

1023 assert not use_batched_format, "batched format not supported"

1024

@@ -1003,6 +1074,7 @@ def cutlass_moe_w4a8_fp8(

1074 s_strides1: torch.Tensor,

1075 s_strides2: torch.Tensor,

1076 quant_config: FusedMoEQuantConfig,

1077 + moe_config: FusedMoEConfig,

1078 activation: str = "silu",

1079 expert_map: torch.Tensor | None = None,

1080 apply_router_weight_on_input: bool = False,

@@ -1076,6 +1148,7 @@ def cutlass_moe_w4a8_fp8(

1148 c_strides2=c_strides2,

1149 s_strides1=s_strides1,

1150 s_strides2=s_strides2,

1151 + moe_config=moe_config,

1152 quant_config=quant_config,

1153 group_size=group_size,

1154 ),

vllm/model_executor/layers/fused_moe/deep_gemm_moe.py

+39 -93

@@ -6,17 +6,15 @@ import torch

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

7 from vllm.logger import init_logger

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

9 + FusedMoEConfig,

10 + FusedMoEParallelConfig,

11 FusedMoEQuantConfig,

10 - fp8_w8a8_moe_quant_config,

12 )

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

14 compute_aligned_M,

15 deepgemm_moe_permute,

16 deepgemm_unpermute_and_reduce,

17 )

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

18 - MoEPrepareAndFinalizeNoEP,

19 -)

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

19 TopKWeightAndReduceNoOP,

20 )

@@ -26,9 +24,15 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (

24 per_token_group_quant_fp8_packed_for_deepgemm,

25 silu_mul_per_token_group_quant_fp8_colmajor,

26 )

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

28 + QuantKey,

29 + kFp8Dynamic128Sym,

30 + kFp8Static128BlockSym,

31 +)

32 from vllm.utils.deep_gemm import (

33 DeepGemmQuantScaleFMT,

34 get_mk_alignment_for_contiguous_layout,

35 + is_deep_gemm_supported,

36 m_grouped_fp8_gemm_nt_contiguous,

37 )

38 from vllm.utils.import_utils import has_deep_gemm

@@ -109,21 +113,42 @@ def _valid_deep_gemm(

113

114

115 class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

112 - def __init__(self, quant_config: FusedMoEQuantConfig):

113 - super().__init__(quant_config)

116 + def __init__(self, moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig):

117 + super().__init__(moe_config=moe_config, quant_config=quant_config)

118 assert quant_config.block_shape == get_mk_alignment_for_contiguous_layout()

119 assert quant_config.quant_dtype == torch.float8_e4m3fn

120 assert not quant_config.per_act_token_quant

121 assert not quant_config.per_out_ch_quant

122

119 - @property

120 - def activation_formats(

121 - self,

122 - ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:

123 - return (

124 - mk.FusedMoEActivationFormat.Standard,

125 - mk.FusedMoEActivationFormat.Standard,

126 - )

123 + @staticmethod

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

125 + return mk.FusedMoEActivationFormat.Standard

126 +

127 + @staticmethod

128 + def _supports_current_device() -> bool:

129 + return is_deep_gemm_supported()

130 +

131 + @staticmethod

132 + def _supports_no_act_and_mul() -> bool:

133 + return False

134 +

135 + @staticmethod

136 + def _supports_quant_scheme(

137 + weight_key: QuantKey | None,

138 + activation_key: QuantKey | None,

139 + ) -> bool:

140 + SUPPORTED_W_A = [

141 + (kFp8Static128BlockSym, kFp8Dynamic128Sym),

142 + ]

143 + return (weight_key, activation_key) in SUPPORTED_W_A

144 +

145 + @staticmethod

146 + def _supports_activation(activation: str) -> bool:

147 + return activation in ["silu"]

148 +

149 + @staticmethod

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

151 + return True

152

153 def supports_chunking(self) -> bool:

154 return True

@@ -283,82 +308,3 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

308 expert_map=expert_map,

309 output=output,

310 )

286 -

287 -

288 -def deep_gemm_moe_fp8(

289 - hidden_states: torch.Tensor,

290 - w1: torch.Tensor,

291 - w2: torch.Tensor,

292 - w1_scale: torch.Tensor,

293 - w2_scale: torch.Tensor,

294 - topk_weights: torch.Tensor,

295 - topk_ids: torch.Tensor,

296 - inplace: bool = False,

297 - activation: str = "silu",

298 - global_num_experts: int = -1,

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

300 - a1_scale: torch.Tensor | None = None,

301 - a2_scale: torch.Tensor | None = None,

302 - apply_router_weight_on_input: bool = False,

303 -) -> torch.Tensor:

304 - """

305 - This function computes a a8w8-quantized Mixture of Experts (MoE) layer

306 - using two sets of quantized weights, w1_q and w2_q, and top-k gating

307 - mechanism. The matrix multiplications are implemented with DeepGemm

308 - grouped gemm.

309 -

310 - Parameters:

311 - - hidden_states (torch.Tensor): The input tensor to the MoE layer.

312 - Shape: [M, K]

313 - - w1 (torch.Tensor): The first set of fp8 quantized expert weights.

314 - Shape: [num_experts, K, 2N] (the weights are passed transposed)

315 - - w2 (torch.Tensor): The second set of fp8 quantized expert weights.

316 - Shape: [num_experts, N, K] (the weights are passed transposed)

317 - - w1_scale (torch.Tensor): The fp32 scale to dequantize w1_q.

318 - Shape: [num_experts] or [num_experts, 2N]

319 - - w2_scale (torch.Tensor): The fp32 scale to dequantize w2_q.

320 - Shape: [num_experts] or [num_experts, K]

321 - - topk_weights (torch.Tensor): The weights of each token->expert mapping.

322 - - topk_ids (torch.Tensor): The token->expert mapping for topk_weights.

323 - - inplace (bool): If True, perform the operation in-place.

324 - Defaults to False.

325 - - activation (str): The activation function to apply after the first

326 - MoE layer.

327 - - global_num_experts (int): The total number of experts in the global

328 - expert space.

329 - - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices

330 - from the global expert space to the local expert space of the expert

331 - parallel shard.

332 - - a1_scale (Optional[torch.Tensor]): The optional fp32 scale to quantize a.

333 - Shape: scalar or [M]

334 - - a2_scale (Optional[torch.Tensor]): The optional fp32 scale to

335 - quantize the intermediate result between the gemms.

336 - Shape: scalar or [M]

337 -

338 - Returns:

339 - - torch.Tensor: The bfloat16 output tensor after applying the MoE layer.

340 - """

341 - quant_config = fp8_w8a8_moe_quant_config(

342 - w1_scale=w1_scale,

343 - w2_scale=w2_scale,

344 - a1_scale=a1_scale,

345 - a2_scale=a2_scale,

346 - block_shape=get_mk_alignment_for_contiguous_layout(),

347 - )

348 -

349 - fn = mk.FusedMoEModularKernel(

350 - MoEPrepareAndFinalizeNoEP(),

351 - DeepGemmExperts(quant_config),

352 - )

353 - return fn(

354 - hidden_states,

355 - w1,

356 - w2,

357 - topk_weights,

358 - topk_ids,

359 - inplace=inplace,

360 - activation=activation,

361 - global_num_experts=global_num_experts,

362 - expert_map=expert_map,

363 - apply_router_weight_on_input=apply_router_weight_on_input,

364 - )

vllm/model_executor/layers/fused_moe/fallback.py

+70 -8

@@ -6,6 +6,8 @@ from abc import ABC, abstractmethod

6 import torch

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

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

10 +from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey

13 class FallbackExperts(mk.FusedMoEPermuteExpertsUnpermute, ABC):

@@ -16,18 +18,78 @@ class FallbackExperts(mk.FusedMoEPermuteExpertsUnpermute, ABC):

18 experts: mk.FusedMoEPermuteExpertsUnpermute,

19 fallback_experts: mk.FusedMoEPermuteExpertsUnpermute,

20 ):

19 - super().__init__(experts.quant_config)

21 + super().__init__(

22 + moe_config=experts.moe_config, quant_config=experts.quant_config

23 + )

24 self.fallback_experts = fallback_experts

25 self.experts = experts

23 - @property

24 - def activation_formats(

25 - self,

26 - ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:

27 - assert (

28 - self.fallback_experts.activation_formats == self.experts.activation_formats

27 + @staticmethod

28 + def get_clses() -> tuple[

29 + type[mk.FusedMoEPermuteExpertsUnpermute],

30 + type[mk.FusedMoEPermuteExpertsUnpermute],

31 + ]:

32 + """

33 + Get the cls for the experts and fallback experts.

34 +

35 + Subclasses should implement this method, so that

36 + we have a consistent way to call the _supports_*

37 + class methods below.

38 + """

39 + raise NotImplementedError(

40 + "Subclasses must return the cls for the experts and fallback experts."

41 + )

42 +

43 + @classmethod

44 + def activation_format(

45 + cls: type["FallbackExperts"],

46 + ) -> mk.FusedMoEActivationFormat:

47 + experts_cls, fallback_cls = cls.get_clses()

48 + assert experts_cls.activation_format() == fallback_cls.activation_format()

49 + return experts_cls.activation_format()

50 +

51 + @classmethod

52 + def _supports_current_device(cls) -> bool:

53 + experts_cls, fallback_cls = cls.get_clses()

54 + return (

55 + experts_cls._supports_current_device()

56 + and fallback_cls._supports_current_device()

57 + )

58 +

59 + @classmethod

60 + def _supports_no_act_and_mul(cls) -> bool:

61 + experts_cls, fallback_cls = cls.get_clses()

62 + return (

63 + experts_cls._supports_no_act_and_mul()

64 + and fallback_cls._supports_no_act_and_mul()

65 )

30 - return self.fallback_experts.activation_formats

66 +

67 + @classmethod

68 + def _supports_quant_scheme(

69 + cls,

70 + weight_key: QuantKey | None,

71 + activation_key: QuantKey | None,

72 + ) -> bool:

73 + experts_cls, fallback_cls = cls.get_clses()

74 + return experts_cls._supports_quant_scheme(

75 + weight_key, activation_key

76 + ) and fallback_cls._supports_quant_scheme(weight_key, activation_key)

77 +

78 + @classmethod

79 + def _supports_activation(cls, activation: str) -> bool:

80 + experts_cls, fallback_cls = cls.get_clses()

81 + return experts_cls._supports_activation(

82 + activation

83 + ) and fallback_cls._supports_activation(activation)

84 +

85 + @classmethod

86 + def _supports_parallel_config(

87 + cls, moe_parallel_config: FusedMoEParallelConfig

88 + ) -> bool:

89 + experts_cls, fallback_cls = cls.get_clses()

90 + return experts_cls._supports_parallel_config(

91 + moe_parallel_config

92 + ) and fallback_cls._supports_parallel_config(moe_parallel_config)

94 def supports_chunking(self) -> bool:

95 assert (

vllm/model_executor/layers/fused_moe/flashinfer_cutedsl_moe.py

+50 -41

@@ -6,13 +6,22 @@ import torch

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

7 from vllm import envs

8 from vllm.logger import init_logger

9 -from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

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

10 + FusedMoEConfig,

11 + FusedMoEParallelConfig,

12 + FusedMoEQuantConfig,

13 +)

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

15 TopKWeightAndReduceDelegate,

16 )

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

18 + QuantKey,

19 + kNvfp4Dynamic,

20 + kNvfp4Static,

21 +)

22 +from vllm.platforms import current_platform

23 from vllm.utils.flashinfer import (

24 flashinfer_cutedsl_grouped_gemm_nt_masked,

15 - has_flashinfer_cutedsl_grouped_gemm_nt_masked,

25 scaled_fp4_grouped_quantize,

26 silu_and_mul_scaled_nvfp4_experts_quantize,

27 )

@@ -20,54 +29,54 @@ from vllm.utils.flashinfer import (

29 logger = init_logger(__name__)

23 -def is_valid_flashinfer_cutedsl_fused_moe(

24 - hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor

25 -) -> bool:

26 - """

27 - Check if the given problem size is supported by the FlashInfer CuteDSL MoE

28 - kernel.

29 - """

30 - if not has_flashinfer_cutedsl_grouped_gemm_nt_masked():

31 - logger.debug_once(

32 - "FlashInferCuteDSLExperts disabled: "

33 - "flashinfer_cutedsl_fused_moe not available."

34 - )

35 - return False

36 - # Data type checks

37 - if (

38 - w1.dtype != torch.uint8

39 - or w2.dtype != torch.uint8

40 - or hidden_states.dtype not in [torch.float32, torch.float16, torch.bfloat16]

41 - ):

42 - logger.debug_once(

43 - "FlashInferCuteDSLExperts disabled: w1/w2 must be torch.uint8 "

44 - f"(got w1={w1.dtype}, w2={w2.dtype}), hidden_states must be "

45 - f"float32, float16, or bfloat16 (got {hidden_states.dtype})."

46 - )

47 - return False

48 - return True

49 -

50 -

32 class FlashInferCuteDSLExperts(mk.FusedMoEPermuteExpertsUnpermute):

33 def __init__(

34 self,

54 - out_dtype: torch.dtype,

35 + moe_config: FusedMoEConfig,

36 quant_config: FusedMoEQuantConfig,

37 + max_num_tokens: int,

38 + num_dispatchers: int,

39 ):

57 - super().__init__(quant_config)

40 + super().__init__(

41 + moe_config=moe_config,

42 + quant_config=quant_config,

43 + max_num_tokens=max_num_tokens,

44 + num_dispatchers=num_dispatchers,

45 + )

46 assert quant_config.quant_dtype == "nvfp4", (

47 "Only nvfp4 quantization are currently supported."

48 )

61 - self.out_dtype = out_dtype

49 + self.out_dtype = moe_config.in_dtype

63 - @property

64 - def activation_formats(

65 - self,

66 - ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:

67 - return (

68 - mk.FusedMoEActivationFormat.BatchedExperts,

69 - mk.FusedMoEActivationFormat.BatchedExperts,

70 - )

51 + @staticmethod

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

53 + return mk.FusedMoEActivationFormat.BatchedExperts

54 +

55 + @staticmethod

56 + def _supports_current_device() -> bool:

57 + return current_platform.is_device_capability_family(100)

58 +

59 + @staticmethod

60 + def _supports_no_act_and_mul() -> bool:

61 + return False

62 +

63 + @staticmethod

64 + def _supports_quant_scheme(

65 + weight_key: QuantKey | None,

66 + activation_key: QuantKey | None,

67 + ) -> bool:

68 + SUPPORTED_W_A = [

69 + (kNvfp4Static, kNvfp4Dynamic),

70 + ]

71 + return (weight_key, activation_key) in SUPPORTED_W_A

72 +

73 + @staticmethod

74 + def _supports_activation(activation: str) -> bool:

75 + return activation in ["silu"]

76 +

77 + @staticmethod

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

79 + return True

81 def supports_expert_map(self) -> bool:

82 return False

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py

+93 -106

@@ -5,13 +5,22 @@ import torch

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

7 from vllm.logger import init_logger

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

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

10 - create_flashinfer_prepare_finalize,

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

9 + FusedMoEParallelConfig,

10 + FusedMoEQuantConfig,

11 )

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

13 TopKWeightAndReduceNoOP,

14 )

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

16 + QuantKey,

17 + kFp8Dynamic128Sym,

18 + kFp8Static128BlockSym,

19 + kFp8StaticTensorSym,

20 + kNvfp4Dynamic,

21 + kNvfp4Static,

22 +)

23 +from vllm.platforms import current_platform

24 from vllm.utils.flashinfer import (

25 flashinfer_cutlass_fused_moe,

26 has_flashinfer_cutlass_fused_moe,

@@ -50,40 +59,100 @@ def is_valid_flashinfer_cutlass_fused_moe(

59 class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):

60 def __init__(

61 self,

53 - out_dtype: torch.dtype,

62 + moe_config: mk.FusedMoEConfig,

63 quant_config: FusedMoEQuantConfig,

55 - ep_rank: int = 0,

56 - ep_size: int = 1,

57 - tp_rank: int = 0,

58 - tp_size: int = 1,

59 - use_dp: bool = False,

60 - use_deepseek_fp8_block_scale: bool = False,

64 ):

62 - super().__init__(quant_config)

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

66 assert quant_config.quant_dtype in ("nvfp4", torch.float8_e4m3fn, None), (

67 "Only nvfp4, fp8, bfloat16 and"

68 " float16 quantization are currently supported."

69 )

67 - self.ep_rank = ep_rank

68 - self.ep_size = ep_size

69 - self.tp_rank = tp_rank

70 - self.tp_size = tp_size

71 - self.out_dtype = out_dtype

72 - self.use_dp = use_dp

70 + self.ep_rank = moe_config.moe_parallel_config.ep_rank

71 + self.ep_size = moe_config.moe_parallel_config.ep_size

72 + self.tp_rank = moe_config.moe_parallel_config.tp_rank

73 + self.tp_size = moe_config.moe_parallel_config.tp_size

74 + self.out_dtype = moe_config.in_dtype

75 + self.use_dp = moe_config.moe_parallel_config.dp_size > 1

76 # Enables DeepSeek-style FP8 block-scale path:

77 # - pass per-block weight scales to the kernel

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

76 - self.use_deepseek_fp8_block_scale = use_deepseek_fp8_block_scale

79 + self.use_deepseek_fp8_block_scale = quant_config.is_block_quantized

78 - @property

79 - def activation_formats(

80 - self,

81 - ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:

81 + @staticmethod

82 + def expects_unquantized_inputs(

83 + moe_config: mk.FusedMoEConfig, quant_config: FusedMoEQuantConfig

84 + ) -> bool:

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

86 + # input quantization inside FusedMoEPermuteExpertsUnpermute.

87 + return (

88 + quant_config.use_nvfp4_w4a4

89 + and not moe_config.moe_parallel_config.use_all2all_kernels

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

91 +

92 + @staticmethod

93 + def _supports_current_device() -> bool:

94 return (

83 - mk.FusedMoEActivationFormat.Standard,

84 - mk.FusedMoEActivationFormat.Standard,

95 + current_platform.is_cuda()

96 + and (

97 + current_platform.is_device_capability((9, 0))

98 + or current_platform.is_device_capability_family(100)

99 + )

100 + and has_flashinfer_cutlass_fused_moe()

101 )

102

103 + @staticmethod

104 + def _supports_no_act_and_mul() -> bool:

105 + return True

106 +

107 + @staticmethod

108 + def _supports_quant_scheme(

109 + weight_key: QuantKey | None,

110 + activation_key: QuantKey | None,

111 + ) -> bool:

112 + # The following are supported by FlashInferExperts:

113 + # * unquantized

114 + # * fp8 static per-tensor on 9.0+

115 + # * fp8 block on 9.0

116 + # * nvfp4 on 10.0+

117 +

118 + p = current_platform

119 + scheme = (weight_key, activation_key)

120 + return (

121 + (

122 + scheme

123 + in [

124 + (None, None),

125 + (kFp8StaticTensorSym, kFp8StaticTensorSym),

126 + ]

127 + )

128 + or (

129 + (scheme == (kFp8Static128BlockSym, kFp8Dynamic128Sym))

130 + and (p.is_device_capability((9, 0)))

131 + )

132 + or (

133 + (scheme == (kNvfp4Static, kNvfp4Dynamic))

134 + and (p.is_device_capability_family(100))

135 + )

136 + )

137 +

138 + @staticmethod

139 + def _supports_activation(activation: str) -> bool:

140 + return activation in ["silu", "relu2_no_mul"]

141 +

142 + @staticmethod

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

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

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

146 + # rid of the FlashInfer specific P/F function.

147 + return (

148 + moe_parallel_config.dp_size == 1

149 + or moe_parallel_config.dp_size == moe_parallel_config.ep_size

150 + )

151 +

152 + @staticmethod

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

154 + return mk.FusedMoEActivationFormat.Standard

155 +

156 def supports_expert_map(self) -> bool:

157 return False

158

@@ -231,85 +300,3 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):

300 # No support for LoRA in flashinfer_cutlass_fused_moe.

301 # See TODOs in flashinfer functions runMoe and runMoeMinLantency.

302 raise NotImplementedError("LoRA is not supported for flashinfer_cutlass_moe")

234 -

235 -

236 -def flashinfer_cutlass_moe_fp4(

237 - hidden_states: torch.Tensor,

238 - w1: torch.Tensor,

239 - w2: torch.Tensor,

240 - topk_weights: torch.Tensor,

241 - topk_ids: torch.Tensor,

242 - quant_config: FusedMoEQuantConfig,

243 - inplace: bool = False,

244 - activation: str = "silu",

245 - global_num_experts: int = -1,

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

247 - apply_router_weight_on_input: bool = False,

248 -) -> torch.Tensor:

249 - fused_experts = mk.FusedMoEModularKernel(

250 - create_flashinfer_prepare_finalize(

251 - use_dp=False, use_nvfp4=True, enable_alltoallv=False

252 - ),

253 - FlashInferExperts(

254 - out_dtype=hidden_states.dtype,

255 - quant_config=quant_config,

256 - use_dp=False,

257 - ),

258 - )

259 -

260 - return fused_experts(

261 - hidden_states=hidden_states,

262 - w1=w1,

263 - w2=w2,

264 - topk_weights=topk_weights,

265 - topk_ids=topk_ids,

266 - inplace=inplace,

267 - activation=activation,

268 - global_num_experts=global_num_experts,

269 - expert_map=expert_map,

270 - apply_router_weight_on_input=apply_router_weight_on_input,

271 - )

272 -

273 -

274 -def flashinfer_cutlass_moe(

275 - hidden_states: torch.Tensor,

276 - w1: torch.Tensor,

277 - w2: torch.Tensor,

278 - topk_weights: torch.Tensor,

279 - topk_ids: torch.Tensor,

280 - quant_config: FusedMoEQuantConfig,

281 - inplace: bool = False,

282 - activation: str = "silu",

283 - global_num_experts: int = -1,

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

285 - apply_router_weight_on_input: bool = False,

286 - tp_rank: int = 0,

287 - tp_size: int = 1,

288 - ep_rank: int = 0,

289 - ep_size: int = 1,

290 - use_dp: bool = False,

291 -) -> torch.Tensor:

292 - fused_experts = mk.FusedMoEModularKernel(

293 - create_flashinfer_prepare_finalize(use_dp=use_dp),

294 - FlashInferExperts(

295 - out_dtype=hidden_states.dtype,

296 - quant_config=quant_config,

297 - tp_rank=tp_rank,

298 - tp_size=tp_size,

299 - ep_rank=ep_rank,

300 - ep_size=ep_size,

301 - ),

302 - )

303 -

304 - return fused_experts(

305 - hidden_states=hidden_states,

306 - w1=w1,

307 - w2=w2,

308 - topk_weights=topk_weights,

309 - topk_ids=topk_ids,

310 - inplace=inplace,

311 - activation=activation,

312 - global_num_experts=global_num_experts,

313 - expert_map=expert_map,

314 - apply_router_weight_on_input=apply_router_weight_on_input,

315 - )

vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py

+105 -1

@@ -3,7 +3,12 @@

4 import torch

6 -from vllm.model_executor.layers.fused_moe.config import RoutingMethodType

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

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

8 + FusedMoEConfig,

9 + FusedMoEParallelConfig,

10 + RoutingMethodType,

11 +)

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

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

14 calculate_tile_tokens_dim,

@@ -11,8 +16,107 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

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

17 per_token_group_quant_fp8,

18 )

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

20 + QuantKey,

21 + kFp8Dynamic128Sym,

22 + kFp8Static128BlockSym,

23 + kFp8StaticTensorSym,

24 +)

25 +from vllm.platforms import current_platform

26 from vllm.utils.torch_utils import direct_register_custom_op

28 +#

29 +# Methods used by the oracle for kernel selection.

30 +#

31 +

32 +

33 +def _supports_current_device() -> bool:

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

35 + p = current_platform

36 + # Add check flashinfer trtllm is available

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

38 +

39 +

40 +def _supports_no_act_and_mul() -> bool:

41 + """Does not support non-gated MoE (i.e. Nanotron-Mini)."""

42 + return False

43 +

44 +

45 +def _supports_quant_scheme(

46 + weight_key: QuantKey | None,

47 + activation_key: QuantKey | None,

48 +) -> bool:

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

50 + SUPPORTED_W_A = [

51 + (kFp8Static128BlockSym, kFp8Dynamic128Sym),

52 + (kFp8StaticTensorSym, kFp8StaticTensorSym),

53 + ]

54 + return (weight_key, activation_key) in SUPPORTED_W_A

55 +

56 +

57 +def _supports_activation(activation: str) -> bool:

58 + """Supports silu activation only."""

59 + return activation in ["silu"]

60 +

61 +

62 +def _supports_routing_method(

63 + weight_key: QuantKey | None,

64 + activation_key: QuantKey | None,

65 + routing_method: RoutingMethodType,

66 +) -> bool:

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

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

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

70 + return routing_method in [

71 + RoutingMethodType.DeepSeekV3,

72 + RoutingMethodType.Renormalize,

73 + RoutingMethodType.RenormalizeNaive,

74 + ]

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

76 + # NOTE(rob): kernel requires Llama4.

77 + return routing_method == RoutingMethodType.Llama4

78 +

79 + else:

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

81 +

82 +

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

84 + """Supports TRTLLM Kernel does not support EPLB."""

85 + return not moe_parallel_config.enable_eplb

86 +

87 +

88 +def is_supported_config_trtllm(

89 + moe_config: FusedMoEConfig,

90 + weight_key: QuantKey | None,

91 + activation_key: QuantKey | None,

92 + activation_format: mk.FusedMoEActivationFormat,

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

94 + """

95 + This method mirrors mk.FusedMoEPermuteExpertsUnpermute.is_supported_config

96 + """

97 +

98 + def _make_reason(reason: str) -> str:

99 + return f"kernel does not support {reason}"

100 +

101 + if not _supports_current_device():

102 + return False, _make_reason("current device")

103 + elif not (moe_config.is_act_and_mul or _supports_no_act_and_mul()):

104 + return False, _make_reason("no act_and_mul MLP layer")

105 + elif not _supports_activation(moe_config.activation):

106 + return False, _make_reason(f"{moe_config.activation} activation")

107 + elif not _supports_quant_scheme(weight_key, activation_key):

108 + return False, _make_reason("quantization scheme")

109 + elif not _supports_parallel_config(moe_config.moe_parallel_config):

110 + return False, _make_reason("parallel config")

111 + elif not _supports_routing_method(

112 + weight_key, activation_key, moe_config.routing_method

113 + ):

114 + return False, _make_reason("routing method")

115 + elif activation_format != mk.FusedMoEActivationFormat.Standard:

116 + return False, _make_reason("activation format")

117 +

118 + return True, None

119 +

120

121 def flashinfer_fused_moe_blockscale_fp8(

122 routing_logits: torch.Tensor,

vllm/model_executor/layers/fused_moe/fused_batched_moe.py

+121 -26

@@ -5,7 +5,11 @@

5 import torch

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

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

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

9 + FusedMoEConfig,

10 + FusedMoEParallelConfig,

11 + FusedMoEQuantConfig,

12 +)

13 from vllm.model_executor.layers.fused_moe.fused_moe import try_get_optimal_moe_config

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

15 TopKWeightAndReduceDelegate,

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

21 normalize_batched_scales_shape,

22 normalize_scales_shape,

23 )

20 -from vllm.model_executor.layers.quantization.utils.quant_utils import group_broadcast

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

25 + QuantKey,

26 + group_broadcast,

27 + kFp8Dynamic128Sym,

28 + kFp8DynamicTensorSym,

29 + kFp8DynamicTokenSym,

30 + kFp8Static128BlockSym,

31 + kFp8StaticChannelSym,

32 + kFp8StaticTensorSym,

33 +)

34 +from vllm.platforms import current_platform

35 from vllm.triton_utils import tl, triton

@@ -633,25 +647,62 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):

647

648 def __init__(

649 self,

650 + moe_config: FusedMoEConfig,

651 + quant_config: FusedMoEQuantConfig,

652 max_num_tokens: int,

653 num_dispatchers: int,

638 - quant_config: FusedMoEQuantConfig,

654 ):

640 - super().__init__(quant_config)

655 + super().__init__(

656 + moe_config=moe_config,

657 + quant_config=quant_config,

658 + max_num_tokens=max_num_tokens,

659 + num_dispatchers=num_dispatchers,

660 + )

661 assert not self.quant_config.use_int8_w8a8, "NYI"

662 assert not self.quant_config.use_int8_w8a16, "NYI"

663 assert not self.quant_config.use_int4_w4a16, "NYI"

664 assert self.quant_config.ocp_mx_scheme is None, "NYI"

645 - self.max_num_tokens = max_num_tokens

646 - self.num_dispatchers = num_dispatchers

665

648 - @property

649 - def activation_formats(

650 - self,

651 - ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:

652 - return (

653 - mk.FusedMoEActivationFormat.BatchedExperts,

654 - mk.FusedMoEActivationFormat.BatchedExperts,

666 + @staticmethod

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

668 + return mk.FusedMoEActivationFormat.BatchedExperts

669 +

670 + @staticmethod

671 + def _supports_current_device() -> bool:

672 + raise NotImplementedError(

673 + "NaiveBatchedExperts is not yet used by an Oracle. "

674 + "This method should not be called."

675 + )

676 +

677 + @staticmethod

678 + def _supports_no_act_and_mul() -> bool:

679 + raise NotImplementedError(

680 + "NaiveBatchedExperts is not yet used by an Oracle. "

681 + "This method should not be called."

682 + )

683 +

684 + @staticmethod

685 + def _supports_quant_scheme(

686 + weight_key: QuantKey | None,

687 + activation_key: QuantKey | None,

688 + ) -> bool:

689 + raise NotImplementedError(

690 + "NaiveBatchedExperts is not yet used by an Oracle. "

691 + "This method should not be called."

692 + )

693 +

694 + @staticmethod

695 + def _supports_activation(activation: str) -> bool:

696 + raise NotImplementedError(

697 + "NaiveBatchedExperts is not yet used by an Oracle. "

698 + "This method should not be called."

699 + )

700 +

701 + @staticmethod

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

703 + raise NotImplementedError(

704 + "NaiveBatchedExperts is not yet used by an Oracle. "

705 + "This method should not be called."

706 )

707

708 def supports_chunking(self) -> bool:

@@ -675,6 +726,8 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):

726 expert_tokens_meta: mk.ExpertTokensMetadata | None,

727 activation: str,

728 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

729 + assert self.num_dispatchers is not None

730 + assert self.max_num_tokens is not None

731 num_dp = self.num_dispatchers

732 num_experts = local_num_experts

733 workspace13 = (num_experts, self.max_num_tokens * num_dp, K)

@@ -826,29 +879,69 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

879

880 def __init__(

881 self,

882 + moe_config: FusedMoEConfig,

883 + quant_config: FusedMoEQuantConfig,

884 max_num_tokens: int,

885 num_dispatchers: int,

831 - quant_config: FusedMoEQuantConfig,

886 ):

833 - super().__init__(quant_config)

887 + super().__init__(

888 + moe_config=moe_config,

889 + quant_config=quant_config,

890 + max_num_tokens=max_num_tokens,

891 + num_dispatchers=num_dispatchers,

892 + )

893 assert not self.quant_config.use_int8_w8a8, "NYI"

894 assert not self.quant_config.use_int8_w8a16, "NYI"

895 assert not self.quant_config.use_int4_w4a16, "NYI"

896 assert self.quant_config.ocp_mx_scheme is None, "NYI"

838 - assert max_num_tokens > 0

839 - assert num_dispatchers > 0

840 - self.max_num_tokens = max_num_tokens

841 - self.num_dispatchers = num_dispatchers

897

843 - @property

844 - def activation_formats(

845 - self,

846 - ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:

847 - return (

848 - mk.FusedMoEActivationFormat.BatchedExperts,

849 - mk.FusedMoEActivationFormat.BatchedExperts,

898 + @staticmethod

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

900 + return mk.FusedMoEActivationFormat.BatchedExperts

901 +

902 + @staticmethod

903 + def _supports_current_device() -> bool:

904 + return current_platform.is_cuda_alike()

905 +

906 + @staticmethod

907 + def _supports_no_act_and_mul() -> bool:

908 + return False

909 +

910 + @staticmethod

911 + def _supports_quant_scheme(

912 + weight_key: QuantKey | None,

913 + activation_key: QuantKey | None,

914 + ) -> bool:

915 + p = current_platform

916 + device_supports_fp8 = (p.is_rocm() and p.rocm.on_gfx9()) or (

917 + p.is_cuda() and p.has_device_capability((8, 9))

918 + )

919 +

920 + SUPPORTED_W_A_FP8 = [

921 + (kFp8Static128BlockSym, kFp8Dynamic128Sym),

922 + (kFp8StaticChannelSym, kFp8DynamicTokenSym),

923 + (kFp8StaticTensorSym, kFp8StaticTensorSym),

924 + (kFp8StaticTensorSym, kFp8DynamicTensorSym),

925 + ]

926 + return (weight_key, activation_key) == (None, None) or (

927 + device_supports_fp8 and (weight_key, activation_key) in SUPPORTED_W_A_FP8

928 )

929

930 + @staticmethod

931 + def _supports_activation(activation: str) -> bool:

932 + return activation in [

933 + "silu",

934 + "gelu",

935 + "swigluoai",

936 + "silu_no_mul",

937 + "gelu_no_mul",

938 + "relu2_no_mul",

939 + ]

940 +

941 + @staticmethod

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

943 + return True

944 +

945 def supports_chunking(self) -> bool:

946 return False

947

@@ -870,6 +963,8 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

963 expert_tokens_meta: mk.ExpertTokensMetadata | None,

964 activation: str,

965 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

966 + assert self.num_dispatchers is not None

967 + assert self.max_num_tokens is not None

968 num_dp = self.num_dispatchers

969 num_experts = local_num_experts

970 max_num_tokens = self.max_num_tokens

vllm/model_executor/layers/fused_moe/fused_marlin_moe.py

+80 -46

@@ -8,7 +8,11 @@ import torch

9 import vllm._custom_ops as ops

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

11 -from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

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

12 + FusedMoEConfig,

13 + FusedMoEParallelConfig,

14 + FusedMoEQuantConfig,

15 +)

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

17 batched_moe_align_block_size,

18 moe_align_block_size,

@@ -27,6 +31,13 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (

31 marlin_moe_intermediate_size,

32 marlin_quant_input,

33 )

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

35 + QuantKey,

36 + kFp8Static128BlockSym,

37 + kFp8StaticChannelSym,

38 + kFp8StaticTensorSym,

39 + kNvfp4Static,

40 +)

41 from vllm.platforms import current_platform

42 from vllm.scalar_type import ScalarType, scalar_types

@@ -522,7 +533,10 @@ def batched_fused_marlin_moe(

533 class MarlinExpertsBase(mk.FusedMoEPermuteExpertsUnpermute):

534 def __init__(

535 self,

536 + moe_config: FusedMoEConfig,

537 quant_config: FusedMoEQuantConfig,

538 + max_num_tokens: int | None = None,

539 + num_dispatchers: int | None = None,

540 w13_g_idx: torch.Tensor | None = None,

541 w2_g_idx: torch.Tensor | None = None,

542 w13_g_idx_sort_indices: torch.Tensor | None = None,

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

555 self.w13_g_idx_sort_indices = w13_g_idx_sort_indices

556 self.w2_g_idx_sort_indices = w2_g_idx_sort_indices

557 self.is_k_full = is_k_full

544 - super().__init__(quant_config)

558 + super().__init__(

559 + moe_config=moe_config,

560 + quant_config=quant_config,

561 + max_num_tokens=max_num_tokens,

562 + num_dispatchers=num_dispatchers,

563 + )

564 +

565 + @staticmethod

566 + def _supports_current_device() -> bool:

567 + p = current_platform

568 + return p.is_cuda() and p.has_device_capability((7, 5))

569 +

570 + @staticmethod

571 + def _supports_no_act_and_mul() -> bool:

572 + return False

573 +

574 + @staticmethod

575 + def _supports_quant_scheme(

576 + weight_key: QuantKey | None,

577 + activation_key: QuantKey | None,

578 + ) -> bool:

579 + # TODO(rob): add int4, mxfp4, int8 as integrations

580 + # are migrated to use the oracle one-by-one.

581 + SUPPORTED_W = [

582 + kFp8Static128BlockSym,

583 + kFp8StaticChannelSym,

584 + kFp8StaticTensorSym,

585 + kNvfp4Static,

586 + ]

587 + return weight_key in SUPPORTED_W

588 +

589 + @staticmethod

590 + def _supports_activation(activation: str) -> bool:

591 + return activation in [

592 + "silu",

593 + "gelu",

594 + "swigluoai",

595 + "silu_no_mul",

596 + "gelu_no_mul",

597 + "relu2_no_mul",

598 + ]

599 +

600 + @staticmethod

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

602 + return True

603

604 @property

605 def quant_type_id(self) -> int:

@@ -587,38 +645,15 @@ class MarlinExpertsBase(mk.FusedMoEPermuteExpertsUnpermute):

645

646

647 class MarlinExperts(MarlinExpertsBase):

590 - def __init__(

591 - self,

592 - quant_config: FusedMoEQuantConfig,

593 - w13_g_idx: torch.Tensor | None = None,

594 - w2_g_idx: torch.Tensor | None = None,

595 - w13_g_idx_sort_indices: torch.Tensor | None = None,

596 - w2_g_idx_sort_indices: torch.Tensor | None = None,

597 - is_k_full: bool = True,

598 - ):

599 - super().__init__(

600 - quant_config,

601 - w13_g_idx,

602 - w2_g_idx,

603 - w13_g_idx_sort_indices,

604 - w2_g_idx_sort_indices,

605 - is_k_full,

606 - )

607 -

648 def supports_expert_map(self) -> bool:

649 return True

650

651 def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:

652 return TopKWeightAndReduceNoOP()

653

614 - @property

615 - def activation_formats(

616 - self,

617 - ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:

618 - return (

619 - mk.FusedMoEActivationFormat.Standard,

620 - mk.FusedMoEActivationFormat.Standard,

621 - )

654 + @staticmethod

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

656 + return mk.FusedMoEActivationFormat.Standard

657

658 def supports_chunking(self) -> bool:

659 return True

@@ -714,9 +749,10 @@ class MarlinExperts(MarlinExpertsBase):

749 class BatchedMarlinExperts(MarlinExpertsBase):

750 def __init__(

751 self,

752 + moe_config: FusedMoEConfig,

753 + quant_config: FusedMoEQuantConfig,

754 max_num_tokens: int,

755 num_dispatchers: int,

719 - quant_config: FusedMoEQuantConfig,

756 w13_g_idx: torch.Tensor | None = None,

757 w2_g_idx: torch.Tensor | None = None,

758 w13_g_idx_sort_indices: torch.Tensor | None = None,

@@ -724,15 +760,16 @@ class BatchedMarlinExperts(MarlinExpertsBase):

760 is_k_full: bool = True,

761 ):

762 super().__init__(

727 - quant_config,

728 - w13_g_idx,

729 - w2_g_idx,

730 - w13_g_idx_sort_indices,

731 - w2_g_idx_sort_indices,

732 - is_k_full,

763 + moe_config=moe_config,

764 + quant_config=quant_config,

765 + max_num_tokens=max_num_tokens,

766 + num_dispatchers=num_dispatchers,

767 + w13_g_idx=w13_g_idx,

768 + w2_g_idx=w2_g_idx,

769 + w13_g_idx_sort_indices=w13_g_idx_sort_indices,

770 + w2_g_idx_sort_indices=w2_g_idx_sort_indices,

771 + is_k_full=is_k_full,

772 )

734 - self.max_num_tokens = max_num_tokens

735 - self.num_dispatchers = num_dispatchers

773

774 def supports_expert_map(self) -> bool:

775 return True

@@ -740,14 +777,9 @@ class BatchedMarlinExperts(MarlinExpertsBase):

777 def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:

778 return TopKWeightAndReduceDelegate()

779

743 - @property

744 - def activation_formats(

745 - self,

746 - ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:

747 - return (

748 - mk.FusedMoEActivationFormat.BatchedExperts,

749 - mk.FusedMoEActivationFormat.BatchedExperts,

750 - )

780 + @staticmethod

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

782 + return mk.FusedMoEActivationFormat.BatchedExperts

783

784 def supports_chunking(self) -> bool:

785 return False

@@ -763,9 +795,11 @@ class BatchedMarlinExperts(MarlinExpertsBase):

795 expert_tokens_meta: mk.ExpertTokensMetadata | None,

796 activation: str,

797 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:

798 + assert self.num_dispatchers is not None

799 + assert self.max_num_tokens is not None

800 num_dispatchers = self.num_dispatchers

801 num_experts = local_num_experts

768 - max_num_tokens = M if self.max_num_tokens is None else self.max_num_tokens

802 + max_num_tokens = self.max_num_tokens

803 workspace13 = (num_experts * max_num_tokens * num_dispatchers, max(K, N * 2))

804 workspace2 = (num_experts * max_num_tokens * num_dispatchers, N)

805 output = (num_experts, max_num_tokens * num_dispatchers, K)

vllm/model_executor/layers/fused_moe/fused_moe.py

+120 -77

@@ -19,13 +19,11 @@ from vllm.model_executor.layers.batch_invariant import (

19 )

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

21 FUSED_MOE_UNQUANTIZED_CONFIG,

22 + FusedMoEConfig,

23 + FusedMoEParallelConfig,

24 FusedMoEQuantConfig,

25 _get_config_dtype_str,

26 )

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

26 - _valid_deep_gemm,

27 - deep_gemm_moe_fp8,

28 -)

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

28 moe_align_block_size,

29 )

@@ -44,9 +42,16 @@ from vllm.model_executor.layers.fused_moe.utils import (

42 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import dequant_mxfp4

43 from vllm.model_executor.layers.quantization.utils.mxfp6_utils import dequant_mxfp6

44 from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import OCP_MX_Scheme

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

46 + QuantKey,

47 + kFp8Dynamic128Sym,

48 + kFp8DynamicTokenSym,

49 + kFp8Static128BlockSym,

50 + kFp8StaticChannelSym,

51 + kFp8StaticTensorSym,

52 +)

53 from vllm.platforms import current_platform

54 from vllm.triton_utils import tl, triton

49 -from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used

55 from vllm.utils.torch_utils import direct_register_custom_op, is_torch_equal_or_newer

57 logger = init_logger(__name__)

@@ -1534,66 +1539,36 @@ def fused_experts(

1539 global_num_experts: int = -1,

1540 expert_map: torch.Tensor | None = None,

1541 quant_config: FusedMoEQuantConfig | None = None,

1537 - allow_deep_gemm: bool = False,

1542 ) -> torch.Tensor:

1543 if quant_config is None:

1544 quant_config = FUSED_MOE_UNQUANTIZED_CONFIG

1545

1542 - # For now, disable DeepGemm for small N (<= 512) until better

1543 - # permute/unpermute ops are available.

1544 - # However, on B200, we use DeepGemm for all cases because they only support

1545 - # E8M0 scale, which means we requantize the weight and input to the specific

1546 - # scale. Fallen back to cutlass or triton for some cases would cause

1547 - # accuracy issue.

1548 - if (

1549 - allow_deep_gemm

1550 - and quant_config.use_fp8_w8a8

1551 - and (is_deep_gemm_e8m0_used() or _valid_deep_gemm(hidden_states, w1, w2))

1552 - ):

1553 - assert quant_config is not None

1554 - return deep_gemm_moe_fp8(

1555 - hidden_states=hidden_states,

1556 - w1=w1,

1557 - w2=w2,

1558 - topk_weights=topk_weights,

1559 - topk_ids=topk_ids,

1560 - inplace=inplace,

1561 - activation=activation,

1562 - global_num_experts=global_num_experts,

1563 - expert_map=expert_map,

1564 - w1_scale=quant_config.w1_scale,

1565 - w2_scale=quant_config.w2_scale,

1566 - a1_scale=quant_config.a1_scale,

1567 - a2_scale=quant_config.a2_scale,

1568 - apply_router_weight_on_input=apply_router_weight_on_input,

1569 - )

1570 - else:

1571 - return dispatch_fused_experts_func(inplace)(

1572 - hidden_states=hidden_states,

1573 - w1=w1,

1574 - w2=w2,

1575 - topk_weights=topk_weights,

1576 - topk_ids=topk_ids,

1577 - activation=activation,

1578 - apply_router_weight_on_input=apply_router_weight_on_input,

1579 - use_fp8_w8a8=quant_config.use_fp8_w8a8,

1580 - use_int8_w8a8=quant_config.use_int8_w8a8,

1581 - use_int8_w8a16=quant_config.use_int8_w8a16,

1582 - use_int4_w4a16=quant_config.use_int4_w4a16,

1583 - ocp_mx_scheme=quant_config.ocp_mx_scheme,

1584 - per_channel_quant=quant_config.per_act_token_quant,

1585 - global_num_experts=global_num_experts,

1586 - expert_map=expert_map,

1587 - w1_scale=quant_config.w1_scale,

1588 - w2_scale=quant_config.w2_scale,

1589 - w1_zp=quant_config.w1_zp,

1590 - w2_zp=quant_config.w2_zp,

1591 - a1_scale=quant_config.a1_scale,

1592 - a2_scale=quant_config.a2_scale,

1593 - block_shape=quant_config.block_shape,

1594 - w1_bias=quant_config.w1_bias,

1595 - w2_bias=quant_config.w2_bias,

1596 - )

1546 + return dispatch_fused_experts_func(inplace)(

1547 + hidden_states=hidden_states,

1548 + w1=w1,

1549 + w2=w2,

1550 + topk_weights=topk_weights,

1551 + topk_ids=topk_ids,

1552 + activation=activation,

1553 + apply_router_weight_on_input=apply_router_weight_on_input,

1554 + use_fp8_w8a8=quant_config.use_fp8_w8a8,

1555 + use_int8_w8a8=quant_config.use_int8_w8a8,

1556 + use_int8_w8a16=quant_config.use_int8_w8a16,

1557 + use_int4_w4a16=quant_config.use_int4_w4a16,

1558 + ocp_mx_scheme=quant_config.ocp_mx_scheme,

1559 + per_channel_quant=quant_config.per_act_token_quant,

1560 + global_num_experts=global_num_experts,

1561 + expert_map=expert_map,

1562 + w1_scale=quant_config.w1_scale,

1563 + w2_scale=quant_config.w2_scale,

1564 + w1_zp=quant_config.w1_zp,

1565 + w2_zp=quant_config.w2_zp,

1566 + a1_scale=quant_config.a1_scale,

1567 + a2_scale=quant_config.a2_scale,

1568 + block_shape=quant_config.block_shape,

1569 + w1_bias=quant_config.w1_bias,

1570 + w2_bias=quant_config.w2_bias,

1571 + )

1572

1573

1574 def _get_config_quant_dtype(

@@ -1924,19 +1899,53 @@ def fused_experts_impl(

1899 class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

1900 def __init__(

1901 self,

1902 + moe_config: FusedMoEConfig,

1903 quant_config: FusedMoEQuantConfig,

1904 ):

1929 - super().__init__(quant_config)

1930 -

1931 - @property

1932 - def activation_formats(

1933 - self,

1934 - ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:

1935 - return (

1936 - mk.FusedMoEActivationFormat.Standard,

1937 - mk.FusedMoEActivationFormat.Standard,

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

1906 +

1907 + @staticmethod

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

1909 + return mk.FusedMoEActivationFormat.Standard

1910 +

1911 + @staticmethod

1912 + def _supports_current_device() -> bool:

1913 + return current_platform.is_cuda_alike()

1914 +

1915 + @staticmethod

1916 + def _supports_no_act_and_mul() -> bool:

1917 + return False

1918 +

1919 + @staticmethod

1920 + def _supports_quant_scheme(

1921 + weight_key: QuantKey | None,

1922 + activation_key: QuantKey | None,

1923 + ) -> bool:

1924 + p = current_platform

1925 + device_supports_fp8 = (p.is_rocm() and p.rocm.on_gfx9()) or (

1926 + p.is_cuda() and p.has_device_capability((8, 9))

1927 )

1928

1929 + if not device_supports_fp8:

1930 + return (weight_key, activation_key) == (None, None)

1931 +

1932 + SUPPORTED_W_A = [

1933 + (None, None),

1934 + (kFp8Static128BlockSym, kFp8Dynamic128Sym),

1935 + (kFp8StaticChannelSym, kFp8DynamicTokenSym),

1936 + (kFp8StaticTensorSym, kFp8DynamicTokenSym),

1937 + (kFp8StaticTensorSym, kFp8StaticTensorSym),

1938 + ]

1939 + return (weight_key, activation_key) in SUPPORTED_W_A

1940 +

1941 + @staticmethod

1942 + def _supports_activation(activation: str) -> bool:

1943 + return activation in ["silu", "gelu", "swigluoai"]

1944 +

1945 + @staticmethod

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

1947 + return True

1948 +

1949 def supports_chunking(self) -> bool:

1950 return True

1951

@@ -2111,11 +2120,43 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

2120

2121

2122 class TritonWNA16Experts(TritonExperts):

2114 - def __init__(

2115 - self,

2116 - quant_config: FusedMoEQuantConfig,

2117 - ):

2118 - super().__init__(quant_config)

2123 + @staticmethod

2124 + def _supports_current_device() -> bool:

2125 + raise NotImplementedError(

2126 + "TritonWNA16Experts is not yet used by an Oracle. "

2127 + "This method should not be called."

2128 + )

2129 +

2130 + @staticmethod

2131 + def _supports_no_act_and_mul() -> bool:

2132 + raise NotImplementedError(

2133 + "TritonWNA16Experts is not yet used by an Oracle. "

2134 + "This method should not be called."

2135 + )

2136 +

2137 + @staticmethod

2138 + def _supports_quant_scheme(

2139 + weight_key: QuantKey | None,

2140 + activation_key: QuantKey | None,

2141 + ) -> bool:

2142 + raise NotImplementedError(

2143 + "TritonWNA16Experts is not yet used by an Oracle. "

2144 + "This method should not be called."

2145 + )

2146 +

2147 + @staticmethod

2148 + def _supports_activation(activation: str) -> bool:

2149 + raise NotImplementedError(

2150 + "TritonWNA16Experts is not yet used by an Oracle. "

2151 + "This method should not be called."

2152 + )

2153 +

2154 + @staticmethod

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

2156 + raise NotImplementedError(

2157 + "TritonWNA16Experts is not yet used by an Oracle. "

2158 + "This method should not be called."

2159 + )

2160

2161 def apply(

2162 self,

@@ -2254,10 +2295,12 @@ class TritonWNA16Experts(TritonExperts):

2295

2296

2297 def modular_triton_fused_moe(

2257 - quant_config: FusedMoEQuantConfig, shared_experts: torch.nn.Module | None = None

2298 + moe_config: FusedMoEConfig,

2299 + quant_config: FusedMoEQuantConfig,

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

2301 ) -> mk.FusedMoEModularKernel:

2302 return mk.FusedMoEModularKernel(

2303 MoEPrepareAndFinalizeNoEP(),

2261 - TritonExperts(quant_config),

2304 + TritonExperts(moe_config, quant_config),

2305 shared_experts,

2306 )

vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py

+47 -28

@@ -9,12 +9,16 @@ from vllm import _custom_ops as ops

9 from vllm.logger import init_logger

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

11 FUSED_MOE_UNQUANTIZED_CONFIG,

12 + FusedMoEParallelConfig,

13 FusedMoEQuantConfig,

14 )

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

16 TopKWeightAndReduceNoOP,

17 )

18 from vllm.model_executor.layers.fused_moe.utils import _resize_cache

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

20 + QuantKey,

21 +)

22 from vllm.triton_utils import tl, triton

23 from vllm.utils.import_utils import has_triton_kernels

@@ -241,8 +245,43 @@ def make_routing_data(

245

246

247 class BaseOAITritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

244 - def __init__(self, quant_config: FusedMoEQuantConfig):

245 - super().__init__(quant_config)

248 + @staticmethod

249 + def _supports_current_device() -> bool:

250 + raise NotImplementedError(

251 + "OAITritonExperts is not yet used by an Oracle. "

252 + "This method should not be called."

253 + )

254 +

255 + @staticmethod

256 + def _supports_no_act_and_mul() -> bool:

257 + raise NotImplementedError(

258 + "OAITritonExperts is not yet used by an Oracle. "

259 + "This method should not be called."

260 + )

261 +

262 + @staticmethod

263 + def _supports_quant_scheme(

264 + weight_key: QuantKey | None,

265 + activation_key: QuantKey | None,

266 + ) -> bool:

267 + raise NotImplementedError(

268 + "OAITritonExperts is not yet used by an Oracle. "

269 + "This method should not be called."

270 + )

271 +

272 + @staticmethod

273 + def _supports_activation(activation: str) -> bool:

274 + raise NotImplementedError(

275 + "OAITritonExperts is not yet used by an Oracle. "

276 + "This method should not be called."

277 + )

278 +

279 + @staticmethod

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

281 + raise NotImplementedError(

282 + "OAITritonExperts is not yet used by an Oracle. "

283 + "This method should not be called."

284 + )

285

286 def supports_expert_map(self) -> bool:

287 return True

@@ -297,19 +336,9 @@ class BaseOAITritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

336

337

338 class OAITritonExperts(BaseOAITritonExperts):

300 - def __init__(self, quant_config: FusedMoEQuantConfig):

301 - # TODO (varun) : Enable activation quantization

302 - assert quant_config.use_mxfp4_w4a16, "Supports only mxfp4_w4a16"

303 - super().__init__(quant_config)

304 -

305 - @property

306 - def activation_formats(

307 - self,

308 - ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:

309 - return (

310 - mk.FusedMoEActivationFormat.Standard,

311 - mk.FusedMoEActivationFormat.Standard,

312 - )

339 + @staticmethod

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

341 + return mk.FusedMoEActivationFormat.Standard

342

343 def supports_chunking(self) -> bool:

344 return True

@@ -391,19 +420,9 @@ class UnfusedOAITritonExperts(BaseOAITritonExperts):

420 One use case for it is to inject LoRA modules on the activation and moe_sum.

421 """

422

394 - def __init__(self, quant_config: FusedMoEQuantConfig):

395 - # TODO (varun) : Enable activation quantization

396 - assert quant_config.use_mxfp4_w4a16, "Supports only mxfp4_w4a16"

397 - super().__init__(quant_config)

398 -

399 - @property

400 - def activation_formats(

401 - self,

402 - ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:

403 - return (

404 - mk.FusedMoEActivationFormat.Standard,

405 - mk.FusedMoEActivationFormat.Standard,

406 - )

423 + @staticmethod

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

425 + return mk.FusedMoEActivationFormat.Standard

426

427 def supports_chunking(self) -> bool:

428 return True

vllm/model_executor/layers/fused_moe/layer.py

+36 -34

@@ -330,7 +330,6 @@ class FusedMoE(CustomOp):

330 is_sequence_parallel=False,

331 expert_mapping: list[tuple[str, str, int, str]] | None = None,

332 n_shared_experts: int | None = None,

333 - routing_method_type: RoutingMethodType | None = None,

333 router_logits_dtype: torch.dtype | None = None,

334 ):

335 super().__init__()

@@ -519,10 +518,43 @@ class FusedMoE(CustomOp):

518 self.apply_router_weight_on_input = apply_router_weight_on_input

519 self.activation = activation

520

521 + # TODO(bnell): in next PR move capture back to layer

522 + capture: Callable[[torch.Tensor], None] | None = None

523 + if (

524 + self.vllm_config.model_config is not None

525 + and self.vllm_config.model_config.enable_return_routed_experts

526 + ):

527 + # In dummy runs, the capturer is not initialized.

528 + capturer = RoutedExpertsCapturer.get_instance()

529 + if capturer is not None:

530 + capture = lambda topk_ids: capturer.capture(self.layer_id, topk_ids)

531 +

532 + self.router = create_fused_moe_router(

533 + top_k=top_k,

534 + global_num_experts=self.global_num_experts,

535 + eplb_state=self.eplb_state,

536 + renormalize=renormalize,

537 + use_grouped_topk=use_grouped_topk,

538 + num_expert_group=num_expert_group,

539 + topk_group=topk_group,

540 + custom_routing_function=custom_routing_function,

541 + scoring_func=scoring_func,

542 + routed_scaling_factor=routed_scaling_factor,

543 + e_score_correction_bias=e_score_correction_bias,

544 + num_fused_shared_experts=self.num_fused_shared_experts,

545 + enable_eplb=enable_eplb,

546 + # TODO(bnell): once we can construct the MK at init time, we

547 + # can make this a value.

548 + indices_type_getter=lambda: self.quant_method.topk_indices_dtype,

549 + capture=capture,

550 + )

551 + self.routing_method_type: RoutingMethodType = self.router.routing_method_type

552 +

553 self.moe_config: FusedMoEConfig = FusedMoEConfig(

554 num_experts=self.global_num_experts,

555 experts_per_token=top_k,

556 hidden_dim=hidden_size,

557 + intermediate_size_per_partition=self.intermediate_size_per_partition,

558 num_local_experts=self.local_num_experts,

559 moe_parallel_config=self.moe_parallel_config,

560 in_dtype=moe_in_dtype,

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

563 has_bias=has_bias,

564 is_act_and_mul=is_act_and_mul,

565 is_lora_enabled=vllm_config.lora_config is not None,

566 + activation=activation,

567 + device=vllm_config.device_config.device,

568 + routing_method=self.routing_method_type,

569 )

570 self.moe_config_use_flashinfer_cutlass_kernels = (

571 self.moe_config.use_flashinfer_cutlass_kernels

@@ -594,39 +629,6 @@ class FusedMoE(CustomOp):

629 self.batched_hidden_states: torch.Tensor | None = None

630 self.batched_router_logits: torch.Tensor | None = None

631

597 - # TODO(bnell): in next PR move capture back to layer

598 - capture: Callable[[torch.Tensor], None] | None = None

599 - if (

600 - self.vllm_config.model_config is not None

601 - and self.vllm_config.model_config.enable_return_routed_experts

602 - ):

603 - # In dummy runs, the capturer is not initialized.

604 - capturer = RoutedExpertsCapturer.get_instance()

605 - if capturer is not None:

606 - capture = lambda topk_ids: capturer.capture(self.layer_id, topk_ids)

607 -

608 - self.router = create_fused_moe_router(

609 - top_k=top_k,

610 - global_num_experts=self.global_num_experts,

611 - eplb_state=self.eplb_state,

612 - renormalize=renormalize,

613 - use_grouped_topk=use_grouped_topk,

614 - num_expert_group=num_expert_group,

615 - topk_group=topk_group,

616 - custom_routing_function=custom_routing_function,

617 - scoring_func=scoring_func,

618 - routed_scaling_factor=routed_scaling_factor,

619 - e_score_correction_bias=e_score_correction_bias,

620 - num_fused_shared_experts=self.num_fused_shared_experts,

621 - enable_eplb=enable_eplb,

622 - # TODO(bnell): once we can construct the MK at init time, we

623 - # can make this a value.

624 - indices_type_getter=lambda: self.quant_method.topk_indices_dtype,

625 - routing_method_type=routing_method_type,

626 - capture=capture,

627 - )

628 - self.routing_method_type: RoutingMethodType = self.router.routing_method_type

629 -

632 # Note: maybe_init_modular_kernel should only be called by

633 # prepare_communication_buffer_for_model.

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

vllm/model_executor/layers/fused_moe/modular_kernel.py

+115 -6

@@ -13,6 +13,7 @@ import vllm.envs as envs

13 from vllm.forward_context import get_forward_context, is_forward_context_available

14 from vllm.logger import init_logger

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

16 + FusedMoEConfig,

17 FusedMoEParallelConfig,

18 FusedMoEQuantConfig,

19 )

@@ -22,6 +23,9 @@ from vllm.model_executor.layers.fused_moe.utils import (

23 count_expert_num_tokens,

24 disable_inplace,

25 )

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

27 + QuantKey,

28 +)

29 from vllm.utils.math_utils import cdiv

30 from vllm.v1.worker.ubatching import (

31 dbo_enabled,

@@ -374,18 +378,51 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

378

379 def __init__(

380 self,

381 + moe_config: FusedMoEConfig,

382 quant_config: FusedMoEQuantConfig,

383 + max_num_tokens: int | None = None,

384 + num_dispatchers: int | None = None,

385 ):

386 """

387 + moe_config: MoE layer configuration.

388 quant_config: Quantization parameters for this experts instance.

389 """

390 + if self.activation_format() == FusedMoEActivationFormat.Standard and (

391 + max_num_tokens is not None or num_dispatchers is not None

392 + ):

393 + raise ValueError(

394 + "max_num_tokens and num_dispatchers should only be set for "

395 + "BatchedExperts activation format."

396 + )

397 + elif self.activation_format() == FusedMoEActivationFormat.BatchedExperts and (

398 + max_num_tokens is None or num_dispatchers is None

399 + ):

400 + raise ValueError(

401 + "max_num_tokens and num_dispatchers must be set for "

402 + "BatchedExperts activation format."

403 + )

404 +

405 + self.moe_config = moe_config

406 self.quant_config = quant_config

407 + self.max_num_tokens = max_num_tokens

408 + self.num_dispatchers = num_dispatchers

409

384 - @property

410 + @staticmethod

411 + def expects_unquantized_inputs(

412 + moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig

413 + ) -> bool:

414 + """

415 + Whether or not the PrepareFinalize should defer input quantization

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

417 + execute the input quantization itself.

418 +

419 + Sample subclasses that override are AITER and FlashInfer CUTLASS.

420 + """

421 + return False

422 +

423 + @staticmethod

424 @abstractmethod

386 - def activation_formats(

387 - self,

388 - ) -> tuple[FusedMoEActivationFormat, FusedMoEActivationFormat]:

425 + def activation_format() -> FusedMoEActivationFormat:

426 """

427 A property which is a tuple of the input and output activation formats

428 for the 'apply' method.

@@ -435,6 +472,78 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

472

473 return E, M, N, K, topk

474

475 + #

476 + # Various helpers for registering support for various features.

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

478 + #

479 +

480 + @staticmethod

481 + def is_supported_config(

482 + cls: type["FusedMoEPermuteExpertsUnpermute"],

483 + moe_config: FusedMoEConfig,

484 + weight_key: QuantKey | None,

485 + activation_key: QuantKey | None,

486 + activation_format: FusedMoEActivationFormat,

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

488 + def _make_reason(reason: str) -> str:

489 + return f"kernel does not support {reason}"

490 +

491 + if not cls._supports_current_device():

492 + return False, _make_reason("current device")

493 + elif not (moe_config.is_act_and_mul or cls._supports_no_act_and_mul()):

494 + return False, _make_reason("no act_and_mul MLP layer")

495 + elif not cls._supports_activation(moe_config.activation):

496 + return False, _make_reason(f"{moe_config.activation} activation")

497 + elif not cls._supports_quant_scheme(weight_key, activation_key):

498 + return False, _make_reason("quantization scheme")

499 + elif not cls._supports_parallel_config(moe_config.moe_parallel_config):

500 + return False, _make_reason("parallel config")

501 + elif activation_format != cls.activation_format():

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

503 + return True, None

504 +

505 + @staticmethod

506 + @abstractmethod

507 + def _supports_current_device() -> bool:

508 + """

509 + Whether the kernel supports the current device type

510 + (compute cability and current platform).

511 + """

512 + raise NotImplementedError

513 +

514 + @staticmethod

515 + @abstractmethod

516 + def _supports_no_act_and_mul() -> bool:

517 + """

518 + Whether the kernel supports act_and_mul=False, i.e.

519 + non-gated MoE models like Nemotron-Nano.

520 + """

521 + raise NotImplementedError

522 +

523 + @staticmethod

524 + @abstractmethod

525 + def _supports_quant_scheme(

526 + weight_key: QuantKey | None,

527 + activation_key: QuantKey | None,

528 + ) -> bool:

529 + raise NotImplementedError

530 +

531 + @staticmethod

532 + @abstractmethod

533 + def _supports_activation(activation: str) -> bool:

534 + """

535 + Whether the kernel supports a particular act function.

536 + """

537 + raise NotImplementedError

538 +

539 + @staticmethod

540 + @abstractmethod

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

542 + """

543 + Whether the kernel supports deployment in expert parallel.

544 + """

545 + raise NotImplementedError

546 +

547 #

548 # Various helpers for accessing quantization parameters from the

549 # quant_config.

@@ -715,12 +824,12 @@ class FusedMoEModularKernel(torch.nn.Module):

824

825 self._post_init_setup()

826 assert (

718 - prepare_finalize.activation_format == fused_experts.activation_formats[0]

827 + prepare_finalize.activation_format == fused_experts.activation_format()

828 ), (

829 f"{prepare_finalize.__class__.__name__}."

830 f"{prepare_finalize.activation_format} == "

831 f"{fused_experts.__class__.__name__}."

723 - f"{fused_experts.activation_formats[0]}"

832 + f"{fused_experts.activation_format()}"

833 )

834

835 def _post_init_setup(self):

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

+350 -186

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

14 fp8_w8a8_moe_quant_config,

15 fp8_w8a16_moe_quant_config,

16 )

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

18 + is_supported_config_trtllm,

19 +)

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

21 + MoEPrepareAndFinalizeNoEP,

22 +)

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

24 FlashinferMoeBackend,

25 get_flashinfer_moe_backend,

@@ -26,133 +32,307 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (

32 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (

33 prepare_fp8_moe_layer_for_marlin,

34 )

29 -from vllm.model_executor.layers.quantization.utils.w8a8_utils import (

30 - cutlass_group_gemm_supported,

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

36 + QuantKey,

37 )

32 -from vllm.platforms import current_platform

33 -from vllm.utils.deep_gemm import is_deep_gemm_supported

34 -from vllm.utils.flashinfer import has_flashinfer_moe

35 -from vllm.utils.import_utils import has_deep_gemm

39 logger = init_logger(__name__)

42 class Fp8MoeBackend(Enum):

41 - NONE = 0

42 - FLASHINFER_TRTLLM = 1

43 - FLASHINFER_CUTLASS = 2

44 - DEEPGEMM = 3

45 - MARLIN = 4

46 - TRITON = 5

47 - AITER = 6

48 - VLLM_CUTLASS = 7

43 + NONE = "NONE"

44 + FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM"

45 + FLASHINFER_CUTLASS = "FLASHINFER_CUTLASS"

46 + DEEPGEMM = "DEEPGEMM"

47 + BATCHED_DEEPGEMM = "BATCHED_DEEPGEMM"

48 + MARLIN = "MARLIN"

49 + TRITON = "TRITON"

50 + BATCHED_TRITON = "BATCHED_TRITON"

51 + AITER = "AITER"

52 + VLLM_CUTLASS = "VLLM_CUTLASS"

53 + BATCHED_VLLM_CUTLASS = "BATCHED_VLLM_CUTLASS"

54 +

55 +

56 +def backend_to_kernel_cls(

57 + backend: Fp8MoeBackend,

58 +) -> type[mk.FusedMoEPermuteExpertsUnpermute]:

59 + if backend == Fp8MoeBackend.FLASHINFER_TRTLLM:

60 + raise NotImplementedError

61 +

62 + elif backend == Fp8MoeBackend.FLASHINFER_CUTLASS:

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

64 + FlashInferExperts,

65 + )

66 +

67 + return FlashInferExperts

68 +

69 + elif backend == Fp8MoeBackend.DEEPGEMM:

70 + from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (

71 + TritonOrDeepGemmExperts,

72 + )

73 +

74 + return TritonOrDeepGemmExperts

75 +

76 + elif backend == Fp8MoeBackend.BATCHED_DEEPGEMM:

77 + from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (

78 + BatchedDeepGemmExperts,

79 + )

80 +

81 + return BatchedDeepGemmExperts

82 +

83 + elif backend == Fp8MoeBackend.MARLIN:

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

85 + MarlinExperts,

86 + )

87 +

88 + return MarlinExperts

89 +

90 + elif backend == Fp8MoeBackend.TRITON:

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

92 + TritonExperts,

93 + )

94 +

95 + return TritonExperts

96 +

97 + elif backend == Fp8MoeBackend.BATCHED_TRITON:

98 + from vllm.model_executor.layers.fused_moe.fused_batched_moe import (

99 + BatchedTritonExperts,

100 + )

101 +

102 + return BatchedTritonExperts

103 +

104 + elif backend == Fp8MoeBackend.AITER:

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

106 + AiterExperts,

107 + )

108 +

109 + return AiterExperts

110 +

111 + elif backend == Fp8MoeBackend.VLLM_CUTLASS:

112 + from vllm.model_executor.layers.fused_moe.triton_cutlass_moe import (

113 + TritonOrCutlassExperts,

114 + )

115 +

116 + return TritonOrCutlassExperts

117 +

118 + elif backend == Fp8MoeBackend.BATCHED_VLLM_CUTLASS:

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

120 + CutlassBatchedExpertsFp8,

121 + )

122 +

123 + return CutlassBatchedExpertsFp8

124 +

125 + else:

126 + raise ValueError(f"Unknown FP8 MoE backend: {backend.value}")

127

128

129 def select_fp8_moe_backend(

52 - block_quant: bool,

53 - tp_size: int,

54 - with_lora_support: bool,

55 - is_act_and_mul: bool,

130 + config: FusedMoEConfig,

131 + weight_key: QuantKey | None,

132 + activation_key: QuantKey | None,

133 allow_vllm_cutlass: bool = False,

57 -) -> Fp8MoeBackend:

134 +) -> tuple[Fp8MoeBackend, type[mk.FusedMoEPermuteExpertsUnpermute] | None]:

135 """

136 Select the primary FP8 MoE backend

137 Note: Shape-specific fallbacks may still occur at runtime.

138 """

62 - # TODO(rob): in a future PR, we will query each mk for

63 - # supported features and return the mk directly, just like

64 - # we do for the Attention Backend.

139 + k_cls: type[mk.FusedMoEPermuteExpertsUnpermute] | None = None

140 +

141 + if config.is_lora_enabled:

142 + return Fp8MoeBackend.TRITON, backend_to_kernel_cls(Fp8MoeBackend.TRITON)

143 +

144 + # NOTE: the kernels are selected in the following order.

145 + AVAILABLE_BACKENDS = [

146 + Fp8MoeBackend.AITER,

147 + Fp8MoeBackend.FLASHINFER_TRTLLM,

148 + Fp8MoeBackend.FLASHINFER_CUTLASS,

149 + Fp8MoeBackend.DEEPGEMM,

150 + Fp8MoeBackend.BATCHED_DEEPGEMM,

151 + Fp8MoeBackend.VLLM_CUTLASS,

152 + Fp8MoeBackend.BATCHED_VLLM_CUTLASS,

153 + Fp8MoeBackend.TRITON,

154 + Fp8MoeBackend.BATCHED_TRITON,

155 + Fp8MoeBackend.MARLIN,

156 + ]

157 +

158 + # NOTE(rob): We need to peak into the P/F selection to determine

159 + # if we are using the batched or standard expert format, which

160 + # if not ideal. Once we unify TP + DP/EP, we can select P/F first.

161 + activation_format = (

162 + mk.FusedMoEActivationFormat.BatchedExperts

163 + if config.moe_parallel_config.use_batched_activation_format

164 + else mk.FusedMoEActivationFormat.Standard

165 + )

166

66 - if with_lora_support:

67 - return Fp8MoeBackend.TRITON

167 + def _make_log_backend(backend: Fp8MoeBackend):

168 + available_backend_strs = [b.value for b in AVAILABLE_BACKENDS]

169 + return (

170 + f"Using {backend.value} Fp8 MoE backend out "

171 + f"of potential backends: {available_backend_strs}."

172 + )

173

69 - def _make_log_backend(backend_name: str):

70 - return f"Using {backend_name} backend for FP8 MoE"

174 + def _make_log_unsupported(backend: Fp8MoeBackend, reason: str | None) -> str:

175 + if reason:

176 + return (

177 + f"FP8 MoE backend {backend.value} does not support the "

178 + f"deployment configuration since {reason}."

179 + )

180 + else:

181 + return (

182 + f"FP8 MoE backend '{backend.value}' does not support the "

183 + "deployment configuration."

184 + )

185

72 - # Prefer FlashInfer backends on supported GPUs; allow SM90 and SM100.

73 - if (

74 - current_platform.is_cuda()

75 - and (

76 - current_platform.is_device_capability_family(100)

77 - or current_platform.is_device_capability(90)

186 + def _return_or_raise(

187 + backend: Fp8MoeBackend,

188 + config: FusedMoEConfig,

189 + weight_key: QuantKey | None,

190 + activation_key: QuantKey | None,

191 + activation_format: mk.FusedMoEActivationFormat,

192 + ) -> tuple[Fp8MoeBackend, type[mk.FusedMoEPermuteExpertsUnpermute]]:

193 + k_cls = backend_to_kernel_cls(backend)

194 + supported, reason = k_cls.is_supported_config(

195 + k_cls, config, weight_key, activation_key, activation_format

196 )

79 - and envs.VLLM_USE_FLASHINFER_MOE_FP8

80 - and has_flashinfer_moe()

81 - ):

82 - backend = get_flashinfer_moe_backend()

83 - if backend == FlashinferMoeBackend.TENSORRT_LLM:

84 - logger.info_once(_make_log_backend("FlashInfer TRTLLM"))

85 - if not is_act_and_mul:

86 - raise ValueError(

87 - "FlashInfer TRTLLM FP8 MoE backend only supports "

88 - "act_and_mul gate_up_project fusion. Please set "

89 - "VLLM_USE_FLASHINFER_MOE_FP8=throughput to use the "

90 - "FlashInfer CUTLASS backend instead."

197 + if supported:

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

199 + return backend, k_cls

200 + raise ValueError(_make_log_unsupported(backend, reason))

201 +

202 + # Handle explicit FlashInfer FP8 configuration.

203 + if envs.is_set("VLLM_USE_FLASHINFER_MOE_FP8"):

204 + if not envs.VLLM_USE_FLASHINFER_MOE_FP8:

205 + # If the user rejects FlashInfer remove those backends.

206 + AVAILABLE_BACKENDS.remove(Fp8MoeBackend.FLASHINFER_TRTLLM)

207 + AVAILABLE_BACKENDS.remove(Fp8MoeBackend.FLASHINFER_CUTLASS)

208 +

209 + elif envs.is_set("VLLM_FLASHINFER_MOE_BACKEND"):

210 + # If user is explicit about backend, validate it.

211 + fi_backend = get_flashinfer_moe_backend()

212 +

213 + if fi_backend == FlashinferMoeBackend.TENSORRT_LLM:

214 + backend = Fp8MoeBackend.FLASHINFER_TRTLLM

215 + supported, reason = is_supported_config_trtllm(

216 + config, weight_key, activation_key, activation_format

217 )

92 - return Fp8MoeBackend.FLASHINFER_TRTLLM

93 - else:

94 - if block_quant and current_platform.is_device_capability_family(100):

95 - raise ValueError(

96 - "FlashInfer FP8 MoE throughput backend does not "

97 - "support block quantization on SM100. Please use "

98 - "VLLM_FLASHINFER_MOE_BACKEND=latency to use the "

99 - "FlashInfer TRTLLM backend instead."

218 + if supported:

219 + logger.info_once(_make_log_backend(backend))

220 + return backend, None

221 + else:

222 + raise ValueError(_make_log_unsupported(backend, reason))

223 +

224 + elif fi_backend == FlashinferMoeBackend.CUTLASS:

225 + backend = Fp8MoeBackend.FLASHINFER_CUTLASS

226 + return _return_or_raise(

227 + backend, config, weight_key, activation_key, activation_format

228 )

101 - logger.info_once(_make_log_backend("FlashInfer CUTLASS"))

102 - return Fp8MoeBackend.FLASHINFER_CUTLASS

229

104 - # weight-only path for older GPUs without native FP8

105 - if (

106 - current_platform.is_cuda() and not current_platform.has_device_capability(89)

107 - ) or envs.VLLM_TEST_FORCE_FP8_MARLIN:

108 - logger.info_once(_make_log_backend("Marlin"), scope="local")

109 - return Fp8MoeBackend.MARLIN

110 -

111 - # Determine if we should use DeepGEMM with block-quantized weights:

112 - # - If explicitly set by user, respect their choice

113 - # - If not explicitly set (default), disable when TP size is >= 8

114 - moe_use_deep_gemm = envs.VLLM_MOE_USE_DEEP_GEMM

115 - if not envs.is_set("VLLM_MOE_USE_DEEP_GEMM") and tp_size >= 8:

116 - moe_use_deep_gemm = False

117 - logger.info_once(

118 - "DeepGEMM MoE is disabled by default when TP size is >= 8. "

119 - "Set VLLM_MOE_USE_DEEP_GEMM=1 to enable it.",

120 - scope="local",

121 - )

230 + else:

231 + assert fi_backend == FlashinferMoeBackend.CUTEDSL

232 + raise ValueError("FlashInfer MaskedGEMM not supported for FP8")

233 +

234 + else:

235 + # If the user is not explicit about the backend, try both.

236 + for backend in [

237 + Fp8MoeBackend.FLASHINFER_TRTLLM,

238 + Fp8MoeBackend.FLASHINFER_CUTLASS,

239 + ]:

240 + if backend == Fp8MoeBackend.FLASHINFER_TRTLLM:

241 + k_cls = None

242 + supported, reason = is_supported_config_trtllm(

243 + config,

244 + weight_key,

245 + activation_key,

246 + activation_format,

247 + )

248 + else:

249 + k_cls = backend_to_kernel_cls(backend)

250 + supported, reason = k_cls.is_supported_config(

251 + k_cls,

252 + config,

253 + weight_key,

254 + activation_key,

255 + activation_format,

256 + )

257 +

258 + if supported:

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

260 + return backend, k_cls

261 + else:

262 + logger.debug_once(

263 + _make_log_unsupported(backend, reason), scope="local"

264 + )

265 +

266 + raise NotImplementedError(

267 + "Found VLLM_USE_FLASHINFER_MOE_FP8=1, but no "

268 + "FlashInfer FP8 MoE backend supports the configuration."

269 + )

270 +

271 + # Handle explicit DeepGEMM FP8 configuration.

272 + if envs.is_set("VLLM_USE_DEEP_GEMM") or envs.is_set("VLLM_MOE_USE_DEEP_GEMM"):

273 + if not envs.VLLM_USE_DEEP_GEMM or not envs.VLLM_MOE_USE_DEEP_GEMM:

274 + AVAILABLE_BACKENDS.remove(Fp8MoeBackend.DEEPGEMM)

275 + AVAILABLE_BACKENDS.remove(Fp8MoeBackend.BATCHED_DEEPGEMM)

276 + else:

277 + backend = (

278 + Fp8MoeBackend.DEEPGEMM

279 + if activation_format == mk.FusedMoEActivationFormat.Standard

280 + else Fp8MoeBackend.BATCHED_DEEPGEMM

281 + )

282 + return _return_or_raise(

283 + backend, config, weight_key, activation_key, activation_format

284 + )

285

123 - use_deep_gemm = envs.VLLM_USE_DEEP_GEMM

124 - if not is_deep_gemm_supported():

125 - use_deep_gemm = False

126 - logger.info_once(

127 - "DeepGEMM is disabled because the platform does not support it.",

128 - scope="local",

286 + # Handle explicit MARLIN FP8 configuration.

287 + if envs.VLLM_TEST_FORCE_FP8_MARLIN:

288 + backend = Fp8MoeBackend.MARLIN

289 + return _return_or_raise(

290 + backend, config, weight_key, activation_key, activation_format

291 )

292

131 - if use_deep_gemm and moe_use_deep_gemm and block_quant and is_act_and_mul:

132 - if not has_deep_gemm():

133 - logger.warning_once(

134 - "DeepGEMM backend requested but not available.", scope="local"

293 + # Handle explicit AITER FP8 configuration.

294 + if envs.is_set("VLLM_ROCM_USE_AITER") or envs.is_set("VLLM_ROCM_USE_AITER_MOE"):

295 + if not envs.VLLM_ROCM_USE_AITER or not envs.VLLM_ROCM_USE_AITER_MOE:

296 + AVAILABLE_BACKENDS.remove(Fp8MoeBackend.AITER)

297 + else:

298 + backend = Fp8MoeBackend.AITER

299 + return _return_or_raise(

300 + backend, config, weight_key, activation_key, activation_format

301 )

136 - elif is_deep_gemm_supported():

137 - logger.info_once(_make_log_backend("DeepGEMM"), scope="local")

138 - return Fp8MoeBackend.DEEPGEMM

302

140 - if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MOE:

141 - logger.info_once(_make_log_backend("ROCm AITER"), scope="local")

142 - return Fp8MoeBackend.AITER

303 + if not allow_vllm_cutlass:

304 + AVAILABLE_BACKENDS.remove(Fp8MoeBackend.VLLM_CUTLASS)

305 + AVAILABLE_BACKENDS.remove(Fp8MoeBackend.BATCHED_VLLM_CUTLASS)

306 +

307 + # Select kernels in order of backend.

308 + for backend in AVAILABLE_BACKENDS:

309 + if backend == Fp8MoeBackend.FLASHINFER_TRTLLM:

310 + k_cls = None

311 + supported, reason = is_supported_config_trtllm(

312 + config,

313 + weight_key,

314 + activation_key,

315 + activation_format,

316 + )

317 + else:

318 + k_cls = backend_to_kernel_cls(backend)

319 + supported, reason = k_cls.is_supported_config(

320 + k_cls,

321 + config,

322 + weight_key,

323 + activation_key,

324 + activation_format,

325 + )

326

144 - if (

145 - allow_vllm_cutlass

146 - and not block_quant

147 - and cutlass_group_gemm_supported()

148 - and is_act_and_mul

149 - ):

150 - logger.info_once(_make_log_backend("vLLM CUTLASS"), scope="local")

151 - return Fp8MoeBackend.VLLM_CUTLASS

327 + if supported:

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

329 + return backend, k_cls

330 + else:

331 + logger.debug_once(_make_log_unsupported(backend, reason), scope="local")

332

153 - # default to Triton

154 - logger.info_once(_make_log_backend("Triton"), scope="local")

155 - return Fp8MoeBackend.TRITON

333 + raise NotImplementedError(

334 + "No FP8 MoE backend supports the deployment configuration."

335 + )

336

337

338 def convert_to_fp8_moe_kernel_format(

@@ -166,7 +346,7 @@ def convert_to_fp8_moe_kernel_format(

346 w2_input_scale: torch.Tensor | None,

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

348 block_quant = hasattr(layer, "weight_block_size")

169 - if fp8_backend == Fp8MoeBackend.DEEPGEMM:

349 + if fp8_backend in [Fp8MoeBackend.DEEPGEMM, Fp8MoeBackend.BATCHED_DEEPGEMM]:

350 assert block_quant

351 w13, w2, w13_scale, w2_scale = prepare_fp8_moe_layer_for_deepgemm(

352 w13,

@@ -199,6 +379,14 @@ def convert_to_fp8_moe_kernel_format(

379 w2_input_scale=w2_input_scale,

380 is_trtllm=(fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM),

381 )

382 + else:

383 + if fp8_backend not in [

384 + Fp8MoeBackend.TRITON,

385 + Fp8MoeBackend.BATCHED_TRITON,

386 + Fp8MoeBackend.VLLM_CUTLASS,

387 + Fp8MoeBackend.BATCHED_VLLM_CUTLASS,

388 + ]:

389 + raise ValueError(f"Unsupported FP8 MoE backend: {fp8_backend.value}")

390

391 return w13, w2, w13_scale, w2_scale

392

@@ -210,6 +398,8 @@ def make_fp8_moe_quant_config(

398 a1_scale: torch.Tensor | None,

399 a2_scale: torch.Tensor | None,

400 block_shape: list[int] | None = None,

401 + per_act_token_quant: bool = False,

402 + per_out_ch_quant: bool = False,

403 ) -> FusedMoEQuantConfig | None:

404 """

405 Create FusedMoEQuantConfig for the specifed FP8 Backend.

@@ -262,102 +452,76 @@ def make_fp8_moe_quant_config(

452 a1_scale=a1_scale,

453 a2_scale=a2_scale,

454 block_shape=block_shape,

455 + per_act_token_quant=per_act_token_quant,

456 + per_out_ch_quant=per_out_ch_quant,

457 )

458

459

268 -def make_fp8_moe_kernel(

269 - layer: torch.nn.Module,

270 - moe_quant_config: FusedMoEQuantConfig,

460 +def make_fp8_moe_kernel_for_mkm(

461 moe_config: FusedMoEConfig,

272 - fp8_backend: Fp8MoeBackend,

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

274 - # Delayed import is required since the oracle is imported

275 - # by CPU backends which cannot import all of these experts.

276 - # TODO: update the experts to make this not happen.

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

278 - MoEPrepareAndFinalizeNoEP,

279 - )

280 -

281 - # NOTE(rob): this is a WIP refactor. We are first migrating

282 - # all of the kernels in the TP case to use mk. Once this is

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

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

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

286 - use_inplace = True

287 - if fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:

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

289 - FlashInferExperts,

462 + quant_config: FusedMoEQuantConfig,

463 + experts_cls: type[mk.FusedMoEPermuteExpertsUnpermute],

464 + prepare_finalize: mk.FusedMoEPrepareAndFinalize,

465 +) -> mk.FusedMoEPermuteExpertsUnpermute:

466 + if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts:

467 + max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank()

468 + assert max_num_tokens_per_rank is not None

469 + experts = experts_cls(

470 + moe_config=moe_config,

471 + quant_config=quant_config,

472 + max_num_tokens=max_num_tokens_per_rank,

473 + num_dispatchers=prepare_finalize.num_dispatchers(),

474 )

291 -

292 - kernel = mk.FusedMoEModularKernel(

293 - MoEPrepareAndFinalizeNoEP(

294 - defer_input_quant=moe_quant_config.is_block_quantized

295 - ),

296 - FlashInferExperts(

297 - out_dtype=layer.orig_dtype,

298 - quant_config=moe_quant_config,

299 - ep_rank=moe_config.ep_rank,

300 - ep_size=moe_config.ep_size,

301 - tp_rank=moe_config.tp_rank,

302 - tp_size=moe_config.tp_size,

303 - use_dp=(moe_config.dp_size > 1),

304 - use_deepseek_fp8_block_scale=moe_quant_config.is_block_quantized,

305 - ),

475 + else:

476 + experts = experts_cls(

477 + moe_config=moe_config,

478 + quant_config=quant_config,

479 )

307 - use_inplace = False

480

309 - elif fp8_backend == Fp8MoeBackend.AITER:

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

311 - AiterExperts,

312 - )

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

482 + return experts

483

314 - kernel = mk.FusedMoEModularKernel(

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

316 - MoEPrepareAndFinalizeNoEP(defer_input_quant=True),

317 - AiterExperts(quant_config=moe_quant_config),

318 - )

319 - elif fp8_backend == Fp8MoeBackend.MARLIN:

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

321 - MarlinExperts,

322 - )

484

324 - kernel = mk.FusedMoEModularKernel(

325 - MoEPrepareAndFinalizeNoEP(),

326 - MarlinExperts(quant_config=moe_quant_config),

327 - )

328 - elif fp8_backend == Fp8MoeBackend.VLLM_CUTLASS:

329 - from vllm.model_executor.layers.fused_moe.triton_cutlass_moe import (

330 - TritonOrCutlassExperts,

485 +def make_fp8_moe_kernel(

486 + moe_quant_config: FusedMoEQuantConfig,

487 + moe_config: FusedMoEConfig,

488 + fp8_backend: Fp8MoeBackend,

489 + experts_cls: type[mk.FusedMoEPermuteExpertsUnpermute],

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

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

492 + if (

493 + moe_config.moe_parallel_config.use_all2all_kernels

494 + and moe_config.moe_parallel_config.all2all_backend

495 + not in ["allgather_reducescatter", "naive"]

496 + ):

497 + raise ValueError(

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

499 + "This should happen via the ModularKernelMethod."

500 )

501

333 - kernel = mk.FusedMoEModularKernel(

334 - MoEPrepareAndFinalizeNoEP(),

335 - TritonOrCutlassExperts(

336 - out_dtype=moe_config.in_dtype,

337 - e=layer.local_num_experts,

338 - n=layer.intermediate_size_per_partition,

339 - k=layer.hidden_size,

340 - device=layer.w13_weight.device,

341 - quant_config=moe_quant_config,

342 - ),

343 - )

344 - elif fp8_backend == Fp8MoeBackend.DEEPGEMM:

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

346 - TritonOrDeepGemmExperts,

347 - )

502 + # Create Prepare/Finalize.

503 + prepare_finalize = MoEPrepareAndFinalizeNoEP(

504 + defer_input_quant=experts_cls.expects_unquantized_inputs(

505 + moe_config, moe_quant_config

506 + ),

507 + )

508

349 - kernel = mk.FusedMoEModularKernel(

350 - MoEPrepareAndFinalizeNoEP(),

351 - TritonOrDeepGemmExperts(quant_config=moe_quant_config),

352 - )

353 - else:

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

355 - TritonExperts,

356 - )

509 + # Create Experts.

510 + experts = experts_cls(

511 + moe_config=moe_config,

512 + quant_config=moe_quant_config,

513 + )

514

358 - assert fp8_backend == Fp8MoeBackend.TRITON

359 - kernel = mk.FusedMoEModularKernel(

360 - MoEPrepareAndFinalizeNoEP(),

361 - TritonExperts(quant_config=moe_quant_config),

362 - )

363 - return kernel, use_inplace

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

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

517 + # the new MoE runner class.

518 + kernel = mk.FusedMoEModularKernel(

519 + prepare_finalize,

520 + experts,

521 + shared_experts=None,

522 + moe_parallel_config=moe_config.moe_parallel_config,

523 + )

524 +

525 + # TODO(rob): update inplace logic to be part of the kernel.

526 + inplace = fp8_backend != Fp8MoeBackend.FLASHINFER_CUTLASS

527 + return kernel, inplace

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

+258 -91

@@ -14,21 +14,11 @@ from vllm.model_executor.layers.fused_moe.config import (

14 nvfp4_moe_quant_config,

15 nvfp4_w4a16_moe_quant_config,

16 )

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

18 - CutlassExpertsFp4,

19 -)

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

21 - FlashInferExperts,

22 -)

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

24 - MarlinExperts,

25 -)

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

18 MoEPrepareAndFinalizeNoEP,

19 )

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

30 - is_flashinfer_fp4_cutedsl_moe_available,

31 - is_flashinfer_fp4_cutlass_moe_available,

21 + 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 (

@@ -36,27 +26,26 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

26 get_flashinfer_moe_backend,

27 )

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

39 - is_fp4_marlin_supported,

29 prepare_nvfp4_moe_layer_for_marlin,

30 )

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

43 - cutlass_fp4_supported,

32 + QuantKey,

33 )

35 logger = init_logger(__name__)

38 class NvFp4MoeBackend(Enum):

50 - FLASHINFER_CUTLASS = "FlashInfer CUTLASS"

51 - FLASHINFER_TRTLLM = "FlashInfer TRTLLM"

52 - FLASHINFER_CUTEDSL = "FlashInfer CUTEDSL"

53 - VLLM_CUTLASS = "vLLM CUTASS"

54 - MARLIN = "vLLM MARLIN"

39 + FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM"

40 + FLASHINFER_CUTLASS = "FLASHINFER_CUTLASS"

41 + FLASHINFER_CUTEDSL = "FLASHINFER_CUTEDSL"

42 + VLLM_CUTLASS = "VLLM_CUTLASS"

43 + MARLIN = "MARLIN"

46 FLASHINFER_NVFP4_MOE_BACKENDS = [

58 - NvFp4MoeBackend.FLASHINFER_CUTLASS,

47 NvFp4MoeBackend.FLASHINFER_TRTLLM,

48 + NvFp4MoeBackend.FLASHINFER_CUTLASS,

49 NvFp4MoeBackend.FLASHINFER_CUTEDSL,

50 ]

@@ -72,44 +61,208 @@ def is_global_sf_supported_for_nvfp4_backend(backend: NvFp4MoeBackend) -> bool:

61 # of all experts in Expert Parallel Mode when all experts are not

62 # on the same rank.

75 - return backend in [

76 - NvFp4MoeBackend.FLASHINFER_CUTLASS,

64 + return backend in FLASHINFER_NVFP4_MOE_BACKENDS

65 +

66 +

67 +def backend_to_kernel_cls(

68 + backend: NvFp4MoeBackend,

69 +) -> type[mk.FusedMoEPermuteExpertsUnpermute]:

70 + if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM:

71 + raise NotImplementedError(

72 + "FLASHINFER_TRTLLM doesn't support Modular Kernel Interface"

73 + )

74 +

75 + elif backend == NvFp4MoeBackend.FLASHINFER_CUTLASS:

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

77 + FlashInferExperts,

78 + )

79 +

80 + return FlashInferExperts

81 +

82 + elif backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL:

83 + from vllm.model_executor.layers.fused_moe.flashinfer_cutedsl_moe import (

84 + FlashInferCuteDSLExperts,

85 + )

86 +

87 + return FlashInferCuteDSLExperts

88 +

89 + elif backend == NvFp4MoeBackend.VLLM_CUTLASS:

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

91 + CutlassExpertsFp4,

92 + )

93 +

94 + return CutlassExpertsFp4

95 +

96 + elif backend == NvFp4MoeBackend.MARLIN:

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

98 + MarlinExperts,

99 + )

100 +

101 + return MarlinExperts

102 + else:

103 + raise ValueError(f"Unknown NvFP4 MoE backend: {backend.value}")

104 +

105 +

106 +def select_nvfp4_moe_backend(

107 + config: FusedMoEConfig,

108 + weight_key: QuantKey | None,

109 + activation_key: QuantKey | None,

110 +) -> tuple[NvFp4MoeBackend, type[mk.FusedMoEPermuteExpertsUnpermute] | None]:

111 + """

112 + Select the primary NvFP4 MoE backend

113 + Note: Shape-specific fallbacks may still occur at runtime.

114 + """

115 +

116 + # NOTE: the kernels are selected in the following order.

117 + AVAILABLE_BACKENDS = [

118 NvFp4MoeBackend.FLASHINFER_TRTLLM,

119 + NvFp4MoeBackend.FLASHINFER_CUTEDSL,

120 + NvFp4MoeBackend.FLASHINFER_CUTLASS,

121 + NvFp4MoeBackend.VLLM_CUTLASS,

122 + NvFp4MoeBackend.MARLIN,

123 ]

124

125 + # NOTE(rob): this is kind of a hack. We need to peak into

126 + # the prepare-finalize selection to determine if we are using

127 + # the batched or standard expert format.

128 + use_batched = (

129 + config.moe_parallel_config.use_deepep_ll_kernels

130 + or config.moe_parallel_config.use_pplx_kernels

131 + )

132 + activation_format = (

133 + mk.FusedMoEActivationFormat.BatchedExperts

134 + if use_batched

135 + else mk.FusedMoEActivationFormat.Standard

136 + )

137

81 -def select_nvfp4_moe_backend() -> NvFp4MoeBackend:

138 def _make_log_backend(backend: NvFp4MoeBackend):

83 - return f"Using {backend.value} backend for NvFp4 MoE"

139 + available_backend_strs = [b.value for b in AVAILABLE_BACKENDS]

140 + return (

141 + f"Using '{backend.value}' NvFp4 MoE backend out "

142 + f"of potential backends: {available_backend_strs}."

143 + )

144

85 - if cutlass_fp4_supported() and not envs.VLLM_TEST_FORCE_FP8_MARLIN:

86 - allow_flashinfer = (

87 - is_flashinfer_fp4_cutlass_moe_available()

88 - or is_flashinfer_fp4_cutedsl_moe_available()

145 + def _make_log_unsupported(backend: NvFp4MoeBackend, reason: str | None) -> str:

146 + if reason:

147 + return (

148 + f"NvFp4 MoE backend '{backend.value}' does not support the "

149 + f"deployment configuration since {reason}."

150 + )

151 + else:

152 + return (

153 + f"NvFp4 MoE backend '{backend.value}' does not support the "

154 + "deployment configuration."

155 + )

156 +

157 + def _return_or_raise(

158 + backend: NvFp4MoeBackend,

159 + config: FusedMoEConfig,

160 + weight_key: QuantKey | None,

161 + activation_key: QuantKey | None,

162 + activation_format: mk.FusedMoEActivationFormat,

163 + ) -> tuple[NvFp4MoeBackend, type[mk.FusedMoEPermuteExpertsUnpermute]]:

164 + k_cls = backend_to_kernel_cls(backend)

165 + supported, reason = k_cls.is_supported_config(

166 + k_cls, config, weight_key, activation_key, activation_format

167 )

90 - if allow_flashinfer and envs.VLLM_USE_FLASHINFER_MOE_FP4:

91 - backend = fi_2_vllm_backend_map[get_flashinfer_moe_backend()]

168 + if supported:

169 + logger.info_once(_make_log_backend(backend))

170 + return backend, k_cls

171 + raise ValueError(_make_log_unsupported(backend, reason))

172 +

173 + if envs.is_set("VLLM_USE_FLASHINFER_MOE_FP4"):

174 + if not envs.VLLM_USE_FLASHINFER_MOE_FP4:

175 + # If the user rejects FlashInfer remove those backends.

176 + for b in FLASHINFER_NVFP4_MOE_BACKENDS:

177 + AVAILABLE_BACKENDS.remove(b)

178 +

179 + elif envs.is_set("VLLM_FLASHINFER_MOE_BACKEND"):

180 + # If user is explicit about backend, validate it.

181 + fi_backend = get_flashinfer_moe_backend()

182 +

183 + if fi_backend == FlashinferMoeBackend.TENSORRT_LLM:

184 + backend = NvFp4MoeBackend.FLASHINFER_TRTLLM

185 + supported, reason = is_supported_config_trtllm(

186 + config, weight_key, activation_key, activation_format

187 + )

188 + if supported:

189 + logger.info_once(_make_log_backend(backend))

190 + return backend, None

191 + else:

192 + raise ValueError(_make_log_unsupported(backend, reason))

193 + else:

194 + backend = fi_2_vllm_backend_map[fi_backend]

195 + return _return_or_raise(

196 + backend, config, weight_key, activation_key, activation_format

197 + )

198 else:

93 - backend = NvFp4MoeBackend.VLLM_CUTLASS

94 - elif is_fp4_marlin_supported():

199 + # If the user is not explicit about the backend, try each.

200 + for backend in FLASHINFER_NVFP4_MOE_BACKENDS:

201 + if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM:

202 + k_cls = None

203 + supported, reason = is_supported_config_trtllm(

204 + config,

205 + weight_key,

206 + activation_key,

207 + activation_format,

208 + )

209 + else:

210 + k_cls = backend_to_kernel_cls(backend)

211 + supported, reason = k_cls.is_supported_config(

212 + k_cls,

213 + config,

214 + weight_key,

215 + activation_key,

216 + activation_format,

217 + )

218 + if supported:

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

220 + return backend, None

221 + else:

222 + logger.debug_once(

223 + _make_log_unsupported(backend, reason), scope="local"

224 + )

225 +

226 + raise NotImplementedError(

227 + "Found VLLM_USE_FLASHINFER_MOE_FP4=1, but no "

228 + "FlashInfer NVFP4 MoE backend supports the configuration."

229 + )

230 +

231 + if envs.VLLM_TEST_FORCE_FP8_MARLIN:

232 backend = NvFp4MoeBackend.MARLIN

96 - else:

97 - raise ValueError("No NvFp4 kernel backend available for NvFp4 MoE.")

98 -

99 - # Log warning if FI backend requested but not available.

100 - if (

101 - backend not in FLASHINFER_NVFP4_MOE_BACKENDS

102 - and envs.VLLM_USE_FLASHINFER_MOE_FP4

103 - ):

104 - logger.warning_once(

105 - "Requested FlashInfer backend for NvFp4 MoE, but it's not available. "

106 - "Falling back to %s for NvFp4 MoE",

107 - backend.value,

108 - scope="local",

233 + return _return_or_raise(

234 + backend, config, weight_key, activation_key, activation_format

235 )

110 - else:

111 - logger.info_once(_make_log_backend(backend), scope="local")

112 - return backend

236 +

237 + # Select kernels in order of backend.

238 + for backend in AVAILABLE_BACKENDS:

239 + if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM:

240 + k_cls = None # type: ignore[assignment]

241 + supported, reason = is_supported_config_trtllm(

242 + config,

243 + weight_key,

244 + activation_key,

245 + activation_format,

246 + )

247 + else:

248 + k_cls = backend_to_kernel_cls(backend)

249 + supported, reason = k_cls.is_supported_config(

250 + k_cls,

251 + config,

252 + weight_key,

253 + activation_key,

254 + activation_format,

255 + )

256 +

257 + if supported:

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

259 + return backend, k_cls

260 + else:

261 + logger.debug_once(_make_log_unsupported(backend, reason), scope="local")

262 +

263 + raise NotImplementedError(

264 + "No NvFp4 MoE backend supports the deployment configuration."

265 + )

266

267

268 def convert_to_nvfp4_moe_kernel_format(

@@ -238,55 +391,69 @@ def make_nvfp4_moe_quant_config(

391 )

392

393

241 -def make_nvfp4_moe_kernel(

242 - backend: NvFp4MoeBackend,

243 - quant_config: FusedMoEQuantConfig,

394 +def make_nvfp4_moe_kernel_for_mkm(

395 moe_config: FusedMoEConfig,

245 -) -> mk.FusedMoEModularKernel | None:

246 - assert moe_config.dp_size == 1

396 + quant_config: FusedMoEQuantConfig,

397 + experts_cls: type[mk.FusedMoEPermuteExpertsUnpermute],

398 + prepare_finalize: mk.FusedMoEPrepareAndFinalize,

399 +) -> mk.FusedMoEPermuteExpertsUnpermute:

400 + if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts:

401 + max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank()

402 + assert max_num_tokens_per_rank is not None

403 + experts = experts_cls(

404 + moe_config=moe_config,

405 + quant_config=quant_config,

406 + max_num_tokens=max_num_tokens_per_rank,

407 + num_dispatchers=prepare_finalize.num_dispatchers(),

408 + )

409 + else:

410 + experts = experts_cls(

411 + moe_config=moe_config,

412 + quant_config=quant_config,

413 + )

414

248 - UNSUPPORTED_BACKENDS = [

249 - # TRTLLM does not use the modular kernl abstraction.

250 - NvFp4MoeBackend.FLASHINFER_TRTLLM,

251 - # CUTEDSL is used with BATCHED (masked) format only.

252 - # TODO: add here once we support dp/ep via the oracle.

253 - NvFp4MoeBackend.FLASHINFER_CUTEDSL,

254 - ]

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

416 + return experts

417

256 - if backend in UNSUPPORTED_BACKENDS:

257 - return None

418

259 - elif backend == NvFp4MoeBackend.FLASHINFER_CUTLASS:

260 - return mk.FusedMoEModularKernel(

261 - MoEPrepareAndFinalizeNoEP(defer_input_quant=True),

262 - FlashInferExperts(

263 - out_dtype=moe_config.in_dtype,

264 - quant_config=quant_config,

265 - ep_rank=moe_config.ep_rank,

266 - ep_size=moe_config.ep_size,

267 - tp_rank=moe_config.tp_rank,

268 - tp_size=moe_config.tp_size,

269 - use_dp=False,

270 - use_deepseek_fp8_block_scale=False,

271 - ),

419 +def make_nvfp4_moe_kernel(

420 + moe_quant_config: FusedMoEQuantConfig,

421 + moe_config: FusedMoEConfig,

422 + experts_cls: type[mk.FusedMoEPermuteExpertsUnpermute],

423 +) -> mk.FusedMoEModularKernel:

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

425 + if (

426 + moe_config.moe_parallel_config.use_all2all_kernels

427 + and moe_config.moe_parallel_config.all2all_backend

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

429 + ):

430 + raise ValueError(

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

432 + "This should happen via the ModularKernelMethod."

433 )

434

274 - elif backend == NvFp4MoeBackend.VLLM_CUTLASS:

275 - return mk.FusedMoEModularKernel(

276 - MoEPrepareAndFinalizeNoEP(defer_input_quant=True),

277 - CutlassExpertsFp4(

278 - out_dtype=moe_config.in_dtype,

279 - # TODO(rob): see what impact this has on expert map?

280 - max_experts_per_worker=moe_config.num_experts,

281 - quant_config=quant_config,

282 - ),

283 - )

435 + # Create Prepare/Finalize.

436 + prepare_finalize = MoEPrepareAndFinalizeNoEP(

437 + defer_input_quant=experts_cls.expects_unquantized_inputs(

438 + moe_config, moe_quant_config

439 + ),

440 + )

441

285 - elif backend == NvFp4MoeBackend.MARLIN:

286 - return mk.FusedMoEModularKernel(

287 - MoEPrepareAndFinalizeNoEP(),

288 - MarlinExperts(quant_config=quant_config),

289 - )

442 + # Create Experts.

443 + experts = experts_cls(

444 + moe_config=moe_config,

445 + quant_config=moe_quant_config,

446 + )

447

291 - else:

292 - raise ValueError(f"Unknown NvFp4 MoE backend: {backend}")

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

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

450 + # the new MoE runner class.

451 + kernel = mk.FusedMoEModularKernel(

452 + prepare_finalize,

453 + experts,

454 + shared_experts=None,

455 + moe_parallel_config=moe_config.moe_parallel_config,

456 + )

457 +

458 + # TODO(rob): update inplace logic to be part of the kernel.

459 + return kernel

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

+9 -8

@@ -123,7 +123,6 @@ def convert_to_unquantized_kernel_format(

123

124

125 def make_unquantized_moe_kernel(

126 - layer: torch.nn.Module,

126 backend: UnquantizedMoeBackend,

127 quant_config: FusedMoEQuantConfig,

128 moe_config: FusedMoEConfig,

@@ -141,12 +140,8 @@ def make_unquantized_moe_kernel(

140 kernel = mk.FusedMoEModularKernel(

141 MoEPrepareAndFinalizeNoEP(),

142 FlashInferExperts(

144 - out_dtype=layer.params_dtype,

143 + moe_config=moe_config,

144 quant_config=quant_config,

146 - tp_rank=moe_config.moe_parallel_config.tp_rank,

147 - tp_size=moe_config.moe_parallel_config.tp_size,

148 - ep_rank=moe_config.moe_parallel_config.ep_rank,

149 - ep_size=moe_config.moe_parallel_config.ep_size,

145 ),

146 )

147 use_inplace = False

@@ -157,13 +152,19 @@ def make_unquantized_moe_kernel(

152

153 kernel = mk.FusedMoEModularKernel(

154 MoEPrepareAndFinalizeNoEP(),

160 - AiterExperts(quant_config),

155 + AiterExperts(

156 + moe_config=moe_config,

157 + quant_config=quant_config,

158 + ),

159 )

160 elif backend == UnquantizedMoeBackend.TRITON:

161 from vllm.model_executor.layers.fused_moe import TritonExperts

162

163 kernel = mk.FusedMoEModularKernel(

164 MoEPrepareAndFinalizeNoEP(),

167 - TritonExperts(quant_config),

165 + TritonExperts(

166 + moe_config=moe_config,

167 + quant_config=quant_config,

168 + ),

169 )

170 return kernel, use_inplace

vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py

+52 -10

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

9 from vllm._aiter_ops import rocm_aiter_ops

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

11 FUSED_MOE_UNQUANTIZED_CONFIG,

12 + FusedMoEParallelConfig,

13 FusedMoEQuantConfig,

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.quant_utils import (

19 + QuantKey,

20 + kFp8Dynamic128Sym,

21 + kFp8DynamicTensorSym,

22 + kFp8DynamicTokenSym,

23 + kFp8Static128BlockSym,

24 + kFp8StaticChannelSym,

25 + kFp8StaticTensorSym,

26 +)

29 class QuantMethod(IntEnum):

@@ -269,17 +279,49 @@ def rocm_aiter_fused_experts(

279

280

281 class AiterExperts(mk.FusedMoEPermuteExpertsUnpermute):

272 - def __init__(self, quant_config):

273 - super().__init__(quant_config)

282 + @staticmethod

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

284 + return mk.FusedMoEActivationFormat.Standard

285 +

286 + @staticmethod

287 + def expects_unquantized_inputs(

288 + fused_moe_config: mk.FusedMoEConfig, quant_config: FusedMoEQuantConfig

289 + ) -> bool:

290 + # AITER fused MoE kernels handle input quantization internally.

291 + return True

292

275 - @property

276 - def activation_formats(

277 - self,

278 - ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:

279 - return (

280 - mk.FusedMoEActivationFormat.Standard,

281 - mk.FusedMoEActivationFormat.Standard,

282 - )

293 + @staticmethod

294 + def _supports_current_device() -> bool:

295 + return rocm_aiter_ops.is_fused_moe_enabled()

296 +

297 + @staticmethod

298 + def _supports_no_act_and_mul() -> bool:

299 + return False

300 +

301 + @staticmethod

302 + def _supports_quant_scheme(

303 + weight_key: QuantKey | None,

304 + activation_key: QuantKey | None,

305 + ) -> bool:

306 + # TODO(rob): AITER also supports MXFP4, which is not

307 + # yet supported via an Oracle. Once it is, we will add

308 + # MXFP4 to this list.

309 + SUPPORTED_W_A = [

310 + (None, None),

311 + (kFp8Static128BlockSym, kFp8Dynamic128Sym),

312 + (kFp8StaticTensorSym, kFp8StaticTensorSym),

313 + (kFp8StaticTensorSym, kFp8DynamicTensorSym),

314 + (kFp8StaticChannelSym, kFp8DynamicTokenSym),

315 + ]

316 + return (weight_key, activation_key) in SUPPORTED_W_A

317 +

318 + @staticmethod

319 + def _supports_activation(activation: str) -> bool:

320 + return activation in ["silu", "gelu"]

321 +

322 + @staticmethod

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

324 + return True

325

326 def supports_expert_map(self):

327 return True

vllm/model_executor/layers/fused_moe/router/custom_routing_router.py

@@ -34,6 +34,11 @@ class CustomRoutingRouter(BaseRouter):

35 @property

36 def routing_method_type(self) -> RoutingMethodType:

37 + from vllm.model_executor.models.llama4 import Llama4MoE

38 +

39 + # NOTE: FLASHINFER_TRTLLM support the Llama4 router.

40 + if self.custom_routing_function == Llama4MoE.custom_routing_function:

41 + return RoutingMethodType.Llama4

42 return RoutingMethodType.Custom

44 def _compute_routing(

vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py

+4 -6

@@ -261,7 +261,6 @@ class GroupedTopKRouter(BaseRouter):

261 num_fused_shared_experts: int = 0,

262 enable_eplb: bool = False,

263 indices_type_getter: Callable[[], torch.dtype | None] | None = None,

264 - routing_method_type: RoutingMethodType | None = None,

264 ):

265 super().__init__(

266 top_k=top_k,

@@ -278,13 +277,12 @@ class GroupedTopKRouter(BaseRouter):

277 self.e_score_correction_bias = e_score_correction_bias

278 self.num_fused_shared_experts = num_fused_shared_experts

279

281 - # Determine routing method type

282 - if routing_method_type is not None:

283 - self._routing_method_type = routing_method_type

284 - elif scoring_func == "sigmoid":

280 + if scoring_func == "sigmoid":

281 self._routing_method_type = RoutingMethodType.DeepSeekV3

282 else:

287 - self._routing_method_type = RoutingMethodType.TopK

283 + # NOTE: this prohibits the FLASHINFER_TRTLLM kernels from

284 + # being selected, since they only support DeepSeek-style.

285 + self._routing_method_type = RoutingMethodType.Unspecified

286

287 @property

288 def routing_method_type(self) -> RoutingMethodType:

vllm/model_executor/layers/fused_moe/router/router_factory.py

-3

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

7 import vllm.envs as envs

8 from vllm.distributed.eplb.eplb_state import EplbLayerState

9 -from vllm.model_executor.layers.fused_moe.config import RoutingMethodType

9 from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter

10 from vllm.model_executor.layers.fused_moe.router.custom_routing_router import (

11 CustomRoutingRouter,

@@ -36,7 +35,6 @@ def create_fused_moe_router(

35 global_num_experts: int,

36 renormalize: bool = True,

37 indices_type_getter: Callable[[], torch.dtype | None] | None = None,

39 - routing_method_type: RoutingMethodType | None = None,

38 # grouped topk parameters

39 use_grouped_topk: bool = False,

40 num_expert_group: int | None = None,

@@ -128,7 +126,6 @@ def create_fused_moe_router(

126 num_fused_shared_experts=num_fused_shared_experts,

127 enable_eplb=enable_eplb,

128 indices_type_getter=indices_type_getter,

131 - routing_method_type=routing_method_type,

129 )

130 router.capture = capture

131 return router

vllm/model_executor/layers/fused_moe/triton_cutlass_moe.py

+14 -8

@@ -5,7 +5,10 @@

5 import torch

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

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

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

9 + FusedMoEConfig,

10 + FusedMoEQuantConfig,

11 +)

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

13 from vllm.model_executor.layers.fused_moe.fallback import FallbackExperts

14 from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts

@@ -17,19 +20,22 @@ class TritonOrCutlassExperts(FallbackExperts):

21 def __init__(

22 self,

20 - e: int,

21 - n: int,

22 - k: int,

23 - out_dtype: torch.dtype | None,

23 + moe_config: FusedMoEConfig,

24 quant_config: FusedMoEQuantConfig,

25 - device: torch.dtype,

25 ):

26 self.is_sm100 = current_platform.has_device_capability(100)

27 super().__init__(

29 - experts=CutlassExpertsFp8(e, n, k, out_dtype, quant_config, device),

30 - fallback_experts=TritonExperts(quant_config),

28 + experts=CutlassExpertsFp8(moe_config, quant_config),

29 + fallback_experts=TritonExperts(moe_config, quant_config),

30 )

32 + @staticmethod

33 + def get_clses() -> tuple[

34 + type[mk.FusedMoEPermuteExpertsUnpermute],

35 + type[mk.FusedMoEPermuteExpertsUnpermute],

36 + ]:

37 + return (CutlassExpertsFp8, TritonExperts)

38 +

39 def workspace_shapes(

40 self,

41 M: int,

vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py

+14 -4

@@ -4,7 +4,10 @@

4 import torch

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

7 -from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

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

8 + FusedMoEConfig,

9 + FusedMoEQuantConfig,

10 +)

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

12 DeepGemmExperts,

13 _valid_deep_gemm,

@@ -20,12 +23,19 @@ from vllm.utils.deep_gemm import (

23 class TritonOrDeepGemmExperts(FallbackExperts):

24 """DeepGemm with fallback to Triton for low latency shapes."""

23 - def __init__(self, quant_config: FusedMoEQuantConfig):

26 + def __init__(self, moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig):

27 super().__init__(

25 - experts=DeepGemmExperts(quant_config),

26 - fallback_experts=TritonExperts(quant_config),

28 + experts=DeepGemmExperts(moe_config, quant_config),

29 + fallback_experts=TritonExperts(moe_config, quant_config),

30 )

32 + @staticmethod

33 + def get_clses() -> tuple[

34 + type[mk.FusedMoEPermuteExpertsUnpermute],

35 + type[mk.FusedMoEPermuteExpertsUnpermute],

36 + ]:

37 + return (DeepGemmExperts, TritonExperts)

38 +

39 def workspace_shapes(

40 self,

41 M: int,

vllm/model_executor/layers/fused_moe/trtllm_moe.py

+47 -11

@@ -6,37 +6,73 @@ import torch

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

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

8 FusedMoEConfig,

9 + FusedMoEParallelConfig,

10 FusedMoEQuantConfig,

11 )

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

13 TopKWeightAndReduceNoOP,

14 )

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

16 + QuantKey,

17 +)

20 class TrtLlmGenExperts(mk.FusedMoEPermuteExpertsUnpermute):

21 def __init__(

22 self,

19 - moe: FusedMoEConfig,

23 + moe_config: FusedMoEConfig,

24 quant_config: FusedMoEQuantConfig,

25 gemm1_alpha,

26 gemm1_beta,

27 gemm1_clamp_limit,

28 max_capture_size,

29 ):

26 - super().__init__(quant_config)

27 - self.moe = moe

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

31 self.gemm1_alpha = gemm1_alpha

32 self.gemm1_beta = gemm1_beta

33 self.gemm1_clamp_limit = gemm1_clamp_limit

34 self.max_capture_size = max_capture_size

33 - @property

34 - def activation_formats(

35 - self,

36 - ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:

37 - return (

38 - mk.FusedMoEActivationFormat.Standard,

39 - mk.FusedMoEActivationFormat.Standard,

36 + @staticmethod

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

38 + return mk.FusedMoEActivationFormat.Standard

39 +

40 + @staticmethod

41 + def _supports_current_device() -> bool:

42 + raise NotImplementedError(

43 + "TrtLlmGenExperts is not yet used by an Oracle. "

44 + "This method should not be called."

45 + )

46 +

47 + @staticmethod

48 + def _supports_no_act_and_mul() -> bool:

49 + raise NotImplementedError(

50 + "TrtLlmGenExperts is not yet used by an Oracle. "

51 + "This method should not be called."

52 + )

53 +

54 + @staticmethod

55 + def _supports_quant_scheme(

56 + weight_key: QuantKey | None,

57 + activation_key: QuantKey | None,

58 + ) -> bool:

59 + raise NotImplementedError(

60 + "TrtLlmGenExperts is not yet used by an Oracle. "

61 + "This method should not be called."

62 + )

63 +

64 + @staticmethod

65 + def _supports_activation(activation: str) -> bool:

66 + raise NotImplementedError(

67 + "TrtLlmGenExperts is not yet used by an Oracle. "

68 + "This method should not be called."

69 + )

70 +

71 + @staticmethod

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

73 + raise NotImplementedError(

74 + "TrtLlmGenExperts is not yet used by an Oracle. "

75 + "This method should not be called."

76 )

78 def supports_chunking(self) -> bool:

@@ -86,7 +122,7 @@ class TrtLlmGenExperts(mk.FusedMoEPermuteExpertsUnpermute):

122 topk = topk_ids.size(-1)

123 local_num_experts = w1.size(0)

124 intermediate_size = w2.size(1)

89 - local_expert_offset = self.moe.ep_rank * local_num_experts

125 + local_expert_offset = self.moe_config.ep_rank * local_num_experts

126

127 x_quant = hidden_states

128 x_scale = a1q_scale

vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py

+6 -3

@@ -96,13 +96,17 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

96 ):

97 logger.debug("BatchedTritonExperts %s", self.moe)

98 return BatchedTritonExperts(

99 + moe_config=self.moe,

100 + quant_config=self.moe_quant_config,

101 max_num_tokens=self.moe.max_num_tokens,

102 num_dispatchers=prepare_finalize.num_dispatchers(),

101 - quant_config=self.moe_quant_config,

103 )

104 else:

105 logger.debug("TritonExperts %s", self.moe)

105 - return TritonExperts(self.moe_quant_config)

106 + return TritonExperts(

107 + moe_config=self.moe,

108 + quant_config=self.moe_quant_config,

109 + )

110

111 def create_weights(

112 self,

@@ -192,7 +196,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

196 assert self.moe_quant_config is not None

197

198 self.kernel, self.use_inplace = make_unquantized_moe_kernel(

195 - layer=layer,

199 backend=self.unquantized_backend,

200 quant_config=self.moe_quant_config,

201 moe_config=self.moe,

vllm/model_executor/layers/quantization/awq_marlin.py

@@ -739,6 +739,7 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):

739 return BatchedMarlinExperts(

740 max_num_tokens=max_num_tokens_per_rank,

741 num_dispatchers=prepare_finalize.num_dispatchers(),

742 + moe_config=self.moe,

743 quant_config=self.moe_quant_config,

744 w13_g_idx=w13_g_idx,

745 w2_g_idx=w2_g_idx,

@@ -749,6 +750,7 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):

750 else:

751 # Standard Marlin experts for AWQ

752 return MarlinExperts(

753 + moe_config=self.moe,

754 quant_config=self.moe_quant_config,

755 w13_g_idx=w13_g_idx,

756 w2_g_idx=w2_g_idx,

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

+194 -185

@@ -19,7 +19,6 @@ from vllm.logger import init_logger

19 from vllm.model_executor.layers.fused_moe import (

20 FusedMoE,

21 FusedMoEActivationFormat,

22 - FusedMoEConfig,

22 FusedMoEMethodBase,

23 FusedMoEPermuteExpertsUnpermute,

24 FusedMoERouter,

@@ -27,9 +26,9 @@ from vllm.model_executor.layers.fused_moe import (

26 UnquantizedFusedMoEMethod,

27 )

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

29 + FusedMoEConfig,

30 FusedMoEQuantConfig,

31 - fp8_w8a8_moe_quant_config,

32 - fp8_w8a16_moe_quant_config,

31 + RoutingMethodType,

32 int4_w4a16_moe_quant_config,

33 int4_w4afp8_moe_quant_config,

34 int8_w8a8_moe_quant_config,

@@ -45,15 +44,17 @@ from vllm.model_executor.layers.fused_moe.oracle.fp8 import (

44 Fp8MoeBackend,

45 convert_to_fp8_moe_kernel_format,

46 make_fp8_moe_kernel,

47 + make_fp8_moe_kernel_for_mkm,

48 + make_fp8_moe_quant_config,

49 select_fp8_moe_backend,

50 )

51 from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (

51 - FLASHINFER_NVFP4_MOE_BACKENDS,

52 NvFp4MoeBackend,

53 convert_to_nvfp4_moe_kernel_format,

54 is_global_sf_supported_for_nvfp4_backend,

55 make_mxfp4_moe_quant_config,

56 make_nvfp4_moe_kernel,

57 + make_nvfp4_moe_kernel_for_mkm,

58 make_nvfp4_moe_quant_config,

59 select_nvfp4_moe_backend,

60 )

@@ -62,10 +63,12 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compress

63 WNA16_SUPPORTED_TYPES_MAP,

64 )

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

65 - build_flashinfer_fp4_cutlass_moe_prepare_finalize,

66 flashinfer_trtllm_fp4_moe,

67 flashinfer_trtllm_fp4_routed_moe,

68 - select_nvfp4_gemm_impl,

68 +)

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

70 + apply_fi_trtllm_fp8_per_tensor_moe,

71 + build_flashinfer_fp8_cutlass_moe_prepare_finalize,

72 )

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

74 process_fp8_input_tensor_strategy_moe,

@@ -79,12 +82,18 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (

82 marlin_moe_permute_scales,

83 )

84 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (

82 - is_fp4_marlin_supported,

85 prepare_moe_fp4_layer_for_marlin,

86 )

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

88 convert_bf16_scales_to_fp8,

89 convert_packed_uint4b8_to_signed_int4_inplace,

90 + kFp8Dynamic128Sym,

91 + kFp8DynamicTokenSym,

92 + kFp8Static128BlockSym,

93 + kFp8StaticChannelSym,

94 + kFp8StaticTensorSym,

95 + kNvfp4Dynamic,

96 + kNvfp4Static,

97 )

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

99 normalize_e4m3fn_to_e4m3fnuz,

@@ -200,7 +209,7 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):

209 f"or None for NVFP4A16, found {input_quant}",

210 )

211 return CompressedTensorsW4A4Nvfp4MoEMethod(

203 - layer.moe_config, layer_name, use_marlin=input_quant is None

212 + layer.moe_config, layer_name, use_a16=(input_quant is None)

213 )

214 elif (

215 quant_config._is_fp8_w8a8_sm90(weight_quant, input_quant)

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

243 super().__init__(moe)

244 self.group_size = 32

245 self.mxfp4_backend = NvFp4MoeBackend.MARLIN

246 + self.experts_cls = MarlinExperts

247 self.kernel: mk.FusedMoEModularKernel | None = None

248

249 def create_weights(

@@ -327,9 +337,9 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):

337 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

338 if self.moe_quant_config is not None:

339 self.kernel = make_nvfp4_moe_kernel(

330 - backend=self.mxfp4_backend,

331 - quant_config=self.moe_quant_config,

340 + moe_quant_config=self.moe_quant_config,

341 moe_config=self.moe,

342 + experts_cls=self.experts_cls,

343 )

344

345 def apply(

@@ -368,34 +378,30 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

378 self,

379 moe: FusedMoEConfig,

380 layer_name: str | None = None,

371 - use_marlin: bool = False,

381 + use_a16: bool = False,

382 ):

383 super().__init__(moe)

384 self.group_size = 16

375 - if use_marlin:

376 - if is_fp4_marlin_supported():

377 - self.nvfp4_backend = NvFp4MoeBackend.MARLIN

378 - else:

379 - raise ValueError(

380 - "Marlin FP4 MoE kernel requested but not ",

381 - "supported on current platform.",

382 - )

383 - else:

384 - self.nvfp4_backend = select_nvfp4_moe_backend()

385

386 - # TODO: move this type of check into the oracle.

387 - if not self.moe.is_act_and_mul and self.nvfp4_backend not in [

388 - NvFp4MoeBackend.FLASHINFER_CUTLASS,

389 - NvFp4MoeBackend.MARLIN,

390 - ]:

391 - raise NotImplementedError(

392 - "Non-gated activations are only supported by FlashInfer "

393 - f"CUTLASS and Marlin NvFP4 MoE backends, not {self.nvfp4_backend}."

394 - )

386 + # Select experts implementation.

387 + self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(

388 + config=self.moe,

389 + weight_key=kNvfp4Static,

390 + activation_key=None if use_a16 else kNvfp4Dynamic,

391 + )

392 +

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

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

395 +

396 self.use_global_sf = is_global_sf_supported_for_nvfp4_backend(

397 self.nvfp4_backend

398 )

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

399 +

400 + @property

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

402 + if self.kernel is not None:

403 + return self.kernel.prepare_finalize.topk_indices_dtype()

404 + return None

405

406 def create_weights(

407 self,

@@ -571,35 +577,40 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

577 layer.w13_input_scale = a13_scale

578 layer.w2_input_scale = a2_scale

579

574 - # Initialize the kernel that will be called in apply().

580 + # Setup modular kernel for TP case and naive DP/EP case.

581 + # In non-naive DP/EP case, we will create a ModularKernelMethod.

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

583 + # in both cases.

584 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

576 - use_dp = self.moe.dp_size > 1

577 - if self.moe_quant_config is not None and not use_dp:

585 + if self.moe_quant_config and (

586 + (not self.moe.moe_parallel_config.use_all2all_kernels)

587 + or self.moe.moe_parallel_config.use_naive_all2all_kernels

588 + ):

589 + assert self.experts_cls is not None

590 self.kernel = make_nvfp4_moe_kernel(

579 - backend=self.nvfp4_backend,

580 - quant_config=self.moe_quant_config,

591 + moe_quant_config=self.moe_quant_config,

592 moe_config=self.moe,

593 + experts_cls=self.experts_cls,

594 )

595

596 def maybe_make_prepare_finalize(

597 self,

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

599 ) -> mk.FusedMoEPrepareAndFinalize | None:

588 - UNSUPPORTED = [NvFp4MoeBackend.MARLIN, NvFp4MoeBackend.FLASHINFER_TRTLLM]

589 - if self.nvfp4_backend in UNSUPPORTED:

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

601 return None

602 elif self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_CUTLASS:

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

593 - if self.moe.dp_size == 1:

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

604 + if not self.moe.moe_parallel_config.use_all2all_kernels:

605 return None

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

596 - prepare_finalize = build_flashinfer_fp4_cutlass_moe_prepare_finalize(

597 - self.moe

606 +

607 + prepare_finalize = build_flashinfer_fp8_cutlass_moe_prepare_finalize(

608 + self.moe,

609 + use_deepseek_fp8_block_scale=False,

610 )

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

612 return prepare_finalize

601 - else:

602 - return super().maybe_make_prepare_finalize(routing_tables)

613 + return super().maybe_make_prepare_finalize(routing_tables)

614

615 def select_gemm_impl(

616 self,

@@ -607,14 +618,13 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

618 layer: torch.nn.Module,

619 ) -> mk.FusedMoEPermuteExpertsUnpermute:

620 assert self.moe_quant_config is not None

610 - """Return the appropriate GEMM experts implementation."""

611 - experts = select_nvfp4_gemm_impl(

612 - self.moe,

613 - self.moe_quant_config,

614 - allow_flashinfer=(self.nvfp4_backend in FLASHINFER_NVFP4_MOE_BACKENDS),

621 + assert self.experts_cls is not None

622 + return make_nvfp4_moe_kernel_for_mkm(

623 + moe_config=self.moe,

624 + quant_config=self.moe_quant_config,

625 + experts_cls=self.experts_cls,

626 + prepare_finalize=prepare_finalize,

627 )

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

617 - return experts

628

629 def get_fused_moe_quant_config(

630 self, layer: torch.nn.Module

@@ -727,33 +737,41 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

737 "For FP8 Fused MoE layer, we require either per tensor or "

738 "channelwise, dynamic per token quantization."

739 )

730 - self.fp8_backend = select_fp8_moe_backend(

731 - block_quant=self.block_quant,

732 - tp_size=moe.tp_size,

733 - with_lora_support=moe.is_lora_enabled,

734 - is_act_and_mul=moe.is_act_and_mul,

735 - # TODO(rob): enable selecting this externally.

740 +

741 + ct2vllm_weight = {

742 + QuantizationStrategy.CHANNEL: kFp8StaticChannelSym,

743 + QuantizationStrategy.TENSOR: kFp8StaticTensorSym,

744 + QuantizationStrategy.BLOCK: kFp8Static128BlockSym,

745 + }

746 + ct2vllm_act = {

747 + QuantizationStrategy.TOKEN: kFp8DynamicTokenSym,

748 + QuantizationStrategy.TENSOR: (

749 + kFp8StaticTensorSym if self.static_input_scales else kFp8Dynamic128Sym

750 + ),

751 + }

752 + weight_key = ct2vllm_weight[self.weight_quant.strategy]

753 + if weight_key == kFp8Static128BlockSym:

754 + activation_key = kFp8Dynamic128Sym

755 + else:

756 + activation_key = ct2vllm_act[self.input_quant.strategy]

757 +

758 + # Select Fp8 MoE backend

759 + self.fp8_backend, self.experts_cls = select_fp8_moe_backend(

760 + config=self.moe,

761 + weight_key=weight_key,

762 + activation_key=activation_key,

763 allow_vllm_cutlass=True,

764 )

738 - if self.fp8_backend != Fp8MoeBackend.MARLIN:

739 - per_act_token = self.input_quant.strategy == QuantizationStrategy.TOKEN

740 - per_channel_quant = (

741 - self.weight_quant.strategy == QuantizationStrategy.CHANNEL

742 - )

743 - if per_act_token != per_channel_quant:

744 - raise NotImplementedError(

745 - "For FP8 Fused MoE layers, per-token and per-channel must be "

746 - "used together."

747 - )

748 - # TODO(rob): hook this up in a follow up PR.

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

750 - raise NotImplementedError(

751 - "FlashInfer TRTLLM backend not supported for compressed-tensors yet."

752 - )

753 - self.disable_expert_map = False

765

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

767 self.kernel: mk.FusedMoEModularKernel | None = None

768

769 + @property

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

771 + if self.kernel is not None:

772 + return self.kernel.prepare_finalize.topk_indices_dtype()

773 + return None

774 +

775 def create_weights(

776 self,

777 layer: torch.nn.Module,

@@ -970,140 +988,75 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

988 replace_parameter(layer, "w13_weight_scale", w13_scale)

989 replace_parameter(layer, "w2_weight_scale", w2_scale)

990

991 + # Setup modular kernel for TP case and naive DP/EP case.

992 + # In non-naive DP/EP case, we will create a ModularKernelMethod.

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

994 + # in both cases.

995 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

974 - if self.moe_quant_config:

996 + if self.moe_quant_config and (

997 + (not self.moe.moe_parallel_config.use_all2all_kernels)

998 + or self.moe.moe_parallel_config.use_naive_all2all_kernels

999 + ):

1000 + assert self.experts_cls is not None

1001 self.kernel, self.use_inplace = make_fp8_moe_kernel(

976 - layer=layer,

1002 moe_quant_config=self.moe_quant_config,

1003 moe_config=self.moe,

1004 fp8_backend=self.fp8_backend,

1005 + experts_cls=self.experts_cls,

1006 )

1007

1008 def maybe_make_prepare_finalize(

1009 self,

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

1011 ) -> mk.FusedMoEPrepareAndFinalize | None:

986 - if self.fp8_backend in [Fp8MoeBackend.MARLIN, Fp8MoeBackend.AITER]:

1012 + if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:

1013 return None

988 - else:

989 - return super().maybe_make_prepare_finalize(routing_tables)

1014 + elif self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:

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

1016 + if not self.moe.moe_parallel_config.use_all2all_kernels:

1017 + return None

1018 +

1019 + prepare_finalize = build_flashinfer_fp8_cutlass_moe_prepare_finalize(

1020 + self.moe,

1021 + use_deepseek_fp8_block_scale=self.block_quant,

1022 + )

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

1024 + return prepare_finalize

1025 + return super().maybe_make_prepare_finalize(routing_tables)

1026

1027 def select_gemm_impl(

1028 self,

1029 prepare_finalize: mk.FusedMoEPrepareAndFinalize,

1030 layer: torch.nn.Module,

1031 ) -> FusedMoEPermuteExpertsUnpermute:

996 - # cutlass path

1032 assert self.moe_quant_config is not None

998 - if self.fp8_backend == Fp8MoeBackend.VLLM_CUTLASS:

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

1000 - CutlassBatchedExpertsFp8,

1001 - CutlassExpertsFp8,

1002 - )

1003 -

1004 - experts: FusedMoEPermuteExpertsUnpermute

1005 -

1006 - num_dispatchers = prepare_finalize.num_dispatchers()

1007 -

1008 - if (

1009 - prepare_finalize.activation_format

1010 - == FusedMoEActivationFormat.BatchedExperts

1011 - ):

1012 - logger.debug("CutlassBatchedExpertsFp8(%s)", self.__class__.__name__)

1013 - experts = CutlassBatchedExpertsFp8(

1014 - max_experts_per_worker=self.moe.num_local_experts,

1015 - num_dispatchers=num_dispatchers,

1016 - out_dtype=self.moe.in_dtype,

1017 - e=layer.local_num_experts,

1018 - n=layer.intermediate_size_per_partition,

1019 - k=layer.hidden_size,

1020 - device=layer.w13_weight.device,

1021 - quant_config=self.moe_quant_config,

1022 - )

1023 - else:

1024 - logger.debug("CutlassExpertsFp8(%s)", self.__class__.__name__)

1025 - experts = CutlassExpertsFp8(

1026 - out_dtype=self.moe.in_dtype,

1027 - e=layer.local_num_experts,

1028 - n=layer.intermediate_size_per_partition,

1029 - k=layer.hidden_size,

1030 - device=layer.w13_weight.device,

1031 - quant_config=self.moe_quant_config,

1032 - )

1033 -

1034 - # TODO(rob): investigate disable_expert_map

1035 - self.disable_expert_map = (

1036 - num_dispatchers > 1 or not experts.supports_expert_map()

1037 - )

1038 -

1039 - return experts

1040 -

1041 - from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (

1042 - BatchedDeepGemmExperts,

1043 - )

1044 - from vllm.model_executor.layers.fused_moe.fused_batched_moe import (

1045 - BatchedTritonExperts,

1046 - )

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

1048 - TritonExperts,

1049 - )

1050 - from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (

1051 - TritonOrDeepGemmExperts,

1033 + assert self.experts_cls is not None

1034 + return make_fp8_moe_kernel_for_mkm(

1035 + moe_config=self.moe,

1036 + quant_config=self.moe_quant_config,

1037 + experts_cls=self.experts_cls,

1038 + prepare_finalize=prepare_finalize,

1039 )

1040

1054 - assert self.fp8_backend not in [Fp8MoeBackend.AITER, Fp8MoeBackend.MARLIN]

1055 -

1056 - if (

1057 - prepare_finalize.activation_format

1058 - == FusedMoEActivationFormat.BatchedExperts

1059 - ):

1060 - max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank()

1061 - assert max_num_tokens_per_rank is not None

1062 -

1063 - if self.fp8_backend == Fp8MoeBackend.DEEPGEMM:

1064 - logger.debug("BatchedDeepGemmExperts(%s)", self.__class__.__name__)

1065 - return BatchedDeepGemmExperts(

1066 - max_num_tokens=max_num_tokens_per_rank,

1067 - num_dispatchers=prepare_finalize.num_dispatchers(),

1068 - quant_config=self.moe_quant_config,

1069 - )

1070 - else:

1071 - logger.debug("BatchedTritonExperts(%s)", self.__class__.__name__)

1072 - return BatchedTritonExperts(

1073 - max_num_tokens=max_num_tokens_per_rank,

1074 - num_dispatchers=prepare_finalize.num_dispatchers(),

1075 - quant_config=self.moe_quant_config,

1076 - )

1077 -

1078 - else:

1079 - if self.fp8_backend == Fp8MoeBackend.DEEPGEMM:

1080 - logger.debug("TritonOrDeepGemmExperts(%s)", self.__class__.__name__)

1081 - return TritonOrDeepGemmExperts(self.moe_quant_config)

1082 - else:

1083 - logger.debug("TritonExperts(%s)", self.__class__.__name__)

1084 - return TritonExperts(self.moe_quant_config)

1085 -

1041 def get_fused_moe_quant_config(

1042 self, layer: torch.nn.Module

1043 ) -> FusedMoEQuantConfig | None:

1089 - if self.fp8_backend == Fp8MoeBackend.MARLIN:

1090 - return fp8_w8a16_moe_quant_config(

1091 - w1_scale=layer.w13_weight_scale,

1092 - w2_scale=layer.w2_weight_scale,

1093 - block_shape=self.weight_block_size,

1094 - )

1095 -

1096 - per_act_token = self.input_quant.strategy == QuantizationStrategy.TOKEN

1097 - per_channel_quant = self.weight_quant.strategy == QuantizationStrategy.CHANNEL

1044 + w1_scale = layer.w13_weight_scale

1045 + w2_scale = layer.w2_weight_scale

1046 + a1_scale = layer.w13_input_scale

1047 + a2_scale = layer.w2_input_scale

1048

1099 - return fp8_w8a8_moe_quant_config(

1100 - w1_scale=layer.w13_weight_scale,

1101 - w2_scale=layer.w2_weight_scale,

1102 - a1_scale=layer.w13_input_scale,

1103 - a2_scale=layer.w2_input_scale,

1104 - per_act_token_quant=per_act_token,

1105 - per_out_ch_quant=per_channel_quant,

1106 - block_shape=layer.weight_block_size,

1049 + return make_fp8_moe_quant_config(

1050 + fp8_backend=self.fp8_backend,

1051 + w1_scale=w1_scale,

1052 + w2_scale=w2_scale,

1053 + a1_scale=a1_scale,

1054 + a2_scale=a2_scale,

1055 + per_act_token_quant=(

1056 + self.input_quant.strategy == QuantizationStrategy.TOKEN

1057 + ),

1058 + per_out_ch_quant=(self.input_quant.strategy == QuantizationStrategy.TOKEN),

1059 + block_shape=self.weight_block_size,

1060 )

1061

1062 def apply(

@@ -1113,6 +1066,56 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

1066 x: torch.Tensor,

1067 router_logits: torch.Tensor,

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

1069 + if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:

1070 + if layer.enable_eplb:

1071 + raise NotImplementedError(

1072 + "EPLB not supported for `FlashInfer TRTLLM FP8 MoE`."

1073 + )

1074 + assert layer.activation == "silu"

1075 +

1076 + if self.block_quant:

1077 + import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401

1078 +

1079 + e_score_correction_bias = (

1080 + layer.e_score_correction_bias.to(x.dtype)

1081 + if layer.e_score_correction_bias is not None

1082 + else None

1083 + )

1084 + routing_method_type = layer.routing_method_type

1085 + return torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8(

1086 + routing_logits=router_logits.to(torch.float32)

1087 + if routing_method_type == RoutingMethodType.DeepSeekV3

1088 + else router_logits,

1089 + routing_bias=e_score_correction_bias,

1090 + x=x,

1091 + w13_weight=layer.w13_weight,

1092 + w13_weight_scale_inv=layer.w13_weight_scale,

1093 + w2_weight=layer.w2_weight,

1094 + w2_weight_scale_inv=layer.w2_weight_scale,

1095 + global_num_experts=layer.global_num_experts,

1096 + top_k=layer.top_k,

1097 + num_expert_group=layer.num_expert_group,

1098 + topk_group=layer.topk_group,

1099 + intermediate_size=layer.intermediate_size_per_partition,

1100 + expert_offset=layer.ep_rank * layer.local_num_experts,

1101 + local_num_experts=layer.local_num_experts,

1102 + block_shape=self.weight_block_size,

1103 + routing_method_type=routing_method_type,

1104 + routed_scaling=layer.routed_scaling_factor,

1105 + )

1106 + else:

1107 + return apply_fi_trtllm_fp8_per_tensor_moe(

1108 + layer=layer,

1109 + hidden_states=x,

1110 + router_logits=router_logits,

1111 + routing_bias=layer.e_score_correction_bias,

1112 + global_num_experts=layer.global_num_experts,

1113 + top_k=layer.top_k,

1114 + num_expert_group=layer.num_expert_group,

1115 + topk_group=layer.topk_group,

1116 + apply_router_weight_on_input=layer.apply_router_weight_on_input,

1117 + )

1118 +

1119 topk_weights, topk_ids = router.select_experts(

1120 hidden_states=x,

1121 router_logits=router_logits,

@@ -1130,7 +1133,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

1133 global_num_experts=layer.global_num_experts,

1134 # TODO(rob): investigate the disable_expert_map introduced by:

1135 # https://github.com/vllm-project/vllm/commit/84166fee9770e6fba71a96978b3e7d149392fb28 # noqa: E501

1133 - expert_map=None if self.disable_expert_map else layer.expert_map,

1136 + expert_map=layer.expert_map,

1137 apply_router_weight_on_input=layer.apply_router_weight_on_input,

1138 )

1139

@@ -1596,6 +1599,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

1599 return BatchedMarlinExperts(

1600 max_num_tokens=max_num_tokens_per_rank,

1601 num_dispatchers=prepare_finalize.num_dispatchers(),

1602 + moe_config=self.moe,

1603 quant_config=self.moe_quant_config,

1604 w13_g_idx=layer.w13_weight_g_idx,

1605 w2_g_idx=layer.w2_weight_g_idx,

@@ -1605,6 +1609,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

1609 )

1610 else:

1611 return MarlinExperts(

1612 + moe_config=self.moe,

1613 quant_config=self.moe_quant_config,

1614 w13_g_idx=layer.w13_weight_g_idx,

1615 w2_g_idx=layer.w2_weight_g_idx,

@@ -1854,7 +1859,9 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):

1859

1860 layer.w13_weight = layer.w13_weight_packed

1861 layer.w2_weight = layer.w2_weight_packed

1857 - return TritonWNA16Experts(quant_config=self.moe_quant_config)

1862 + return TritonWNA16Experts(

1863 + moe_config=self.moe, quant_config=self.moe_quant_config

1864 + )

1865 else:

1866 raise NotImplementedError(

1867 "TritonExperts requires Triton. "

@@ -2467,6 +2474,7 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):

2474 c_strides2=self.a_strides1_c_strides2,

2475 s_strides1=self.s_strides1,

2476 s_strides2=self.s_strides2,

2477 + moe_config=self.moe,

2478 quant_config=self.moe_quant_config,

2479 group_size=self.group_size,

2480 )

@@ -2505,6 +2513,7 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):

2513 layer.w2_weight_packed,

2514 topk_weights,

2515 topk_ids,

2516 + moe_config=self.moe,

2517 quant_config=self.moe_quant_config,

2518 activation=layer.activation,

2519 global_num_experts=layer.global_num_experts,

vllm/model_executor/layers/quantization/fp8.py

+51 -104

@@ -19,7 +19,6 @@ from vllm.model_executor.layers.batch_invariant import (

19 )

20 from vllm.model_executor.layers.fused_moe import (

21 FusedMoE,

22 - FusedMoEActivationFormat,

22 FusedMoEMethodBase,

23 FusedMoEPermuteExpertsUnpermute,

24 FusedMoEPrepareAndFinalize,

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

34 Fp8MoeBackend,

35 convert_to_fp8_moe_kernel_format,

36 make_fp8_moe_kernel,

37 + make_fp8_moe_kernel_for_mkm,

38 make_fp8_moe_quant_config,

39 select_fp8_moe_backend,

40 )

@@ -55,7 +55,6 @@ from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod

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

56 apply_fi_trtllm_fp8_per_tensor_moe,

57 build_flashinfer_fp8_cutlass_moe_prepare_finalize,

58 - select_cutlass_fp8_gemm_impl,

58 )

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

60 W8A8BlockFp8LinearOp,

@@ -79,8 +78,10 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (

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

79 GroupShape,

80 is_layer_skipped,

81 + kFp8Dynamic128Sym,

82 kFp8DynamicTensorSym,

83 kFp8DynamicTokenSym,

84 + kFp8Static128BlockSym,

85 kFp8StaticTensorSym,

86 )

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

@@ -658,38 +659,36 @@ class Fp8MoEMethod(FusedMoEMethodBase):

659 self.weight_scale_name = (

660 "weight_scale_inv" if self.block_quant else "weight_scale"

661 )

661 - self.fp8_backend = select_fp8_moe_backend(

662 - block_quant=self.block_quant,

663 - tp_size=layer.moe_parallel_config.tp_size,

664 - with_lora_support=self.moe.is_lora_enabled,

665 - is_act_and_mul=self.moe.is_act_and_mul,

666 - )

662

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

669 - if self.block_quant and self.weight_block_size != [128, 128]:

670 - raise NotImplementedError(

671 - "FlashInfer CUTLASS FP8 MoE backend only supports block "

672 - "size [128, 128]."

673 - )

674 - if layer.activation != "silu":

675 - raise NotImplementedError(

676 - "FlashInfer CUTLASS FP8 MoE backend only supports SiLU "

677 - "activation function, but got {layer.activation}."

678 - )

679 - dynamic_per_token = (

680 - not self.block_quant and self.quant_config.activation_scheme != "static"

681 - )

682 - if dynamic_per_token and self.fp8_backend in [

683 - Fp8MoeBackend.FLASHINFER_TRTLLM,

684 - Fp8MoeBackend.FLASHINFER_CUTLASS,

685 - ]:

686 - raise NotImplementedError(

687 - "FlashInfer FP8 MoE backend does not support dynamic per token "

688 - "activation quantization."

663 + # Set weight key and activation key for kernel compatibility

664 + if self.block_quant:

665 + weight_key = kFp8Static128BlockSym

666 + activation_key = kFp8Dynamic128Sym

667 + else:

668 + weight_key = kFp8StaticTensorSym

669 + activation_key = (

670 + kFp8StaticTensorSym

671 + if self.quant_config.activation_scheme == "static"

672 + else kFp8DynamicTensorSym

673 )

674

675 + # Select Fp8 MoE backend

676 + self.fp8_backend, self.experts_cls = select_fp8_moe_backend(

677 + config=self.moe,

678 + weight_key=weight_key,

679 + activation_key=activation_key,

680 + allow_vllm_cutlass=False,

681 + )

682 +

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

684 self.kernel: mk.FusedMoEModularKernel | None = None

685

686 + @property

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

688 + if self.kernel is not None:

689 + return self.kernel.prepare_finalize.topk_indices_dtype()

690 + return None

691 +

692 def create_weights(

693 self,

694 layer: Module,

@@ -842,14 +841,21 @@ class Fp8MoEMethod(FusedMoEMethodBase):

841 replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_scale)

842 replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_scale)

843

845 - # Setup modular kernel for TP case.

844 + # Setup modular kernel for TP case and naive DP/EP case.

845 + # In non-naive DP/EP case, we will create a ModularKernelMethod.

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

847 + # in both cases.

848 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

847 - if self.moe_quant_config:

849 + if self.moe_quant_config and (

850 + (not self.moe.moe_parallel_config.use_all2all_kernels)

851 + or self.moe.moe_parallel_config.use_naive_all2all_kernels

852 + ):

853 + assert self.experts_cls is not None

854 self.kernel, self.use_inplace = make_fp8_moe_kernel(

849 - layer=layer,

855 moe_quant_config=self.moe_quant_config,

856 moe_config=self.moe,

857 fp8_backend=self.fp8_backend,

858 + experts_cls=self.experts_cls,

859 )

860

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

@@ -904,13 +910,13 @@ class Fp8MoEMethod(FusedMoEMethodBase):

910 self,

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

912 ) -> mk.FusedMoEPrepareAndFinalize | None:

907 - if self.fp8_backend in [

908 - Fp8MoeBackend.AITER,

909 - Fp8MoeBackend.MARLIN,

910 - Fp8MoeBackend.FLASHINFER_TRTLLM,

911 - ]:

913 + if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:

914 return None

915 elif self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:

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

917 + if not self.moe.moe_parallel_config.use_all2all_kernels:

918 + return None

919 +

920 prepare_finalize = build_flashinfer_fp8_cutlass_moe_prepare_finalize(

921 self.moe,

922 use_deepseek_fp8_block_scale=self.block_quant,

@@ -924,73 +930,14 @@ class Fp8MoEMethod(FusedMoEMethodBase):

930 prepare_finalize: FusedMoEPrepareAndFinalize,

931 layer: torch.nn.Module,

932 ) -> FusedMoEPermuteExpertsUnpermute:

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

928 - BatchedDeepGemmExperts,

929 - BatchedTritonExperts,

930 - TritonExperts,

931 - TritonOrDeepGemmExperts,

932 - )

933 -

934 - if self.fp8_backend in [Fp8MoeBackend.MARLIN, Fp8MoeBackend.AITER]:

935 - raise NotImplementedError(

936 - "Marlin and ROCm AITER are not supported with all2all yet."

937 - )

938 -

933 assert self.moe_quant_config is not None

940 -

941 - if (

942 - prepare_finalize.activation_format

943 - == FusedMoEActivationFormat.BatchedExperts

944 - ):

945 - max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank()

946 - assert max_num_tokens_per_rank is not None

947 -

948 - experts_impl = (

949 - BatchedDeepGemmExperts

950 - if self.fp8_backend == Fp8MoeBackend.DEEPGEMM

951 - else BatchedTritonExperts

952 - )

953 - logger.debug(

954 - "%s(%s): max_tokens_per_rank=%s, block_size=%s, per_act_token=%s",

955 - experts_impl.__name__,

956 - self.__class__.__name__,

957 - max_num_tokens_per_rank,

958 - self.weight_block_size,

959 - False,

960 - )

961 - return experts_impl(

962 - max_num_tokens=max_num_tokens_per_rank,

963 - num_dispatchers=prepare_finalize.num_dispatchers(),

964 - quant_config=self.moe_quant_config,

965 - )

966 - elif self.moe.is_lora_enabled:

967 - return TritonExperts(quant_config=self.moe_quant_config)

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

969 - # Select GEMM experts with block-scale when weights are block-quantized

970 - experts = select_cutlass_fp8_gemm_impl(

971 - self.moe,

972 - self.moe_quant_config,

973 - use_deepseek_fp8_block_scale=self.block_quant,

974 - )

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

976 - return experts

977 - elif self.fp8_backend == Fp8MoeBackend.DEEPGEMM:

978 - logger.debug(

979 - "TritonOrDeepGemmExperts(%s): block_size=%s, per_act_token=%s",

980 - self.__class__.__name__,

981 - self.weight_block_size,

982 - False,

983 - )

984 - return TritonOrDeepGemmExperts(self.moe_quant_config)

985 - else:

986 - assert self.fp8_backend == Fp8MoeBackend.TRITON

987 - logger.debug(

988 - "TritonExperts(%s): block_size=%s, per_act_token=%s",

989 - self.__class__.__name__,

990 - self.weight_block_size,

991 - False,

992 - )

993 - return TritonExperts(self.moe_quant_config)

934 + assert self.experts_cls is not None

935 + return make_fp8_moe_kernel_for_mkm(

936 + moe_config=self.moe,

937 + quant_config=self.moe_quant_config,

938 + experts_cls=self.experts_cls,

939 + prepare_finalize=prepare_finalize,

940 + )

941

942 def get_fused_moe_quant_config(

943 self, layer: torch.nn.Module

@@ -1067,7 +1014,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1014 routed_scaling=layer.routed_scaling_factor,

1015 )

1016 else:

1070 - result = apply_fi_trtllm_fp8_per_tensor_moe(

1017 + return apply_fi_trtllm_fp8_per_tensor_moe(

1018 layer=layer,

1019 hidden_states=x,

1020 router_logits=router_logits,

vllm/model_executor/layers/quantization/gptq_marlin.py

@@ -875,6 +875,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):

875 return BatchedMarlinExperts(

876 max_num_tokens=max_num_tokens_per_rank,

877 num_dispatchers=prepare_finalize.num_dispatchers(),

878 + moe_config=self.moe,

879 quant_config=self.moe_quant_config,

880 w13_g_idx=w13_g_idx,

881 w2_g_idx=w2_g_idx,

@@ -885,6 +886,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):

886 else:

887 # Standard Marlin experts for GPTQ

888 return MarlinExperts(

889 + moe_config=self.moe,

890 quant_config=self.moe_quant_config,

891 w13_g_idx=w13_g_idx,

892 w2_g_idx=w2_g_idx,

vllm/model_executor/layers/quantization/modelopt.py

+64 -42

@@ -27,15 +27,16 @@ from vllm.model_executor.layers.fused_moe.oracle.fp8 import (

27 Fp8MoeBackend,

28 convert_to_fp8_moe_kernel_format,

29 make_fp8_moe_kernel,

30 + make_fp8_moe_kernel_for_mkm,

31 make_fp8_moe_quant_config,

32 select_fp8_moe_backend,

33 )

34 from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (

34 - FLASHINFER_NVFP4_MOE_BACKENDS,

35 NvFp4MoeBackend,

36 convert_to_nvfp4_moe_kernel_format,

37 is_global_sf_supported_for_nvfp4_backend,

38 make_nvfp4_moe_kernel,

39 + make_nvfp4_moe_kernel_for_mkm,

40 make_nvfp4_moe_quant_config,

41 select_nvfp4_moe_backend,

42 )

@@ -57,12 +58,10 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (

58 build_flashinfer_fp4_cutlass_moe_prepare_finalize,

59 flashinfer_trtllm_fp4_moe,

60 flashinfer_trtllm_fp4_routed_moe,

60 - select_nvfp4_gemm_impl,

61 )

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

63 apply_fi_trtllm_fp8_per_tensor_moe,

64 build_flashinfer_fp8_cutlass_moe_prepare_finalize,

65 - select_cutlass_fp8_gemm_impl,

65 )

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

67 W8A8BlockFp8LinearOp,

@@ -84,6 +83,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

83 kFp8DynamicTokenSym,

84 kFp8StaticTensorSym,

85 kFp8StaticTokenSym,

86 + kNvfp4Dynamic,

87 + kNvfp4Static,

88 swizzle_blockscale,

89 )

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

@@ -728,14 +729,23 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

729 super().__init__(moe_config)

730 self.quant_config = quant_config

731 assert self.quant_config.is_checkpoint_fp8_serialized

731 - self.fp8_backend = select_fp8_moe_backend(

732 - block_quant=False,

733 - tp_size=moe_config.moe_parallel_config.tp_size,

734 - with_lora_support=self.moe.is_lora_enabled,

735 - is_act_and_mul=self.moe.is_act_and_mul,

732 +

733 + # Select Fp8 MoE backend

734 + self.fp8_backend, self.experts_cls = select_fp8_moe_backend(

735 + config=self.moe,

736 + weight_key=kFp8StaticTensorSym,

737 + activation_key=kFp8StaticTensorSym,

738 )

739 +

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

741 self.kernel: mk.FusedMoEModularKernel | None = None

742

743 + @property

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

745 + if self.kernel is not None:

746 + return self.kernel.prepare_finalize.topk_indices_dtype()

747 + return None

748 +

749 def maybe_make_prepare_finalize(

750 self,

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

@@ -744,8 +754,8 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

754 if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:

755 return None

756 elif self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:

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

748 - if self.moe.dp_size == 1:

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

758 + if not self.moe.moe_parallel_config.use_all2all_kernels:

759 return None

760

761 prepare_finalize = build_flashinfer_fp8_cutlass_moe_prepare_finalize(

@@ -762,12 +772,13 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

772 layer: torch.nn.Module,

773 ) -> mk.FusedMoEPermuteExpertsUnpermute:

774 assert self.moe_quant_config is not None

765 - experts = select_cutlass_fp8_gemm_impl(

766 - self.moe,

767 - self.moe_quant_config,

775 + assert self.experts_cls is not None

776 + return make_fp8_moe_kernel_for_mkm(

777 + moe_config=self.moe,

778 + quant_config=self.moe_quant_config,

779 + experts_cls=self.experts_cls,

780 + prepare_finalize=prepare_finalize,

781 )

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

770 - return experts

782

783 def create_weights(

784 self,

@@ -876,14 +887,15 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

887 replace_parameter(layer, "w13_weight_scale", w13_scale)

888 replace_parameter(layer, "w2_weight_scale", w2_scale)

889

879 - # Setup modular kernel for TP case.

890 + # Setup modular kernel.

891 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

892 if self.moe_quant_config:

893 + assert self.experts_cls is not None

894 self.kernel, self.use_inplace = make_fp8_moe_kernel(

883 - layer=layer,

895 moe_quant_config=self.moe_quant_config,

896 moe_config=self.moe,

897 fp8_backend=self.fp8_backend,

898 + experts_cls=self.experts_cls,

899 )

900

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

@@ -1335,32 +1347,35 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1347 ) -> None:

1348 super().__init__(moe_config)

1349 self.quant_config = quant_config

1338 - self.nvfp4_backend = select_nvfp4_moe_backend()

1339 - # TODO: move this type of check into the oracle.

1340 - if not self.moe.is_act_and_mul and self.nvfp4_backend not in [

1341 - NvFp4MoeBackend.FLASHINFER_CUTLASS,

1342 - NvFp4MoeBackend.MARLIN,

1343 - ]:

1344 - raise NotImplementedError(

1345 - "Non-gated activations are only supported by FlashInfer "

1346 - f"CUTLASS and Marlin NvFP4 MoE backends, not {self.nvfp4_backend}."

1347 - )

1350 + # Select experts implementation.

1351 + self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(

1352 + config=self.moe,

1353 + weight_key=kNvfp4Static,

1354 + activation_key=kNvfp4Dynamic,

1355 + )

1356 +

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

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

1359

1360 self.use_global_sf = is_global_sf_supported_for_nvfp4_backend(

1361 self.nvfp4_backend

1362 )

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

1363 +

1364 + @property

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

1366 + if self.kernel is not None:

1367 + return self.kernel.prepare_finalize.topk_indices_dtype()

1368 + return None

1369

1370 def maybe_make_prepare_finalize(

1371 self,

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

1373 ) -> mk.FusedMoEPrepareAndFinalize | None:

1358 - UNSUPPORTED = [NvFp4MoeBackend.MARLIN, NvFp4MoeBackend.FLASHINFER_TRTLLM]

1359 - if self.nvfp4_backend in UNSUPPORTED:

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

1375 return None

1376 elif self.nvfp4_backend == NvFp4MoeBackend.FLASHINFER_CUTLASS:

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

1363 - if self.moe.dp_size == 1:

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

1378 + if not self.moe.moe_parallel_config.use_all2all_kernels:

1379 return None

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

1381 prepare_finalize = build_flashinfer_fp4_cutlass_moe_prepare_finalize(

@@ -1377,13 +1392,13 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1392 layer: torch.nn.Module,

1393 ) -> mk.FusedMoEPermuteExpertsUnpermute:

1394 assert self.moe_quant_config is not None

1380 - experts = select_nvfp4_gemm_impl(

1381 - self.moe,

1382 - self.moe_quant_config,

1383 - allow_flashinfer=self.nvfp4_backend in FLASHINFER_NVFP4_MOE_BACKENDS,

1395 + assert self.experts_cls is not None

1396 + return make_nvfp4_moe_kernel_for_mkm(

1397 + moe_config=self.moe,

1398 + quant_config=self.moe_quant_config,

1399 + experts_cls=self.experts_cls,

1400 + prepare_finalize=prepare_finalize,

1401 )

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

1386 - return experts

1402

1403 def uses_weight_scale_2_pattern(self) -> bool:

1404 """

@@ -1554,13 +1569,20 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1569 replace_parameter(layer, "w2_weight_scale_2", w2_scale_2)

1570 replace_parameter(layer, "w2_input_scale", a2_scale)

1571

1572 + # Setup modular kernel for TP case and naive DP/EP case.

1573 + # In non-naive DP/EP case, we will create a ModularKernelMethod.

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

1575 + # in both cases.

1576 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

1558 - use_dp = self.moe.dp_size > 1

1559 - if self.moe_quant_config is not None and not use_dp:

1577 + if self.moe_quant_config and (

1578 + (not self.moe.moe_parallel_config.use_all2all_kernels)

1579 + or self.moe.moe_parallel_config.use_naive_all2all_kernels

1580 + ):

1581 + assert self.experts_cls is not None

1582 self.kernel = make_nvfp4_moe_kernel(

1561 - backend=self.nvfp4_backend,

1562 - quant_config=self.moe_quant_config,

1583 + moe_quant_config=self.moe_quant_config,

1584 moe_config=self.moe,

1585 + experts_cls=self.experts_cls,

1586 )

1587

1588 @property

vllm/model_executor/layers/quantization/mxfp4.py

+4 -3

@@ -853,6 +853,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

853 max_num_tokens=max_num_tokens_per_rank,

854 num_dispatchers=prepare_finalize.num_dispatchers(),

855 quant_config=self.moe_quant_config,

856 + moe_config=self.moe,

857 )

858 else:

859 raise NotImplementedError(

@@ -875,11 +876,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

876 }

877 return TrtLlmGenExperts(self.moe, self.moe_quant_config, **kwargs)

878 elif self.mxfp4_backend == Mxfp4Backend.MARLIN:

878 - return MarlinExperts(self.moe_quant_config)

879 + return MarlinExperts(self.moe, self.moe_quant_config)

880 elif self.mxfp4_backend == Mxfp4Backend.TRITON:

881 if self.moe.is_lora_enabled:

881 - return UnfusedOAITritonExperts(self.moe_quant_config)

882 - return OAITritonExperts(self.moe_quant_config)

882 + return UnfusedOAITritonExperts(self.moe, self.moe_quant_config)

883 + return OAITritonExperts(self.moe, self.moe_quant_config)

884 else:

885 raise NotImplementedError(

886 f"Incompatible Mxfp4 backend ({self.mxfp4_backend}) for EP"

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

+84 -38

@@ -11,19 +11,16 @@ 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,

14 + FusedMoEParallelConfig,

15 RoutingMethodType,

16 )

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

18 - FlashInferCuteDSLExperts,

19 -)

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

21 - FlashInferExperts,

22 -)

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

18 create_flashinfer_prepare_finalize,

19 )

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

21 + QuantKey,

22 + kNvfp4Dynamic,

23 + kNvfp4Static,

24 swizzle_blockscale,

25 )

26 from vllm.platforms import current_platform

@@ -47,6 +44,86 @@ __all__ = [

44 "build_flashinfer_fp4_cutlass_moe_prepare_finalize",

45 ]

47 +#

48 +# Methods used by the oracle for kernel selection.

49 +#

50 +

51 +

52 +def _supports_current_device() -> bool:

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

54 + p = current_platform

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

56 +

57 +

58 +def _supports_no_act_and_mul() -> bool:

59 + """Does not support non-gated MoE (i.e. Nemotron-Nano)."""

60 + return False

61 +

62 +

63 +def _supports_quant_scheme(

64 + weight_key: QuantKey | None,

65 + activation_key: QuantKey | None,

66 +) -> bool:

67 + """Supports Nvfp4 quantization."""

68 + SUPPORTED_W_A = [

69 + (kNvfp4Static, kNvfp4Dynamic),

70 + ]

71 + return (weight_key, activation_key) in SUPPORTED_W_A

72 +

73 +

74 +def _supports_activation(activation: str) -> bool:

75 + """Supports silu activation only."""

76 + return activation in ["silu"]

77 +

78 +

79 +def _supports_routing_method(

80 + routing_method: RoutingMethodType,

81 +) -> bool:

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

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

84 + return routing_method in [

85 + RoutingMethodType.DeepSeekV3,

86 + RoutingMethodType.Renormalize,

87 + RoutingMethodType.RenormalizeNaive,

88 + RoutingMethodType.Llama4,

89 + ]

90 +

91 +

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

93 + """Supports EP."""

94 + return True

95 +

96 +

97 +def is_supported_config_trtllm(

98 + moe_config: FusedMoEConfig,

99 + weight_key: QuantKey | None,

100 + activation_key: QuantKey | None,

101 + activation_format: mk.FusedMoEActivationFormat,

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

103 + """

104 + This method mirrors mk.FusedMoEPermuteExpertsUnpermute.is_supported_config

105 + """

106 +

107 + def _make_reason(reason: str) -> str:

108 + return f"kernel does not support {reason}"

109 +

110 + if not _supports_current_device():

111 + return False, _make_reason("current device")

112 + elif not (moe_config.is_act_and_mul or _supports_no_act_and_mul()):

113 + return False, _make_reason("no act_and_mul MLP layer")

114 + elif not _supports_activation(moe_config.activation):

115 + return False, _make_reason(f"{moe_config.activation} activation")

116 + elif not _supports_quant_scheme(weight_key, activation_key):

117 + return False, _make_reason("quantization scheme")

118 + elif not _supports_parallel_config(moe_config.moe_parallel_config):

119 + return False, _make_reason("parallel config")

120 + elif not _supports_routing_method(moe_config.routing_method):

121 + return False, _make_reason("routing method")

122 + elif activation_format != mk.FusedMoEActivationFormat.Standard:

123 + return False, _make_reason("activation format")

124 +

125 + return True, None

126 +

127

128 def is_flashinfer_fp4_cutlass_moe_available() -> bool:

129 """Return `True` when FlashInfer CUTLASS NV-FP4 kernels can be used."""

@@ -96,37 +173,6 @@ def build_flashinfer_fp4_cutlass_moe_prepare_finalize(

173 )

174

175

99 -def select_nvfp4_gemm_impl(

100 - moe: FusedMoEConfig,

101 - moe_quant_config: FusedMoEQuantConfig,

102 - allow_flashinfer: bool,

103 -) -> mk.FusedMoEPermuteExpertsUnpermute:

104 - """Return a GEMM *experts* implementation for NV-FP4 fused-MoE layers"""

105 -

106 - if allow_flashinfer:

107 - if envs.VLLM_FLASHINFER_MOE_BACKEND == "masked_gemm":

108 - return FlashInferCuteDSLExperts(

109 - out_dtype=moe.in_dtype,

110 - quant_config=moe_quant_config,

111 - )

112 - elif envs.VLLM_FLASHINFER_MOE_BACKEND == "throughput":

113 - return FlashInferExperts(

114 - out_dtype=moe.in_dtype,

115 - quant_config=moe_quant_config,

116 - ep_rank=moe.moe_parallel_config.ep_rank,

117 - ep_size=moe.moe_parallel_config.ep_size,

118 - tp_rank=moe.moe_parallel_config.tp_rank,

119 - tp_size=moe.moe_parallel_config.tp_size,

120 - use_dp=moe.moe_parallel_config.dp_size > 1,

121 - )

122 -

123 - # native cutlass experts currently don't support DP; TP case won't call this

124 - raise ValueError(

125 - "CutlassExpertsFp4 doesn't support DP. Use flashinfer CUTLASS "

126 - "Fused MoE backend instead (set VLLM_USE_FLASHINFER_MOE_FP4=1)"

127 - )

128 -

129 -

176 def prepare_static_weights_for_trtllm_fp4_moe(

177 # args_dequant,

178 # args,

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

-31

@@ -9,10 +9,6 @@ from vllm import envs

9 from vllm.logger import init_logger

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

11 FusedMoEConfig,

12 - FusedMoEQuantConfig,

13 -)

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

15 - FlashInferExperts,

12 )

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

14 create_flashinfer_prepare_finalize,

@@ -203,33 +199,6 @@ def build_flashinfer_fp8_cutlass_moe_prepare_finalize(

199 )

200

201

206 -def select_cutlass_fp8_gemm_impl(

207 - moe: FusedMoEConfig | None,

208 - quant_config: FusedMoEQuantConfig,

209 - out_dtype: torch.dtype | None = None,

210 - use_deepseek_fp8_block_scale: bool = False,

211 -) -> mk.FusedMoEPermuteExpertsUnpermute:

212 - """Return a GEMM *experts* implementation for fused-MoE layers"""

213 -

214 - if moe is not None:

215 - return FlashInferExperts(

216 - out_dtype=moe.in_dtype,

217 - quant_config=quant_config,

218 - ep_rank=moe.moe_parallel_config.ep_rank,

219 - ep_size=moe.moe_parallel_config.ep_size,

220 - tp_rank=moe.moe_parallel_config.tp_rank,

221 - tp_size=moe.moe_parallel_config.tp_size,

222 - use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale,

223 - )

224 -

225 - assert out_dtype is not None, "If moe config is None, out_dtype must be passed"

226 - return FlashInferExperts(

227 - out_dtype=out_dtype,

228 - quant_config=quant_config,

229 - use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale,

230 - )

231 -

232 -

202 def get_flashinfer_moe_backend() -> FlashinferMoeBackend:

203 backend_map = {

204 "throughput": FlashinferMoeBackend.CUTLASS,

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

+26 -12

@@ -48,6 +48,7 @@ class GroupShape(_GroupShape):

48 # Aliases for common quantization group shapes

49 PER_TENSOR: ClassVar["GroupShape"]

50 PER_TOKEN: ClassVar["GroupShape"]

51 + PER_CHANNEL: ClassVar["GroupShape"]

53 def is_per_tensor(self) -> bool:

54 return self.row == -1 and self.col == -1

@@ -55,12 +56,16 @@ class GroupShape(_GroupShape):

56 def is_per_token(self) -> bool:

57 return self.row == 1 and self.col == -1

59 + def is_per_channel(self) -> bool:

60 + return self.row == -1 and self.col == 1

61 +

62 def is_per_group(self) -> bool:

63 return self.row == 1 and self.col >= 1

66 GroupShape.PER_TENSOR = GroupShape(-1, -1)

67 GroupShape.PER_TOKEN = GroupShape(1, -1)

68 +GroupShape.PER_CHANNEL = GroupShape(-1, 1)

71 @dataclass(frozen=True)

@@ -77,16 +82,12 @@ class ScaleDesc:

82 group_shape: GroupShape

84 def __str__(self):

80 - group_shape = (

81 - "per_tensor"

82 - if self.group_shape == GroupShape.PER_TENSOR

83 - else (

84 - "per_token"

85 - if self.group_shape == GroupShape.PER_TOKEN

86 - else str(self.group_shape)

87 - )

88 - )

89 -

85 + d = {

86 + GroupShape.PER_TENSOR: "per_tensor",

87 + GroupShape.PER_TOKEN: "per_token",

88 + GroupShape.PER_CHANNEL: "per_channel",

89 + }

90 + group_shape = d.get(self.group_shape, str(self.group_shape))

91 return (

92 f"{fx.graph.dtype_abbrs[self.dtype]},"

93 f"{'static' if self.static else 'dynamic'},{group_shape}"

@@ -126,15 +127,28 @@ kFp8DynamicTensorSym = QuantKey(FP8_DTYPE, kDynamicTensorScale, symmetric=True)

127 kStaticTokenScale = ScaleDesc(torch.float32, True, GroupShape.PER_TOKEN)

128 kFp8StaticTokenSym = QuantKey(FP8_DTYPE, kStaticTokenScale, symmetric=True)

129

130 +kStaticChannelScale = ScaleDesc(torch.float32, True, GroupShape.PER_CHANNEL)

131 +kFp8StaticChannelSym = QuantKey(FP8_DTYPE, kStaticChannelScale, symmetric=True)

132 +

133 kDynamicTokenScale = ScaleDesc(torch.float32, False, GroupShape.PER_TOKEN)

134 kFp8DynamicTokenSym = QuantKey(FP8_DTYPE, kDynamicTokenScale, symmetric=True)

135

132 -kNvfp4GroupScale = ScaleDesc(FP8_DTYPE, False, GroupShape(1, 16))

133 -kNvfp4Quant = QuantKey(FP4_DTYPE, scale=kNvfp4GroupScale, scale2=kStaticTensorScale)

136 +kNvfp4DynamicGroupScale = ScaleDesc(FP8_DTYPE, False, GroupShape(1, 16))

137 +kNvfp4Dynamic = QuantKey(

138 + FP4_DTYPE, scale=kNvfp4DynamicGroupScale, scale2=kStaticTensorScale

139 +)

140 +

141 +kNvfp4StaticGroupScale = ScaleDesc(FP8_DTYPE, True, GroupShape(1, 16))

142 +kNvfp4Static = QuantKey(

143 + FP4_DTYPE, scale=kNvfp4StaticGroupScale, scale2=kStaticTensorScale

144 +)

145

146 kDynamic128Scale = ScaleDesc(torch.float32, False, GroupShape(1, 128))

147 kFp8Dynamic128Sym = QuantKey(FP8_DTYPE, kDynamic128Scale, symmetric=True)

148

149 +kStatic128BlockScale = ScaleDesc(torch.float32, True, GroupShape(128, 128))

150 +kFp8Static128BlockSym = QuantKey(FP8_DTYPE, kStatic128BlockScale, symmetric=True)

151 +

152 kDynamic64Scale = ScaleDesc(torch.float32, False, GroupShape(1, 64))

153 kFp8Dynamic64Sym = QuantKey(FP8_DTYPE, kDynamic64Scale, symmetric=True)

154

vllm/model_executor/models/qwen3_moe.py

-2

@@ -43,7 +43,6 @@ from vllm.distributed import (

43 from vllm.logger import init_logger

44 from vllm.model_executor.layers.activation import SiluAndMul

45 from vllm.model_executor.layers.fused_moe import FusedMoE

46 -from vllm.model_executor.layers.fused_moe.config import RoutingMethodType

46 from vllm.model_executor.layers.layernorm import RMSNorm

47 from vllm.model_executor.layers.linear import (

48 MergedColumnParallelLinear,

@@ -172,7 +171,6 @@ class Qwen3MoeSparseMoeBlock(nn.Module):

171 enable_eplb=self.enable_eplb,

172 num_redundant_experts=self.n_redundant_experts,

173 is_sequence_parallel=self.is_sequence_parallel,

175 - routing_method_type=RoutingMethodType.Renormalize,

174 )

175

176 self.gate = ReplicatedLinear(

vllm/model_executor/models/qwen3_next.py

-2

@@ -34,7 +34,6 @@ from vllm.model_executor.layers.fla.ops import (

34 fused_recurrent_gated_delta_rule,

35 )

36 from vllm.model_executor.layers.fused_moe import SharedFusedMoE

37 -from vllm.model_executor.layers.fused_moe.config import RoutingMethodType

37 from vllm.model_executor.layers.layernorm import (

38 GemmaRMSNorm as Qwen3NextRMSNorm,

39 )

@@ -181,7 +180,6 @@ class Qwen3NextSparseMoeBlock(nn.Module):

180 enable_eplb=self.enable_eplb,

181 num_redundant_experts=self.n_redundant_experts,

182 is_sequence_parallel=self.is_sequence_parallel,

184 - routing_method_type=RoutingMethodType.Renormalize,

183 )

184

185 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

vllm/model_executor/warmup/deep_gemm_warmup.py

@@ -128,11 +128,15 @@ def _fp8_linear_may_use_deep_gemm(module: torch.nn.Module) -> bool:

128 """

129 Return True if the input module/layer could be processed with DeepGEMM.

130 """

131 +

132 + # FIXME: this logic is brittle and incorrect - since we

133 + # could use DeepGEMM with for than just Fp8LinearMethod

134 block_size = get_mk_alignment_for_contiguous_layout()[0]

135 if not (

136 isinstance(module, LinearBase)

137 and isinstance(module.quant_method, Fp8LinearMethod)

138 and module.quant_method.block_quant

139 + and not module.quant_method.use_marlin

140 ):

141 return False

142

vllm/v1/attention/backends/flashinfer.py

+2 -2

@@ -29,7 +29,7 @@ from vllm.model_executor.layers.batch_invariant import (

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

30 QuantKey,

31 kFp8StaticTensorSym,

32 - kNvfp4Quant,

32 + kNvfp4Dynamic,

33 )

34 from vllm.platforms import current_platform

35 from vllm.platforms.interface import DeviceCapability

@@ -1184,7 +1184,7 @@ class FlashInferImpl(AttentionImpl):

1184 return (

1185 self.support_trtllm_attn

1186 and self.kv_cache_dtype.startswith("fp8")

1187 - and quant_key in (kFp8StaticTensorSym, kNvfp4Quant)

1187 + and quant_key in (kFp8StaticTensorSym, kNvfp4Dynamic)

1188 )

1189

1190 # FlashInfer requires attention sinks to be float32