@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 089e412878

[Perf] Integrate TRTLLM BF16 MoE Modular Kernel (#45182)

Signed-off-by: Kaihang Jiang <kaihangj@login-lyris02.lyris.clusters.nvidia.com>

Kaihang Jiang committed Jul 8, 2026 at 17:36 UTC 089e41287824763fc2024c715f2e2c5b4276f8e3
5 files changed +485 -55
tests/kernels/moe/test_trtllm_bf16_moe.py new
+145
@@ -0,0 +1,145 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Tests for the FlashInfer TRTLLM BF16 MoE backend
5 +(`TrtLlmBf16ExpertsModular`).
6 +
7 +This mirrors the TRTLLM NvFP4 modular test shape: construct the modular
8 +expert wrapper directly, pass production-format BlockMajorK weights, and
9 +compare against a torch MoE reference using the original BF16 weights.
10 +"""
11 +
12 +import pytest
13 +import torch
14 +
15 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
16 +from tests.kernels.utils import torch_moe
17 +from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
18 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation
19 +from vllm.model_executor.layers.fused_moe.all2all_utils import (
20 + maybe_make_prepare_finalize,
21 +)
22 +from vllm.model_executor.layers.fused_moe.config import (
23 + FUSED_MOE_UNQUANTIZED_CONFIG,
24 + FusedMoEConfig,
25 + FusedMoEParallelConfig,
26 + RoutingMethodType,
27 +)
28 +from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import (
29 + TrtLlmBf16ExpertsModular,
30 +)
31 +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
32 + convert_moe_weights_to_flashinfer_trtllm_block_layout,
33 +)
34 +from vllm.platforms import current_platform
35 +from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
36 +from vllm.utils.math_utils import next_power_of_2
37 +from vllm.utils.torch_utils import set_random_seed
38 +
39 +if pytest and (
40 + not has_flashinfer_trtllm_fused_moe()
41 + or not current_platform.has_device_capability(100)
42 +):
43 + pytest.skip(
44 + "Requires flashinfer TRTLLM fused MoE BF16 backend (SM100)",
45 + allow_module_level=True,
46 + )
47 +
48 +# (m, n, k) = (tokens, intermediate_size_per_partition, hidden_dim).
49 +# Covers larger NvFP4-like shapes while keeping BF16's FlashInfer TRTLLM
50 +# intermediate-size multiple-of-128 requirement.
51 +MNK_FACTORS = [
52 + (2, 1024, 1024),
53 + (64, 2048, 1536),
54 + (64, 1024, 4096),
55 +]
56 +
57 +
58 +@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
59 +@pytest.mark.parametrize("e", [128])
60 +@pytest.mark.parametrize("topk", [8])
61 +@pytest.mark.parametrize("dtype", [torch.bfloat16])
62 +@torch.inference_mode()
63 +def test_trtllm_bf16_moe_modular_no_graph(
64 + m: int,
65 + n: int,
66 + k: int,
67 + e: int,
68 + topk: int,
69 + dtype: torch.dtype,
70 + workspace_init,
71 +):
72 + set_random_seed(7)
73 + with set_current_vllm_config(
74 + VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
75 + ):
76 + a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
77 + w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
78 + w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
79 + score = torch.randn((m, e), device="cuda", dtype=dtype)
80 + scores = torch.softmax(score, dim=-1, dtype=torch.float32)
81 + topk_weights, topk_ids = torch.topk(scores, topk)
82 + topk_weights = topk_weights.contiguous()
83 + topk_ids = topk_ids.to(torch.int32).contiguous()
84 +
85 + moe_config = FusedMoEConfig(
86 + num_experts=e,
87 + experts_per_token=topk,
88 + hidden_dim=k,
89 + intermediate_size=n,
90 + num_local_experts=e,
91 + num_logical_experts=e,
92 + activation=MoEActivation.SILU,
93 + device="cuda",
94 + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
95 + in_dtype=dtype,
96 + routing_method=RoutingMethodType.TopK,
97 + max_num_tokens=next_power_of_2(m),
98 + )
99 +
100 + trtllm_w1, trtllm_w2 = convert_moe_weights_to_flashinfer_trtllm_block_layout(
101 + {},
102 + w1,
103 + w2,
104 + )
105 +
106 + trtllm_experts = mk.FusedMoEKernel(
107 + maybe_make_prepare_finalize(
108 + moe=moe_config,
109 + quant_config=FUSED_MOE_UNQUANTIZED_CONFIG,
110 + allow_new_interface=True,
111 + use_monolithic=False,
112 + ),
113 + TrtLlmBf16ExpertsModular(
114 + moe_config=moe_config,
115 + quant_config=FUSED_MOE_UNQUANTIZED_CONFIG,
116 + ),
117 + )
118 +
119 + trtllm_output = trtllm_experts.apply(
120 + hidden_states=a,
121 + w1=trtllm_w1,
122 + w2=trtllm_w2,
123 + topk_weights=topk_weights,
124 + topk_ids=topk_ids,
125 + activation=MoEActivation.SILU,
126 + global_num_experts=e,
127 + expert_map=None,
128 + apply_router_weight_on_input=False,
129 + )
130 +
131 + torch_output = torch_moe(
132 + a,
133 + w1,
134 + w2,
135 + score,
136 + topk,
137 + activation=MoEActivation.SILU,
138 + )
139 +
140 + torch.testing.assert_close(
141 + torch_output,
142 + trtllm_output,
143 + atol=1e-1,
144 + rtol=2e-1,
145 + )
tests/kernels/moe/test_unquantized_backend_selection.py
+169 -5
@@ -5,6 +5,7 @@ from unittest.mock import patch
5 import pytest
6
7 from tests.kernels.moe.utils import make_dummy_moe_config
8 +from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
9 from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
10 UnquantizedMoeBackend,
11 select_unquantized_moe_backend,
@@ -117,13 +118,15 @@ def test_select_rocm_aiter_backend(mock_aiter_enabled, mock_has_flashinfer):
118
119
120 @patch(
120 - "vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16Experts.is_supported_config",
121 + "vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsMonolithic.is_supported_config",
122 return_value=(True, None),
123 )
124 @pytest.mark.skipif(
125 not current_platform.is_cuda(), reason="Only supported on NVIDIA platforms."
126 )
126 -def test_select_cuda_flashinfer_trtllm_backend(mock_is_supported_trtllm):
127 +def test_select_cuda_flashinfer_trtllm_backend(
128 + mock_is_supported_trtllm_monolithic,
129 +):
130 """Test CUDA backend selection when FlashInfer TRTLLM is available and enabled."""
131 with (
132 patch.object(current_platform, "is_cuda", return_value=True),
@@ -146,6 +149,162 @@ def test_select_cuda_flashinfer_trtllm_backend(mock_is_supported_trtllm):
149
150 assert selected_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM
151 assert experts_cls is not None
152 + assert experts_cls.__name__ == "TrtLlmBf16ExpertsMonolithic"
153 +
154 +
155 +@patch(
156 + "vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsMonolithic.is_supported_config",
157 + return_value=(False, "monolithic unsupported"),
158 +)
159 +@patch(
160 + "vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsModular.is_supported_config",
161 + return_value=(True, None),
162 +)
163 +@pytest.mark.skipif(
164 + not current_platform.is_cuda(), reason="Only supported on NVIDIA platforms."
165 +)
166 +def test_select_cuda_flashinfer_trtllm_modular_backend(
167 + mock_is_supported_trtllm_modular,
168 + mock_is_supported_trtllm_monolithic,
169 +):
170 + """Test CUDA backend selection falls back to FlashInfer TRTLLM modular."""
171 + with (
172 + patch.object(current_platform, "is_cuda", return_value=True),
173 + patch.object(current_platform, "is_rocm", return_value=False),
174 + patch.object(current_platform, "is_cpu", return_value=False),
175 + patch.object(current_platform, "is_xpu", return_value=False),
176 + patch.object(current_platform, "is_tpu", return_value=False),
177 + patch.object(current_platform, "is_out_of_tree", return_value=False),
178 + patch.object(current_platform, "has_device_capability", return_value=True),
179 + ):
180 + moe_config = make_dummy_moe_config()
181 + moe_config.moe_backend = "flashinfer_trtllm"
182 + moe_config.moe_parallel_config.use_ep = True
183 + moe_config.moe_parallel_config.use_dp = False
184 +
185 + selected_backend, experts_cls = select_unquantized_moe_backend(
186 + moe_config=moe_config
187 + )
188 +
189 + assert selected_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM
190 + assert experts_cls is not None
191 + assert experts_cls.__name__ == "TrtLlmBf16ExpertsModular"
192 +
193 +
194 +@patch(
195 + "vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsBase._supports_current_device",
196 + return_value=True,
197 +)
198 +@pytest.mark.parametrize(
199 + "all2all_backend",
200 + [
201 + "mori_high_throughput",
202 + "mori_low_latency",
203 + "flashinfer_nvlink_two_sided",
204 + "flashinfer_nvlink_one_sided",
205 + ],
206 +)
207 +def test_select_cuda_flashinfer_trtllm_modular_for_standard_all2all(
208 + mock_supports_current_device,
209 + all2all_backend,
210 +):
211 + """Test non-AG/RS standard-format all2all backends select modular BF16."""
212 + with (
213 + patch.object(current_platform, "is_cuda", return_value=True),
214 + patch.object(current_platform, "is_rocm", return_value=False),
215 + patch.object(current_platform, "is_cpu", return_value=False),
216 + patch.object(current_platform, "is_xpu", return_value=False),
217 + patch.object(current_platform, "is_tpu", return_value=False),
218 + patch.object(current_platform, "is_out_of_tree", return_value=False),
219 + patch.object(
220 + current_platform, "is_device_capability_family", return_value=False
221 + ),
222 + ):
223 + moe_config = make_dummy_moe_config(num_experts=4, num_local_experts=2)
224 + moe_config.moe_backend = "flashinfer_trtllm"
225 + moe_config.moe_parallel_config.use_ep = True
226 + moe_config.moe_parallel_config.dp_size = 2
227 + moe_config.moe_parallel_config.ep_size = 2
228 + moe_config.moe_parallel_config.all2all_backend = all2all_backend
229 +
230 + selected_backend, experts_cls = select_unquantized_moe_backend(
231 + moe_config=moe_config
232 + )
233 +
234 + assert selected_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM
235 + assert experts_cls is not None
236 + assert experts_cls.__name__ == "TrtLlmBf16ExpertsModular"
237 +
238 +
239 +@patch(
240 + "vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsBase._supports_current_device",
241 + return_value=True,
242 +)
243 +def test_select_cuda_deepep_ht_falls_back_from_trtllm(
244 + mock_supports_current_device,
245 +):
246 + """Test DeepEP HT avoids the unsupported BF16 TRTLLM modular path."""
247 + with (
248 + patch.object(current_platform, "is_cuda", return_value=True),
249 + patch.object(current_platform, "is_rocm", return_value=False),
250 + patch.object(current_platform, "is_cpu", return_value=False),
251 + patch.object(current_platform, "is_xpu", return_value=False),
252 + patch.object(current_platform, "is_tpu", return_value=False),
253 + patch.object(current_platform, "is_out_of_tree", return_value=False),
254 + patch.object(
255 + current_platform, "is_device_capability_family", return_value=False
256 + ),
257 + ):
258 + moe_config = make_dummy_moe_config(num_experts=4, num_local_experts=2)
259 + moe_config.moe_backend = "auto"
260 + moe_config.moe_parallel_config.use_ep = True
261 + moe_config.moe_parallel_config.dp_size = 2
262 + moe_config.moe_parallel_config.ep_size = 2
263 + moe_config.moe_parallel_config.all2all_backend = "deepep_high_throughput"
264 +
265 + selected_backend, experts_cls = select_unquantized_moe_backend(
266 + moe_config=moe_config
267 + )
268 +
269 + assert selected_backend == UnquantizedMoeBackend.TRITON
270 + assert experts_cls is not None
271 + assert experts_cls.__name__ == "TritonExperts"
272 +
273 +
274 +@patch(
275 + "vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsBase._supports_current_device",
276 + return_value=True,
277 +)
278 +def test_select_cuda_flashinfer_trtllm_ag_rs_uses_monolithic(
279 + mock_supports_current_device,
280 +):
281 + """Test AG/RS stays on BF16 TRTLLM monolithic when TRTLLM is supported."""
282 + with (
283 + patch.object(current_platform, "is_cuda", return_value=True),
284 + patch.object(current_platform, "is_rocm", return_value=False),
285 + patch.object(current_platform, "is_cpu", return_value=False),
286 + patch.object(current_platform, "is_xpu", return_value=False),
287 + patch.object(current_platform, "is_tpu", return_value=False),
288 + patch.object(current_platform, "is_out_of_tree", return_value=False),
289 + patch.object(
290 + current_platform, "is_device_capability_family", return_value=False
291 + ),
292 + ):
293 + moe_config = make_dummy_moe_config(num_experts=4, num_local_experts=2)
294 + moe_config.moe_backend = "flashinfer_trtllm"
295 + moe_config.routing_method = RoutingMethodType.Renormalize
296 + moe_config.moe_parallel_config.use_ep = True
297 + moe_config.moe_parallel_config.dp_size = 2
298 + moe_config.moe_parallel_config.ep_size = 2
299 + moe_config.moe_parallel_config.all2all_backend = "allgather_reducescatter"
300 +
301 + selected_backend, experts_cls = select_unquantized_moe_backend(
302 + moe_config=moe_config
303 + )
304 +
305 + assert selected_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM
306 + assert experts_cls is not None
307 + assert experts_cls.__name__ == "TrtLlmBf16ExpertsMonolithic"
308
309
310 @patch(
@@ -153,7 +312,11 @@ def test_select_cuda_flashinfer_trtllm_backend(mock_is_supported_trtllm):
312 return_value=True,
313 )
314 @patch(
156 - "vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16Experts.is_supported_config",
315 + "vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsMonolithic.is_supported_config",
316 + return_value=(False, None),
317 +)
318 +@patch(
319 + "vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16ExpertsModular.is_supported_config",
320 return_value=(False, None),
321 )
322 @patch(
@@ -164,9 +327,10 @@ def test_select_cuda_flashinfer_trtllm_backend(mock_is_supported_trtllm):
327 not current_platform.is_cuda(), reason="Only supported on NVIDIA platforms."
328 )
329 def test_select_cuda_flashinfer_cutlass_backend(
167 - mock_has_flashinfer,
168 - mock_is_supported_trtllm,
330 mock_is_supported_cutlass,
331 + mock_is_supported_trtllm_modular,
332 + mock_is_supported_trtllm_monolithic,
333 + mock_has_flashinfer,
334 ):
335 """Test CUDA backend selection when FlashInfer TRTLLM is not available
336 and FlashInfer CUTLASS is available."""
vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py
+139 -23
@@ -11,7 +11,13 @@ from vllm.model_executor.layers.fused_moe.config import (
11 FusedMoEQuantConfig,
12 RoutingMethodType,
13 )
14 -from vllm.model_executor.layers.fused_moe.utils import fi_moe_largest_bucket
14 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
15 + TopKWeightAndReduceNoOP,
16 +)
17 +from vllm.model_executor.layers.fused_moe.utils import (
18 + fi_moe_largest_bucket,
19 + trtllm_moe_pack_topk_ids_weights,
20 +)
21 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
22 activation_to_flashinfer_int,
23 )
@@ -22,9 +28,10 @@ from vllm.platforms import current_platform
28 from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
29
30
25 -class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
31 +class TrtLlmBf16ExpertsBase:
32 """
27 - BF16 unquantized TRTLLM-Gen MoE kernels. Supports monolithic interface.
33 + BF16 unquantized TRTLLM-Gen MoE kernels. Shared base for modular and
34 + monolithic interfaces.
35 """
36
37 def __init__(
@@ -32,7 +39,6 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
39 moe_config: FusedMoEConfig,
40 quant_config: FusedMoEQuantConfig,
41 ):
35 - super().__init__(moe_config, quant_config)
42 self.routing_method_type = moe_config.routing_method
43 self.topk = moe_config.experts_per_token
44 self.intermediate_size_per_partition = (
@@ -42,6 +48,9 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
48 self.local_num_experts = moe_config.num_local_experts
49 self.ep_rank = moe_config.moe_parallel_config.ep_rank
50
51 + self.moe_config = moe_config
52 + self.quant_config = quant_config
53 +
54 @staticmethod
55 def activation_format() -> mk.FusedMoEActivationFormat:
56 return mk.FusedMoEActivationFormat.Standard
@@ -75,40 +84,147 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
84 return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]
85
86 @staticmethod
78 - def _supports_routing_method(
87 + def _supports_router_logits_dtype(
88 + router_logits_dtype: torch.dtype | None,
89 routing_method: RoutingMethodType,
80 - weight_key: QuantKey | None,
81 - activation_key: QuantKey | None,
90 ) -> bool:
83 - return routing_method in [
84 - RoutingMethodType.DeepSeekV3,
85 - RoutingMethodType.Llama4,
86 - RoutingMethodType.Renormalize,
87 - RoutingMethodType.RenormalizeNaive,
88 - RoutingMethodType.SigmoidRenorm,
89 - RoutingMethodType.Sigmoid,
90 - ]
91 + return True
92 +
93 + @property
94 + def expects_unquantized_inputs(self) -> bool:
95 + return True
96 +
97 +
98 +class TrtLlmBf16ExpertsModular(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsModular):
99 + """
100 + BF16 unquantized TRTLLM-Gen MoE kernels. Supports modular interface.
101 + """
102
103 @staticmethod
104 def _supports_parallel_config(
105 moe_parallel_config: FusedMoEParallelConfig,
106 ) -> bool:
96 - """Monolithic kernel so only use with naive DP/EP and TP."""
107 + return (
108 + moe_parallel_config.use_all2all_kernels
109 + and not moe_parallel_config.use_ag_rs_all2all_kernels
110 + and not moe_parallel_config.use_deepep_ht_kernels
111 + and not moe_parallel_config.enable_eplb
112 + )
113 +
114 + def moe_problem_size(
115 + self,
116 + a1: torch.Tensor,
117 + w1: torch.Tensor,
118 + w2: torch.Tensor,
119 + topk_ids: torch.Tensor,
120 + ) -> tuple[int, int, int, int, int]:
121 + """Override to handle 4D BlockMajorK weights (E, K/bk, Mn, bk)."""
122 + if w1.dim() == 4:
123 + E = w1.shape[0]
124 + N = w1.shape[2]
125 + K = a1.size(-1)
126 + M = a1.size(0) if a1.dim() == 2 else a1.size(1)
127 + topk = topk_ids.size(1)
128 + return E, M, N, K, topk
129 + return super().moe_problem_size(a1, w1, w2, topk_ids)
130 +
131 + def workspace_shapes(
132 + self,
133 + M: int,
134 + N: int,
135 + K: int,
136 + topk: int,
137 + global_num_experts: int,
138 + local_num_experts: int,
139 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
140 + activation: MoEActivation,
141 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
142 + # The workspaces for this implementation are managed by FlashInfer.
143 + workspace1 = (0,)
144 + workspace2 = (0,)
145 + output = (M, K)
146 +
147 + return workspace1, workspace2, output
148 +
149 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
150 + return TopKWeightAndReduceNoOP()
151 +
152 + def apply(
153 + self,
154 + output: torch.Tensor,
155 + hidden_states: torch.Tensor,
156 + w1: torch.Tensor,
157 + w2: torch.Tensor,
158 + topk_weights: torch.Tensor,
159 + topk_ids: torch.Tensor,
160 + activation: MoEActivation,
161 + global_num_experts: int,
162 + expert_map: torch.Tensor | None,
163 + a1q_scale: torch.Tensor | None,
164 + a2_scale: torch.Tensor | None,
165 + workspace13: torch.Tensor,
166 + workspace2: torch.Tensor,
167 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
168 + apply_router_weight_on_input: bool,
169 + ):
170 + import flashinfer
171 + from flashinfer.fused_moe import WeightLayout
172 +
173 + # Pack topk ids and weights into format expected by the TRTLLM kernel.
174 + packed_topk_ids = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)
175 +
176 + result = flashinfer.fused_moe.trtllm_bf16_routed_moe(
177 + topk_ids=packed_topk_ids,
178 + hidden_states=hidden_states,
179 + gemm1_weights=w1,
180 + gemm2_weights=w2,
181 + num_experts=global_num_experts,
182 + top_k=self.topk,
183 + n_group=None,
184 + topk_group=None,
185 + intermediate_size=self.intermediate_size_per_partition,
186 + local_expert_offset=self.ep_rank * self.local_num_experts,
187 + local_num_experts=self.local_num_experts,
188 + routed_scaling_factor=None,
189 + routing_method_type=1, # not used
190 + use_shuffled_weight=True,
191 + weight_layout=WeightLayout.BlockMajorK,
192 + do_finalize=True,
193 + activation_type=activation_to_flashinfer_int(activation),
194 + )
195 + # FlashInfer's BF16 routed wrapper does not expose an output= argument.
196 + output.copy_(result[0] if isinstance(result, list) else result)
197 +
198 +
199 +class TrtLlmBf16ExpertsMonolithic(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsMonolithic):
200 + """
201 + BF16 unquantized TRTLLM-Gen MoE kernels. Supports monolithic interface.
202 + """
203 +
204 + @staticmethod
205 + def _supports_parallel_config(
206 + moe_parallel_config: FusedMoEParallelConfig,
207 + ) -> bool:
208 + """Monolithic kernel supports no-all2all and AG/RS paths."""
209 return (
210 not moe_parallel_config.use_all2all_kernels
211 or moe_parallel_config.use_ag_rs_all2all_kernels
212 ) and not moe_parallel_config.enable_eplb
213
214 @staticmethod
103 - def _supports_router_logits_dtype(
104 - router_logits_dtype: torch.dtype | None,
215 + def _supports_routing_method(
216 routing_method: RoutingMethodType,
217 + weight_key: QuantKey | None,
218 + activation_key: QuantKey | None,
219 ) -> bool:
107 - return True
108 -
109 - @property
110 - def expects_unquantized_inputs(self) -> bool:
111 - return True
220 + return routing_method in [
221 + RoutingMethodType.DeepSeekV3,
222 + RoutingMethodType.Llama4,
223 + RoutingMethodType.Renormalize,
224 + RoutingMethodType.RenormalizeNaive,
225 + RoutingMethodType.SigmoidRenorm,
226 + RoutingMethodType.Sigmoid,
227 + ]
228
229 def apply(
230 self,
vllm/model_executor/layers/fused_moe/oracle/base.py
+4 -2
@@ -63,9 +63,11 @@ class MoEKernelOracle(ABC, Generic[BackendT]):
63 this `moe_config`."""
64
65 @abstractmethod
66 - def backend_to_kernel_cls(self, backend: BackendT) -> type[mk.FusedMoEExperts]:
66 + def backend_to_kernel_cls(
67 + self, backend: BackendT
68 + ) -> list[type[mk.FusedMoEExperts]]:
69 """Map a backend enum value to its concrete `FusedMoEExperts`
68 - subclass."""
70 + subclasses, in selection priority order."""
71
72 @abstractmethod
73 def map_backend(self, runner_backend: MoEBackend) -> BackendT:
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+28 -25
@@ -101,46 +101,47 @@ def _get_priority_backends(moe_config: FusedMoEConfig) -> list[UnquantizedMoeBac
101
102 def backend_to_kernel_cls(
103 backend: UnquantizedMoeBackend,
104 -) -> type[mk.FusedMoEExperts]:
104 +) -> list[type[mk.FusedMoEExperts]]:
105 if backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM:
106 from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import (
107 - TrtLlmBf16Experts,
107 + TrtLlmBf16ExpertsModular,
108 + TrtLlmBf16ExpertsMonolithic,
109 )
110
110 - return TrtLlmBf16Experts
111 + return [TrtLlmBf16ExpertsMonolithic, TrtLlmBf16ExpertsModular]
112
113 elif backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS:
114 from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import ( # noqa: E501
115 FlashInferExperts,
116 )
117
117 - return FlashInferExperts
118 + return [FlashInferExperts]
119
120 elif backend == UnquantizedMoeBackend.AITER:
121 from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
122 AiterExperts,
123 )
124
124 - return AiterExperts
125 + return [AiterExperts]
126
127 elif backend == UnquantizedMoeBackend.TRITON:
128 from vllm.model_executor.layers.fused_moe.experts.triton_moe import (
129 TritonExperts,
130 )
131
131 - return TritonExperts
132 + return [TritonExperts]
133
134 elif backend == UnquantizedMoeBackend.BATCHED_TRITON:
135 from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import (
136 BatchedTritonExperts,
137 )
138
138 - return BatchedTritonExperts
139 + return [BatchedTritonExperts]
140
141 elif backend == UnquantizedMoeBackend.XPU:
142 from vllm.model_executor.layers.fused_moe.experts.xpu_moe import XPUExperts
143
143 - return XPUExperts
144 + return [XPUExperts]
145
146 else:
147 raise ValueError(f"Unknown unquantized MoE backend: {backend.value}")
@@ -183,7 +184,7 @@ def select_unquantized_moe_backend(
184 if moe_config.is_lora_enabled:
185 return UnquantizedMoeBackend.TRITON, backend_to_kernel_cls(
186 UnquantizedMoeBackend.TRITON
186 - )
187 + )[0]
188
189 # NOTE: the kernels are selected in the following order.
190 AVAILABLE_BACKENDS = _get_priority_backends(moe_config)
@@ -222,13 +223,14 @@ def select_unquantized_moe_backend(
223 config: FusedMoEConfig,
224 activation_format: mk.FusedMoEActivationFormat,
225 ) -> tuple[UnquantizedMoeBackend, type[mk.FusedMoEExperts] | None]:
225 - k_cls = backend_to_kernel_cls(backend)
226 - supported, reason = k_cls.is_supported_config(
227 - k_cls, config, None, None, activation_format
228 - )
229 - if supported:
230 - logger.info_once(_make_log_backend(backend))
231 - return backend, k_cls
226 + reason = None
227 + for k_cls in backend_to_kernel_cls(backend):
228 + supported, reason = k_cls.is_supported_config(
229 + k_cls, config, None, None, activation_format
230 + )
231 + if supported:
232 + logger.info_once(_make_log_backend(backend))
233 + return backend, k_cls
234 raise ValueError(_make_log_unsupported(backend, reason))
235
236 runner_backend = moe_config.moe_backend
@@ -254,15 +256,15 @@ def select_unquantized_moe_backend(
256 return _return_or_raise(backend, moe_config, activation_format)
257
258 for backend in AVAILABLE_BACKENDS:
257 - k_cls = backend_to_kernel_cls(backend)
258 - supported, reason = k_cls.is_supported_config(
259 - k_cls, moe_config, None, None, activation_format
260 - )
261 - if supported:
262 - logger.info_once(_make_log_backend(backend))
263 - return backend, k_cls
259 + for k_cls in backend_to_kernel_cls(backend):
260 + supported, reason = k_cls.is_supported_config(
261 + k_cls, moe_config, None, None, activation_format
262 + )
263 + if supported:
264 + logger.info_once(_make_log_backend(backend))
265 + return backend, k_cls
266
265 - logger.debug_once(_make_log_unsupported(backend, reason))
267 + logger.debug_once(_make_log_unsupported(backend, reason))
268
269 raise NotImplementedError(
270 "No Unquantized MoE backend supports the deployment configuration."
@@ -332,6 +334,7 @@ def make_unquantized_moe_kernel(
334 assert prepare_finalize is not None
335
336 logger.info_once("Using %s", prepare_finalize.__class__.__name__)
337 + logger.info_once("Using %s MoE backend", experts_cls.__name__)
338
339 # Create Experts
340 if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
@@ -383,7 +386,7 @@ class UnquantizedMoEKernelOracle(MoEKernelOracle[UnquantizedMoeBackend]):
386
387 def backend_to_kernel_cls(
388 self, backend: UnquantizedMoeBackend
386 - ) -> type[mk.FusedMoEExperts]:
389 + ) -> list[type[mk.FusedMoEExperts]]:
390 return backend_to_kernel_cls(backend)
391
392 def map_backend(self, runner_backend: MoEBackend) -> UnquantizedMoeBackend: