@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d9b90a07ac
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 5, 17:15 UTC

[MoE Refactor] Migrate Unquantized to Full Oracle Flow (#36286)

Signed-off-by: Yifan Zong <yzong@redhat.com> Signed-off-by: Robert Shaw <robshaw@redhat.com> Signed-off-by: yzong-rh <yzong@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Robert Shaw <robshaw@redhat.com>

yzong-rh committed Mar 31, 2026 at 15:43 UTC d9b90a07aced1789998f97f14e5e1456d4e671f1
11 files changed +608 -504
tests/kernels/moe/test_moe.py
+15 -3
@@ -1664,7 +1664,7 @@ def test_unquantized_bf16_flashinfer_trtllm_backend(
1664 intermediate_size_per_partition=n,
1665 num_local_experts=e,
1666 num_logical_experts=e,
1667 - activation="silu",
1667 + activation=MoEActivation.SILU,
1668 device="cuda",
1669 moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
1670 in_dtype=dtype,
@@ -1695,13 +1695,25 @@ def test_unquantized_bf16_flashinfer_trtllm_backend(
1695 layer.topk_group = 1
1696 layer.intermediate_size_per_partition = n
1697 layer.ep_rank = 0
1698 - layer.activation = "silu"
1698 + layer.activation = MoEActivation.SILU
1699 layer.e_score_correction_bias = None
1700 layer.routing_method_type = RoutingMethodType.Renormalize
1701 + layer.expert_map = None
1702 + layer.apply_router_weight_on_input = False
1703 + layer.routed_scaling_factor = None
1704 + layer.shared_experts = None
1705 + layer._maybe_init_expert_routing_tables = lambda: None
1706
1707 quant_method.process_weights_after_loading(layer)
1708
1704 - trtllm_output = quant_method.forward_monolithic_cuda(
1709 + assert quant_method.moe_kernel is not None, (
1710 + "moe_kernel should be set after process_weights_after_loading"
1711 + )
1712 + assert quant_method.supports_internal_mk, (
1713 + "supports_internal_mk should be True after setup"
1714 + )
1715 +
1716 + trtllm_output = quant_method.apply_monolithic(
1717 layer=layer,
1718 x=a,
1719 router_logits=router_logits,
tests/kernels/moe/test_unquantized_backend_selection.py
+64 -46
@@ -24,7 +24,7 @@ from vllm.platforms import current_platform
24 ],
25 )
26 @patch(
27 - "vllm.model_executor.layers.fused_moe.oracle.unquantized.has_flashinfer",
27 + "vllm.utils.flashinfer.has_flashinfer",
28 return_value=False,
29 )
30 @patch(
@@ -54,13 +54,29 @@ def test_select_default_backend_by_platform(
54 # Set only the specified platform to True
55 getattr(mock_platform, platform_method).return_value = True
56
57 + with (
58 + patch.object(current_platform, "is_cuda", return_value=False),
59 + patch.object(current_platform, "is_rocm", return_value=False),
60 + patch.object(current_platform, "is_cpu", return_value=False),
61 + patch.object(current_platform, "is_xpu", return_value=False),
62 + patch.object(current_platform, "is_tpu", return_value=False),
63 + patch.object(current_platform, "is_out_of_tree", return_value=False),
64 + patch.object(current_platform, platform_method, return_value=True),
65 + ):
66 moe_config = make_dummy_moe_config()
58 - selected_backend = select_unquantized_moe_backend(
59 - moe_config=moe_config,
60 - use_dp=False,
67 + selected_backend, expert_cls = select_unquantized_moe_backend(
68 + moe_config=moe_config
69 )
70
71 assert selected_backend == expected_backend
72 + if expected_backend in [
73 + UnquantizedMoeBackend.CPU,
74 + UnquantizedMoeBackend.OOT,
75 + UnquantizedMoeBackend.TPU,
76 + ]:
77 + assert expert_cls is None
78 + else:
79 + assert expert_cls is not None
80
81
82 @patch(
@@ -87,88 +103,90 @@ def test_select_rocm_aiter_backend(mock_aiter_enabled, mock_has_flashinfer):
103 mock_platform.is_out_of_tree.return_value = False
104
105 moe_config = make_dummy_moe_config()
90 - selected_backend = select_unquantized_moe_backend(
106 + selected_backend, expert_cls = select_unquantized_moe_backend(
107 moe_config=moe_config,
92 - use_dp=False,
108 )
109
110 assert selected_backend == UnquantizedMoeBackend.AITER
111 + assert expert_cls is not None
112
113
114 @patch(
99 - "vllm.model_executor.layers.fused_moe.oracle.unquantized.has_flashinfer",
100 - return_value=True,
101 -)
102 -@patch(
103 - "vllm.model_executor.layers.fused_moe.oracle.unquantized.is_supported_config_trtllm_bf16",
115 + "vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16Experts.is_supported_config",
116 return_value=(True, None),
117 )
118 @pytest.mark.skipif(
119 not current_platform.is_cuda(), reason="Only supported on NVIDIA platforms."
120 )
109 -def test_select_cuda_flashinfer_trtllm_backend(
110 - mock_has_flashinfer, mock_is_supported_trtllm, monkeypatch
111 -):
121 +def test_select_cuda_flashinfer_trtllm_backend(mock_is_supported_trtllm, monkeypatch):
122 """Test CUDA backend selection when FlashInfer TRTLLM is available and enabled."""
113 - with patch(
114 - "vllm.model_executor.layers.fused_moe.oracle.unquantized.current_platform"
115 - ) as mock_platform:
116 - # Set as CUDA platform
117 - mock_platform.is_cuda.return_value = True
118 - mock_platform.is_rocm.return_value = False
119 - mock_platform.is_cpu.return_value = False
120 - mock_platform.is_xpu.return_value = False
121 - mock_platform.is_tpu.return_value = False
122 - mock_platform.is_out_of_tree.return_value = False
123 -
123 + with (
124 + patch.object(current_platform, "is_cuda", return_value=True),
125 + patch.object(current_platform, "is_rocm", return_value=False),
126 + patch.object(current_platform, "is_cpu", return_value=False),
127 + patch.object(current_platform, "is_xpu", return_value=False),
128 + patch.object(current_platform, "is_tpu", return_value=False),
129 + patch.object(current_platform, "is_out_of_tree", return_value=False),
130 + patch.object(current_platform, "has_device_capability", return_value=True),
131 + ):
132 monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP16", "1")
133
134 moe_config = make_dummy_moe_config()
135 + # TRTLLM requires EP and does not support DP
136 + moe_config.moe_parallel_config.use_ep = True
137 + moe_config.moe_parallel_config.use_dp = False
138
128 - selected_backend = select_unquantized_moe_backend(
129 - moe_config=moe_config,
130 - use_dp=False,
139 + selected_backend, experts_cls = select_unquantized_moe_backend(
140 + moe_config=moe_config
141 )
142
143 assert selected_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM
144 + assert experts_cls is not None
145
146
147 @patch(
137 - "vllm.model_executor.layers.fused_moe.oracle.unquantized.has_flashinfer",
148 + "vllm.utils.flashinfer.has_flashinfer",
149 return_value=True,
150 )
151 @patch(
141 - "vllm.model_executor.layers.fused_moe.oracle.unquantized.is_supported_config_trtllm_bf16",
152 + "vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe.TrtLlmBf16Experts.is_supported_config",
153 return_value=(False, None),
154 )
155 +@patch(
156 + "vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts.is_supported_config",
157 + return_value=(True, None),
158 +)
159 @pytest.mark.skipif(
160 not current_platform.is_cuda(), reason="Only supported on NVIDIA platforms."
161 )
162 def test_select_cuda_flashinfer_cutlass_backend(
148 - mock_has_flashinfer, mock_is_supported_trtllm, monkeypatch
163 + mock_has_flashinfer,
164 + mock_is_supported_trtllm,
165 + mock_is_supported_cutlass,
166 + monkeypatch,
167 ):
168 """Test CUDA backend selection when FlashInfer TRTLLM is not available
169 and FlashInfer CUTLASS is available."""
152 - with patch(
153 - "vllm.model_executor.layers.fused_moe.oracle.unquantized.current_platform"
154 - ) as mock_platform:
155 - # Set as CUDA platform with Hopper capability
156 - mock_platform.is_cuda.return_value = True
157 - mock_platform.is_rocm.return_value = False
158 - mock_platform.is_cpu.return_value = False
159 - mock_platform.is_xpu.return_value = False
160 - mock_platform.is_tpu.return_value = False
161 - mock_platform.is_out_of_tree.return_value = False
162 - mock_platform.has_device_capability.return_value = True # SM90+
163 -
170 + with (
171 + patch.object(current_platform, "is_cuda", return_value=True),
172 + patch.object(current_platform, "is_rocm", return_value=False),
173 + patch.object(current_platform, "is_cpu", return_value=False),
174 + patch.object(current_platform, "is_xpu", return_value=False),
175 + patch.object(current_platform, "is_tpu", return_value=False),
176 + patch.object(current_platform, "is_out_of_tree", return_value=False),
177 + patch.object(current_platform, "has_device_capability", return_value=True),
178 + ):
179 # Enable FlashInfer via env var
180 monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP16", "1")
181
182 moe_config = make_dummy_moe_config()
183 + # CUTLASS requires EP and does not support DP
184 + moe_config.moe_parallel_config.use_ep = True
185 + moe_config.moe_parallel_config.use_dp = False
186
169 - selected_backend = select_unquantized_moe_backend(
170 - moe_config=moe_config,
171 - use_dp=False, # CUTLASS doesn't support DP
187 + selected_backend, experts_cls = select_unquantized_moe_backend(
188 + moe_config=moe_config
189 )
190
191 assert selected_backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS
192 + assert experts_cls is not None
tests/quantization/test_blackwell_moe.py
+7
@@ -210,6 +210,13 @@ def test_gptoss_eager(monkeypatch: pytest.MonkeyPatch):
210 ## Qwen3 Next ##
211
212
213 +@pytest.mark.skip(
214 + reason=(
215 + "FLASHINFER TRTLLM MoE has a bug with all negative router logits "
216 + "for models with RENORMALIZE. This will be re-enabled once the "
217 + "issue is fixed in flashinfer."
218 + )
219 +)
220 def test_qwen3_next_bf16_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch):
221 can_initialize(
222 "Qwen/Qwen3-Next-80B-A3B-Instruct",
vllm/lora/layers/fused_moe.py
+3
@@ -49,6 +49,9 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
49 assert not self.base_layer.use_ep, (
50 "EP support for Fused MoE LoRA is not implemented yet."
51 )
52 + assert not self.base_layer.quant_method.is_monolithic, (
53 + "Monolithic kernels are not supported for Fused MoE LoRA."
54 + )
55 self.tp_size = get_tensor_model_parallel_world_size()
56 self.tp_rank = get_tensor_model_parallel_rank()
57 self.device = _get_lora_device(base_layer)
vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py new
+148
@@ -0,0 +1,148 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation
8 +from vllm.model_executor.layers.fused_moe.config import (
9 + FusedMoEConfig,
10 + FusedMoEParallelConfig,
11 + FusedMoEQuantConfig,
12 + RoutingMethodType,
13 +)
14 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
15 + QuantKey,
16 +)
17 +from vllm.platforms import current_platform
18 +from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
19 +
20 +
21 +class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
22 + """
23 + BF16 unquantized TRTLLM-Gen MoE kernels. Supports monolithic interface.
24 + """
25 +
26 + def __init__(
27 + self,
28 + moe_config: FusedMoEConfig,
29 + quant_config: FusedMoEQuantConfig,
30 + ):
31 + super().__init__(moe_config, quant_config)
32 + self.routing_method_type = moe_config.routing_method
33 + self.topk = moe_config.experts_per_token
34 + self.intermediate_size_per_partition = (
35 + moe_config.intermediate_size_per_partition
36 + )
37 + self.hidden_dim = moe_config.hidden_dim
38 + self.local_num_experts = moe_config.num_local_experts
39 + self.ep_rank = moe_config.moe_parallel_config.ep_rank
40 +
41 + @staticmethod
42 + def activation_format() -> mk.FusedMoEActivationFormat:
43 + return mk.FusedMoEActivationFormat.Standard
44 +
45 + @staticmethod
46 + def _supports_current_device() -> bool:
47 + """Supports only Blackwell-family GPUs."""
48 + p = current_platform
49 + return (
50 + p.is_cuda()
51 + and p.is_device_capability_family(100)
52 + and has_flashinfer_trtllm_fused_moe()
53 + )
54 +
55 + @staticmethod
56 + def _supports_no_act_and_mul() -> bool:
57 + """BF16 kernels do not support non-gated MoE"""
58 + return False
59 +
60 + @staticmethod
61 + def _supports_quant_scheme(
62 + weight_key: QuantKey | None,
63 + activation_key: QuantKey | None,
64 + ) -> bool:
65 + """Supports only unquantized inputs."""
66 + return weight_key is None and activation_key is None
67 +
68 + @staticmethod
69 + def _supports_activation(activation: MoEActivation) -> bool:
70 + return activation in [MoEActivation.SILU]
71 +
72 + @staticmethod
73 + def _supports_routing_method(
74 + routing_method: RoutingMethodType,
75 + weight_key: QuantKey | None,
76 + activation_key: QuantKey | None,
77 + ) -> bool:
78 + return routing_method in [
79 + RoutingMethodType.Default,
80 + RoutingMethodType.DeepSeekV3,
81 + RoutingMethodType.Llama4,
82 + # NOTE: TRTLLM Kernel has issue with Qwen3.5 router.
83 + # Re-enable once the issue is resolved.
84 + # https://github.com/vllm-project/vllm/issues/37591
85 + # RoutingMethodType.Renormalize,
86 + # RoutingMethodType.RenormalizeNaive
87 + ]
88 +
89 + @staticmethod
90 + def _supports_parallel_config(
91 + moe_parallel_config: FusedMoEParallelConfig,
92 + ) -> bool:
93 + """Monolithic kernel so only use with naive DP/EP and TP."""
94 + return (
95 + not moe_parallel_config.use_all2all_kernels
96 + or moe_parallel_config.use_ag_rs_all2all_kernels
97 + ) and not moe_parallel_config.enable_eplb
98 +
99 + @staticmethod
100 + def _supports_router_logits_dtype(
101 + router_logits_dtype: torch.dtype | None,
102 + routing_method: RoutingMethodType,
103 + ) -> bool:
104 + return True
105 +
106 + def supports_chunking(self) -> bool:
107 + return False
108 +
109 + def supports_expert_map(self) -> bool:
110 + return False
111 +
112 + @property
113 + def expects_unquantized_inputs(self) -> bool:
114 + return True
115 +
116 + def apply(
117 + self,
118 + hidden_states: torch.Tensor,
119 + w1: torch.Tensor,
120 + w2: torch.Tensor,
121 + router_logits: torch.Tensor,
122 + activation: MoEActivation,
123 + global_num_experts: int,
124 + expert_map: torch.Tensor | None,
125 + a1q_scale: torch.Tensor | None,
126 + apply_router_weight_on_input: bool,
127 + num_expert_group: int | None = None,
128 + e_score_correction_bias: torch.Tensor | None = None,
129 + routed_scaling_factor: float | None = None,
130 + topk_group: int | None = None,
131 + ) -> torch.Tensor:
132 + import flashinfer
133 +
134 + return flashinfer.fused_moe.trtllm_bf16_moe(
135 + routing_logits=router_logits,
136 + routing_bias=e_score_correction_bias,
137 + hidden_states=hidden_states,
138 + gemm1_weights=w1,
139 + gemm2_weights=w2,
140 + num_experts=global_num_experts,
141 + top_k=self.topk,
142 + n_group=num_expert_group,
143 + topk_group=topk_group,
144 + intermediate_size=self.intermediate_size_per_partition,
145 + local_expert_offset=self.ep_rank * self.local_num_experts,
146 + local_num_experts=self.local_num_experts,
147 + routing_method_type=self.routing_method_type,
148 + )
vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py deleted
-141
@@ -1,141 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import torch
5 -
6 -import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7 -from vllm.model_executor.layers.fused_moe.activation import MoEActivation
8 -from vllm.model_executor.layers.fused_moe.config import (
9 - FusedMoEConfig,
10 - FusedMoEParallelConfig,
11 - RoutingMethodType,
12 -)
13 -from vllm.platforms import current_platform
14 -from vllm.utils.torch_utils import direct_register_custom_op
15 -
16 -#
17 -# Methods used by the oracle for kernel selection.
18 -#
19 -
20 -
21 -def _supports_current_device() -> bool:
22 - """Supports only Blackwell-family GPUs."""
23 - p = current_platform
24 - return p.is_cuda() and p.is_device_capability_family(100)
25 -
26 -
27 -def _supports_no_act_and_mul() -> bool:
28 - """BF16 kernels do not support non-gated MoE"""
29 - return False
30 -
31 -
32 -def _supports_activation(activation: MoEActivation) -> bool:
33 - return activation in [MoEActivation.SILU]
34 -
35 -
36 -def _supports_routing_method_bf16(
37 - routing_method: RoutingMethodType,
38 -) -> bool:
39 - return routing_method in [
40 - RoutingMethodType.Default,
41 - RoutingMethodType.Renormalize,
42 - RoutingMethodType.DeepSeekV3,
43 - RoutingMethodType.Llama4,
44 - RoutingMethodType.RenormalizeNaive,
45 - ]
46 -
47 -
48 -def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
49 - """Supports TRTLLM Kernel does not support EPLB."""
50 - return not moe_parallel_config.enable_eplb
51 -
52 -
53 -def is_supported_config_trtllm_bf16(
54 - moe_config: FusedMoEConfig,
55 - activation_format: mk.FusedMoEActivationFormat,
56 -) -> tuple[bool, str | None]:
57 - """
58 - This method mirrors mk.FusedMoEPermuteExpertsUnpermute.is_supported_config
59 - for BF16 unquantized kernels.
60 - """
61 -
62 - def _make_reason(reason: str) -> str:
63 - return f"kernel does not support {reason}"
64 -
65 - if not _supports_current_device():
66 - return False, _make_reason(f"current device {current_platform.device_name}")
67 - elif not (moe_config.is_act_and_mul or _supports_no_act_and_mul()):
68 - return False, _make_reason("no act_and_mul MLP layer")
69 - elif not _supports_activation(moe_config.activation):
70 - return False, _make_reason(f"{moe_config.activation} activation")
71 - elif not _supports_parallel_config(moe_config.moe_parallel_config):
72 - return False, _make_reason(f"parallel config {moe_config.moe_parallel_config}")
73 - elif not _supports_routing_method_bf16(moe_config.routing_method):
74 - return False, _make_reason(f"routing method {moe_config.routing_method}")
75 - elif activation_format != mk.FusedMoEActivationFormat.Standard:
76 - return False, _make_reason(f"activation format {activation_format}")
77 -
78 - return True, None
79 -
80 -
81 -def flashinfer_fused_moe_bf16(
82 - routing_logits: torch.Tensor,
83 - routing_bias: torch.Tensor | None,
84 - hidden_states: torch.Tensor,
85 - gemm1_weights: torch.Tensor,
86 - gemm2_weights: torch.Tensor,
87 - num_experts: int,
88 - top_k: int,
89 - n_group: int | None,
90 - topk_group: int | None,
91 - intermediate_size: int,
92 - local_expert_offset: int,
93 - local_num_experts: int,
94 - routing_method_type: int,
95 - tune_max_num_tokens: int = 8192,
96 -) -> torch.Tensor:
97 - from vllm.utils.flashinfer import flashinfer_trtllm_bf16_moe
98 -
99 - return flashinfer_trtllm_bf16_moe(
100 - routing_logits=routing_logits,
101 - routing_bias=routing_bias,
102 - hidden_states=hidden_states,
103 - gemm1_weights=gemm1_weights,
104 - gemm2_weights=gemm2_weights,
105 - num_experts=num_experts,
106 - top_k=top_k,
107 - n_group=n_group,
108 - topk_group=topk_group,
109 - intermediate_size=intermediate_size,
110 - local_expert_offset=local_expert_offset,
111 - local_num_experts=local_num_experts,
112 - routing_method_type=routing_method_type,
113 - tune_max_num_tokens=tune_max_num_tokens,
114 - )
115 -
116 -
117 -def flashinfer_fused_moe_bf16_fake(
118 - routing_logits: torch.Tensor,
119 - routing_bias: torch.Tensor | None,
120 - hidden_states: torch.Tensor,
121 - gemm1_weights: torch.Tensor,
122 - gemm2_weights: torch.Tensor,
123 - num_experts: int,
124 - top_k: int,
125 - n_group: int | None,
126 - topk_group: int | None,
127 - intermediate_size: int,
128 - local_expert_offset: int,
129 - local_num_experts: int,
130 - routing_method_type: int = RoutingMethodType.Renormalize,
131 - tune_max_num_tokens: int = 8192,
132 -) -> torch.Tensor:
133 - return torch.empty_like(hidden_states)
134 -
135 -
136 -direct_register_custom_op(
137 - op_name="flashinfer_fused_moe_bf16",
138 - op_func=flashinfer_fused_moe_bf16,
139 - fake_impl=flashinfer_fused_moe_bf16_fake,
140 - tags=(torch.Tag.needs_fixed_stride_order,),
141 -)
vllm/model_executor/layers/fused_moe/fused_moe.py
+4
@@ -1967,6 +1967,10 @@ class TritonExperts(mk.FusedMoEExpertsModular):
1967 or moe_parallel_config.use_fi_nvl_one_sided_kernels
1968 )
1969
1970 + @staticmethod
1971 + def _supports_batch_invariance():
1972 + return True
1973 +
1974 def supports_expert_map(self) -> bool:
1975 return True
1976
vllm/model_executor/layers/fused_moe/modular_kernel.py
+12
@@ -9,6 +9,7 @@ from typing import final
9
10 import torch
11
12 +import vllm.envs as envs
13 from vllm.logger import init_logger
14 from vllm.model_executor.layers.fused_moe.activation import (
15 MoEActivation,
@@ -563,6 +564,8 @@ class FusedMoEExperts(ABC):
564 )
565 elif activation_format != cls.activation_format():
566 return False, _make_reason(f"{activation_format.value} activation format")
567 + elif envs.VLLM_BATCH_INVARIANT and not cls._supports_batch_invariance():
568 + return False, _make_reason("batch invariance")
569 return True, None
570
571 @staticmethod
@@ -645,6 +648,15 @@ class FusedMoEExperts(ABC):
648 """
649 return True
650
651 + @staticmethod
652 + def _supports_batch_invariance() -> bool:
653 + """
654 + Whether the kernel supports batch invariance, i.e. the output does not
655 + depend on the order of the tokens in the input batch. This is useful
656 + for determining if the kernel can used with VLLM_BATCH_INVARIANT=1.
657 + """
658 + return False
659 +
660 #
661 # Various helpers for accessing quantization parameters from the
662 # quant_config.
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+267 -163
@@ -11,21 +11,20 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11 from vllm._aiter_ops import rocm_aiter_ops
12 from vllm.config.kernel import MoEBackend
13 from vllm.logger import init_logger
14 +from vllm.model_executor.layers.fused_moe.all2all_utils import (
15 + maybe_make_prepare_finalize,
16 +)
17 from vllm.model_executor.layers.fused_moe.config import (
18 FusedMoEConfig,
19 FusedMoEQuantConfig,
20 )
18 -from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe import (
19 - is_supported_config_trtllm_bf16,
20 -)
21 -from vllm.model_executor.layers.fused_moe.prepare_finalize import (
22 - MoEPrepareAndFinalizeNoDPEPModular,
23 -)
21 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
22 + FlashinferMoeBackend,
23 + convert_moe_weights_to_flashinfer_trtllm_block_layout,
24 + get_flashinfer_moe_backend,
25 swap_w13_to_w31,
26 )
27 from vllm.platforms import current_platform
28 -from vllm.utils.flashinfer import has_flashinfer, has_flashinfer_cutlass_fused_moe
28
29 logger = init_logger(__name__)
30
@@ -35,21 +34,96 @@ class UnquantizedMoeBackend(Enum):
34 FLASHINFER_CUTLASS = "FlashInfer CUTLASS"
35 AITER = "ROCm AITER"
36 TRITON = "TRITON"
37 + BATCHED_TRITON = "BATCHED_TRITON"
38 CPU = "CPU"
39 XPU = "XPU"
40 TPU = "TPU"
41 OOT = "OOT"
42
43
44 -# NOTE(zyongye): Unsupported backend means backend
45 -# that is not conform with Modular kernel format.
46 -# We will directly call the kernel for those backend
47 -UNSUPPORTED_BACKEND = [
48 - UnquantizedMoeBackend.FLASHINFER_TRTLLM,
49 - UnquantizedMoeBackend.CPU,
50 - UnquantizedMoeBackend.TPU,
51 - UnquantizedMoeBackend.OOT,
52 -]
44 +def _get_priority_backends(moe_config: FusedMoEConfig) -> list[UnquantizedMoeBackend]:
45 + """
46 + Get available backends in priority order based on platform and config.
47 +
48 + This function can be extended to become more complex as needed.
49 + """
50 +
51 + def _move_to_back(
52 + backends: list[UnquantizedMoeBackend],
53 + backend: UnquantizedMoeBackend,
54 + ) -> None:
55 + backends.append(backends.pop(backends.index(backend)))
56 +
57 + if current_platform.is_rocm():
58 + _AVAILABLE_BACKENDS = [
59 + UnquantizedMoeBackend.AITER,
60 + UnquantizedMoeBackend.TRITON,
61 + UnquantizedMoeBackend.BATCHED_TRITON,
62 + ]
63 + elif current_platform.is_cuda():
64 + _AVAILABLE_BACKENDS = [
65 + UnquantizedMoeBackend.FLASHINFER_TRTLLM,
66 + UnquantizedMoeBackend.FLASHINFER_CUTLASS,
67 + UnquantizedMoeBackend.TRITON,
68 + UnquantizedMoeBackend.BATCHED_TRITON,
69 + ]
70 +
71 + # HACK: Qwen3.5 has crash with FLASHINFER_CUTLASS BF16 if DEP.
72 + # Updating the oracle querying logic is out of the scope of this
73 + # PR. Need to fix the kernel or update structure in follow up.
74 + if moe_config.moe_parallel_config.dp_size > 1:
75 + _move_to_back(_AVAILABLE_BACKENDS, UnquantizedMoeBackend.FLASHINFER_CUTLASS)
76 +
77 + elif current_platform.is_xpu():
78 + _AVAILABLE_BACKENDS = [UnquantizedMoeBackend.XPU]
79 + elif current_platform.is_cpu():
80 + _AVAILABLE_BACKENDS = [UnquantizedMoeBackend.CPU]
81 + return _AVAILABLE_BACKENDS
82 +
83 +
84 +def backend_to_kernel_cls(
85 + backend: UnquantizedMoeBackend,
86 +) -> type[mk.FusedMoEExperts]:
87 + if backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM:
88 + from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import (
89 + TrtLlmBf16Experts,
90 + )
91 +
92 + return TrtLlmBf16Experts
93 +
94 + elif backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS:
95 + from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
96 + FlashInferExperts,
97 + )
98 +
99 + return FlashInferExperts
100 +
101 + elif backend == UnquantizedMoeBackend.AITER:
102 + from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
103 + AiterExperts,
104 + )
105 +
106 + return AiterExperts
107 +
108 + elif backend == UnquantizedMoeBackend.TRITON:
109 + from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
110 +
111 + return TritonExperts
112 +
113 + elif backend == UnquantizedMoeBackend.BATCHED_TRITON:
114 + from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
115 + BatchedTritonExperts,
116 + )
117 +
118 + return BatchedTritonExperts
119 +
120 + elif backend == UnquantizedMoeBackend.XPU:
121 + from vllm.model_executor.layers.fused_moe.xpu_fused_moe import XPUExperts
122 +
123 + return XPUExperts
124 +
125 + else:
126 + raise ValueError(f"Unknown unquantized MoE backend: {backend.value}")
127
128
129 def map_unquantized_backend(runner_backend: MoEBackend) -> UnquantizedMoeBackend:
@@ -70,194 +144,224 @@ def map_unquantized_backend(runner_backend: MoEBackend) -> UnquantizedMoeBackend
144
145 def select_unquantized_moe_backend(
146 moe_config: FusedMoEConfig,
73 - use_dp: bool,
74 -) -> UnquantizedMoeBackend:
147 +) -> tuple[UnquantizedMoeBackend, type[mk.FusedMoEExperts] | None]:
148 """
76 - Select the primary Unquantized MoE backend
149 + Select the primary Unquantized MoE backend.
150 Note: Shape-specific fallbacks may still occur at runtime.
151 """
152
80 - def _make_log_backend(backend: UnquantizedMoeBackend):
81 - return f"Using {backend.value} backend for Unquantized MoE"
153 + if current_platform.is_cpu():
154 + # TODO: migrate to MK structure.
155 + return UnquantizedMoeBackend.CPU, None
156 +
157 + if current_platform.is_tpu():
158 + return UnquantizedMoeBackend.TPU, None
159 +
160 + if current_platform.is_out_of_tree():
161 + return UnquantizedMoeBackend.OOT, None
162 +
163 + # NOTE: the kernels are selected in the following order.
164 + AVAILABLE_BACKENDS = _get_priority_backends(moe_config)
165
166 + # NOTE(rob): We need to peak into the P/F selection to determine
167 + # if we are using the batched or standard expert format, which
168 + # if not ideal. Once we unify TP + DP/EP, we can select P/F first.
169 activation_format = (
170 mk.FusedMoEActivationFormat.BatchedExperts
171 if moe_config.moe_parallel_config.use_batched_activation_format
172 else mk.FusedMoEActivationFormat.Standard
173 )
174
89 - # Check if FlashInfer TRTLLM BF16 MoE is supported
90 - trtllm_supported, _ = is_supported_config_trtllm_bf16(
91 - moe_config=moe_config,
92 - activation_format=activation_format,
93 - )
94 - flashinfer_trtllm_available = has_flashinfer() and trtllm_supported
95 - # FlashInfer CUTLASS MoE is only supported on Hopper and later GPUS
96 - flashinfer_cutlass_available = (
97 - has_flashinfer_cutlass_fused_moe()
98 - and (not use_dp)
99 - and current_platform.has_device_capability(90)
100 - )
101 - flashinfer_trtllm_moe_enabled = (
102 - flashinfer_trtllm_available
103 - and envs.VLLM_USE_FLASHINFER_MOE_FP16
104 - and envs.VLLM_FLASHINFER_MOE_BACKEND == "latency"
105 - )
106 - flashinfer_cutlass_moe_enabled = (
107 - flashinfer_cutlass_available and envs.VLLM_USE_FLASHINFER_MOE_FP16
108 - )
109 - rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
175 + def _make_log_backend(backend: UnquantizedMoeBackend) -> str:
176 + available_strs = [b.value for b in AVAILABLE_BACKENDS]
177 + return (
178 + f"Using {backend.value} Unquantized MoE backend out "
179 + f"of potential backends: {available_strs}."
180 + )
181 +
182 + def _make_log_unsupported(
183 + backend: UnquantizedMoeBackend, reason: str | None
184 + ) -> str:
185 + if reason:
186 + return (
187 + f"Unquantized MoE backend {backend.value} does not support the "
188 + f"deployment configuration since {reason}."
189 + )
190 + return (
191 + f"Unquantized MoE backend '{backend.value}' does not support the "
192 + "deployment configuration."
193 + )
194 +
195 + def _return_or_raise(
196 + backend: UnquantizedMoeBackend,
197 + config: FusedMoEConfig,
198 + activation_format: mk.FusedMoEActivationFormat,
199 + ) -> tuple[UnquantizedMoeBackend, type[mk.FusedMoEExperts] | None]:
200 + k_cls = backend_to_kernel_cls(backend)
201 + supported, reason = k_cls.is_supported_config(
202 + k_cls, config, None, None, activation_format
203 + )
204 + if supported:
205 + logger.info_once(_make_log_backend(backend), scope="local")
206 + return backend, k_cls
207 + raise ValueError(_make_log_unsupported(backend, reason))
208
111 - # Handle explicit moe_backend from user.
209 runner_backend = moe_config.moe_backend
210 if runner_backend != "auto":
211 requested_backend = map_unquantized_backend(runner_backend)
115 - if requested_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM:
116 - if not flashinfer_trtllm_available:
212 + if (
213 + activation_format == mk.FusedMoEActivationFormat.BatchedExperts
214 + and requested_backend == UnquantizedMoeBackend.TRITON
215 + ):
216 + requested_backend = UnquantizedMoeBackend.BATCHED_TRITON
217 +
218 + return _return_or_raise(requested_backend, moe_config, activation_format)
219 +
220 + # Handle explicit FlashInfer FP16 configuration.
221 + if envs.is_set("VLLM_USE_FLASHINFER_MOE_FP16"):
222 + if not envs.VLLM_USE_FLASHINFER_MOE_FP16:
223 + if UnquantizedMoeBackend.FLASHINFER_TRTLLM in AVAILABLE_BACKENDS:
224 + AVAILABLE_BACKENDS.remove(UnquantizedMoeBackend.FLASHINFER_TRTLLM)
225 + if UnquantizedMoeBackend.FLASHINFER_CUTLASS in AVAILABLE_BACKENDS:
226 + AVAILABLE_BACKENDS.remove(UnquantizedMoeBackend.FLASHINFER_CUTLASS)
227 +
228 + elif envs.is_set("VLLM_FLASHINFER_MOE_BACKEND"):
229 + # If user is explicit about backend, validate it.
230 + fi_backend = get_flashinfer_moe_backend()
231 + if fi_backend == FlashinferMoeBackend.CUTLASS:
232 + backend = UnquantizedMoeBackend.FLASHINFER_CUTLASS
233 + elif fi_backend == FlashinferMoeBackend.TENSORRT_LLM:
234 + backend = UnquantizedMoeBackend.FLASHINFER_TRTLLM
235 + else:
236 raise ValueError(
118 - "FlashInfer TRTLLM MoE backend is not available for this "
119 - "configuration."
237 + f"FlashInfer MOE backend {fi_backend} "
238 + "does not support unquantized MoE."
239 )
121 - elif requested_backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS:
122 - if not flashinfer_cutlass_available:
123 - raise ValueError(
124 - "FlashInfer CUTLASS MoE backend is not available for this "
125 - "configuration."
240 + k_cls = backend_to_kernel_cls(backend)
241 + return _return_or_raise(backend, moe_config, activation_format)
242 + else:
243 + # If the user is not explicit about the backend, try both.
244 + for backend in [
245 + UnquantizedMoeBackend.FLASHINFER_TRTLLM,
246 + UnquantizedMoeBackend.FLASHINFER_CUTLASS,
247 + ]:
248 + k_cls = backend_to_kernel_cls(backend)
249 + supported, reason = k_cls.is_supported_config(
250 + k_cls, moe_config, None, None, activation_format
251 )
127 - elif requested_backend == UnquantizedMoeBackend.AITER and not (
128 - current_platform.is_rocm() and rocm_aiter_moe_enabled
129 - ):
130 - raise ValueError(
131 - "ROCm AITer MoE backend is not available for this configuration."
252 + if supported:
253 + logger.info_once(_make_log_backend(backend), scope="local")
254 + return backend, k_cls
255 + else:
256 + logger.debug_once(
257 + _make_log_unsupported(backend, reason), scope="local"
258 + )
259 +
260 + raise NotImplementedError(
261 + "Found VLLM_USE_FLASHINFER_MOE_FP16=1, but no "
262 + "FlashInfer unquantized MoE backend supports the configuration."
263 )
133 - logger.info_once(_make_log_backend(requested_backend), scope="local")
134 - return requested_backend
264
136 - if current_platform.is_rocm():
137 - if rocm_aiter_moe_enabled:
138 - backend = UnquantizedMoeBackend.AITER
265 + # Handle explicit AITER FP8 configuration.
266 + if envs.is_set("VLLM_ROCM_USE_AITER") or envs.is_set("VLLM_ROCM_USE_AITER_MOE"):
267 + if not envs.VLLM_ROCM_USE_AITER or not envs.VLLM_ROCM_USE_AITER_MOE:
268 + if UnquantizedMoeBackend.AITER in AVAILABLE_BACKENDS:
269 + AVAILABLE_BACKENDS.remove(UnquantizedMoeBackend.AITER)
270 else:
140 - backend = UnquantizedMoeBackend.TRITON
141 - if current_platform.is_cuda():
142 - if flashinfer_trtllm_moe_enabled:
143 - backend = UnquantizedMoeBackend.FLASHINFER_TRTLLM
144 - elif flashinfer_cutlass_moe_enabled:
145 - backend = UnquantizedMoeBackend.FLASHINFER_CUTLASS
146 - if trtllm_supported:
147 - logger.info_once(
148 - "FlashInfer TRTLLM MoE is available but not enabled, "
149 - "consider setting VLLM_FLASHINFER_MOE_BACKEND=latency "
150 - "to enable it for better performance.",
151 - scope="local",
152 - )
153 - else:
154 - if not envs.VLLM_USE_FLASHINFER_MOE_FP16 and trtllm_supported:
155 - logger.info_once(
156 - "FlashInfer TRTLLM MoE is available but not enabled, "
157 - "consider setting VLLM_USE_FLASHINFER_MOE_FP16=1 "
158 - "and VLLM_FLASHINFER_MOE_BACKEND=latency "
159 - "to enable it for better performance.",
160 - scope="local",
161 - )
162 - elif not use_dp and flashinfer_cutlass_available:
163 - logger.info_once(
164 - "FlashInfer CUTLASS MoE is available"
165 - " but not enabled, consider setting"
166 - " VLLM_USE_FLASHINFER_MOE_FP16=1 to enable it.",
167 - scope="local",
168 - )
169 - elif use_dp:
170 - logger.info_once(
171 - "FlashInfer CUTLASS MoE is currently not available for DP.",
172 - scope="local",
173 - )
174 - backend = UnquantizedMoeBackend.TRITON
175 - if current_platform.is_xpu():
176 - backend = UnquantizedMoeBackend.XPU
177 - if current_platform.is_cpu():
178 - backend = UnquantizedMoeBackend.CPU
179 - if current_platform.is_tpu():
180 - backend = UnquantizedMoeBackend.TPU
181 - if current_platform.is_out_of_tree():
182 - backend = UnquantizedMoeBackend.OOT
271 + backend = UnquantizedMoeBackend.AITER
272 + return _return_or_raise(backend, moe_config, activation_format)
273
184 - logger.info_once(_make_log_backend(backend), scope="local")
185 - return backend
274 + for backend in AVAILABLE_BACKENDS:
275 + k_cls = backend_to_kernel_cls(backend)
276 + supported, reason = k_cls.is_supported_config(
277 + k_cls, moe_config, None, None, activation_format
278 + )
279 + if supported:
280 + logger.info_once(_make_log_backend(backend), scope="local")
281 + return backend, k_cls
282 +
283 + logger.debug_once(_make_log_unsupported(backend, reason), scope="local")
284 +
285 + raise NotImplementedError(
286 + "No Unquantized MoE backend supports the deployment configuration."
287 + )
288
289
290 def convert_to_unquantized_kernel_format(
291 unquantized_backend: UnquantizedMoeBackend,
292 layer: Module,
191 - w13_weight: torch.Tensor | None = None,
192 - w2_weight: torch.Tensor | None = None,
293 + w13_weight: torch.Tensor,
294 + w2_weight: torch.Tensor,
295 ) -> tuple[torch.Tensor, torch.Tensor]:
296 if unquantized_backend == UnquantizedMoeBackend.AITER:
195 - w13_weight, w2_weight = rocm_aiter_ops.shuffle_weights(
196 - layer.w13_weight.data, layer.w2_weight.data
197 - )
297 + w13_weight, w2_weight = rocm_aiter_ops.shuffle_weights(w13_weight, w2_weight)
298
299 elif unquantized_backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS:
300 + if layer.moe_config.is_act_and_mul:
301 + # Swap halves to arrange as [w3; w1] (kernel expectation)
302 + # Non-gated MoE: w13 is a single projection, no need to swap.
303 + w13_weight = swap_w13_to_w31(w13_weight)
304 +
305 + elif unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM:
306 # Swap halves to arrange as [w3; w1] (kernel expectation)
201 - w13_weight = swap_w13_to_w31(layer.w13_weight.data)
307 + w13_weight = swap_w13_to_w31(w13_weight)
308 + _cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
309 + w13_weight, w2_weight = convert_moe_weights_to_flashinfer_trtllm_block_layout(
310 + _cache_permute_indices,
311 + w13_weight,
312 + w2_weight,
313 + )
314
203 - return w13_weight, w2_weight
315 + return w13_weight.contiguous(), w2_weight.contiguous()
316
317
318 def make_unquantized_moe_kernel(
207 - backend: UnquantizedMoeBackend,
319 quant_config: FusedMoEQuantConfig,
320 moe_config: FusedMoEConfig,
210 -) -> mk.FusedMoEKernel | None:
211 - if backend in UNSUPPORTED_BACKEND:
212 - return None
321 + backend: UnquantizedMoeBackend,
322 + experts_cls: type[mk.FusedMoEExperts],
323 + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
324 + shared_experts: torch.nn.Module | None = None,
325 +) -> mk.FusedMoEKernel:
326 + # Create Prepare/Finalize
327 + is_monolithic = issubclass(experts_cls, mk.FusedMoEExpertsMonolithic)
328 + prepare_finalize = maybe_make_prepare_finalize(
329 + moe=moe_config,
330 + quant_config=quant_config,
331 + routing_tables=routing_tables,
332 + allow_new_interface=True,
333 + use_monolithic=is_monolithic,
334 + )
335 + assert prepare_finalize is not None
336
214 - if backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS:
215 - from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
216 - FlashInferExperts,
217 - )
337 + logger.info_once("Using %s", prepare_finalize.__class__.__name__, scope="local")
338
219 - kernel = mk.FusedMoEKernel(
220 - MoEPrepareAndFinalizeNoDPEPModular(),
221 - FlashInferExperts(
222 - moe_config=moe_config,
223 - quant_config=quant_config,
224 - ),
225 - inplace=False,
339 + # Create Experts
340 + if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
341 + max_num_tokens = prepare_finalize.max_num_tokens_per_rank()
342 + assert max_num_tokens is not None
343 + experts = experts_cls(
344 + moe_config=moe_config,
345 + quant_config=quant_config,
346 + max_num_tokens=max_num_tokens,
347 + num_dispatchers=prepare_finalize.num_dispatchers(),
348 )
227 -
228 - elif backend == UnquantizedMoeBackend.AITER:
229 - from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
230 - AiterExperts,
349 + else:
350 + experts = experts_cls(
351 + moe_config=moe_config,
352 + quant_config=quant_config,
353 )
354
233 - kernel = mk.FusedMoEKernel(
234 - MoEPrepareAndFinalizeNoDPEPModular(),
235 - AiterExperts(
236 - moe_config=moe_config,
237 - quant_config=quant_config,
238 - ),
239 - inplace=not moe_config.disable_inplace,
240 - )
241 - elif backend == UnquantizedMoeBackend.TRITON:
242 - from vllm.model_executor.layers.fused_moe import TritonExperts
243 -
244 - kernel = mk.FusedMoEKernel(
245 - MoEPrepareAndFinalizeNoDPEPModular(),
246 - TritonExperts(
247 - moe_config=moe_config,
248 - quant_config=quant_config,
249 - ),
250 - inplace=not moe_config.disable_inplace,
251 - )
252 - elif backend == UnquantizedMoeBackend.XPU:
253 - from vllm.model_executor.layers.fused_moe import XPUExperts
254 -
255 - kernel = mk.FusedMoEKernel(
256 - MoEPrepareAndFinalizeNoDPEPModular(),
257 - XPUExperts(
258 - moe_config=moe_config,
259 - quant_config=quant_config,
260 - ),
261 - inplace=not moe_config.disable_inplace,
262 - )
355 + kernel = mk.FusedMoEKernel(
356 + prepare_finalize,
357 + experts,
358 + shared_experts=(
359 + shared_experts
360 + if moe_config.moe_parallel_config.use_deepep_ll_kernels
361 + else None
362 + ),
363 + moe_parallel_config=moe_config.moe_parallel_config,
364 + inplace=(not moe_config.disable_inplace and not is_monolithic),
365 + )
366 +
367 return kernel
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+87 -151
@@ -6,11 +6,8 @@ from collections.abc import Callable
6 import torch
7 import torch.nn.functional as F
8 from torch.nn import Module
9 -from torch.nn.parameter import Parameter
9
10 import vllm.envs as envs
12 -import vllm.model_executor.layers.fused_moe.modular_kernel as mk
13 -from vllm._aiter_ops import rocm_aiter_ops
11 from vllm.logger import init_logger
12 from vllm.model_executor.custom_op import CustomOp
13 from vllm.model_executor.layers.fused_moe.config import (
@@ -23,7 +20,6 @@ from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
20 FusedMoEMethodBase,
21 )
22 from vllm.model_executor.layers.fused_moe.modular_kernel import (
26 - FusedMoEActivationFormat,
23 FusedMoEExpertsModular,
24 FusedMoEPrepareAndFinalizeModular,
25 )
@@ -33,20 +29,10 @@ from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
29 make_unquantized_moe_kernel,
30 select_unquantized_moe_backend,
31 )
36 -from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
37 - convert_moe_weights_to_flashinfer_trtllm_block_layout,
38 -)
32 from vllm.model_executor.utils import replace_parameter, set_weight_attrs
33 from vllm.platforms import current_platform
34 from vllm.platforms.interface import CpuArchEnum
35
43 -if current_platform.is_cuda_alike() or current_platform.is_xpu():
44 - from .fused_batched_moe import BatchedTritonExperts
45 - from .fused_moe import TritonExperts
46 -else:
47 - TritonExperts = None # type: ignore
48 -
49 -
36 logger = init_logger(__name__)
37
38
@@ -59,45 +45,16 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
45
46 def __init__(self, moe: FusedMoEConfig):
47 super().__init__(moe)
62 - self.unquantized_backend = select_unquantized_moe_backend(
48 + self.unquantized_backend, self.experts_cls = select_unquantized_moe_backend(
49 moe_config=self.moe,
64 - use_dp=self.moe.moe_parallel_config.dp_size > 1,
65 - )
66 -
67 - # AITER only supports gated activations (silu/gelu), so disable it
68 - # for non-gated MoE (is_act_and_mul=False)
69 - self.rocm_aiter_moe_enabled = (
70 - rocm_aiter_ops.is_fused_moe_enabled() and moe.is_act_and_mul
50 )
72 - self.kernel: mk.FusedMoEKernel | None = None
73 - self._is_monolithic = (
74 - current_platform.is_cpu()
75 - or self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM
76 - )
77 -
78 - if self.is_monolithic:
79 - self.apply_monolithic: Callable = self._select_monolithic()
80 -
81 - def _select_monolithic(self) -> Callable:
82 - """Select the monolithic implementation based on platform."""
83 - if current_platform.is_cpu():
84 - return self.forward_monolithic_cpu
85 - else:
86 - return self.forward_monolithic_cuda
87 -
88 - def forward_native(
89 - self,
90 - layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
91 - x: torch.Tensor,
92 - topk_weights: torch.Tensor,
93 - topk_ids: torch.Tensor,
94 - shared_experts_input: torch.Tensor | None,
95 - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
96 - return self.forward_cuda(layer, x, topk_weights, topk_ids, shared_experts_input)
51
52 @property
53 def is_monolithic(self) -> bool:
100 - return self._is_monolithic
54 + # Escape hatch for CPU, which stays on the old monolithic path.
55 + if self.unquantized_backend == UnquantizedMoeBackend.CPU:
56 + return True
57 + return super().is_monolithic
58
59 @property
60 def supports_eplb(self) -> bool:
@@ -106,43 +63,22 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
63 def maybe_make_prepare_finalize(
64 self,
65 routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
109 - ) -> FusedMoEPrepareAndFinalizeModular | None:
110 - return super().maybe_make_prepare_finalize(routing_tables)
66 + ):
67 + raise ValueError(
68 + f"{self.__class__.__name__} uses the new modular kernel initialization "
69 + "logic for all but the CPU backend. CPU backend is monolithic. "
70 + "So this function should not be called."
71 + )
72
73 def select_gemm_impl(
74 self,
75 prepare_finalize: FusedMoEPrepareAndFinalizeModular,
76 layer: torch.nn.Module,
77 ) -> FusedMoEExpertsModular:
117 - assert self.moe_quant_config is not None
118 - if (
119 - prepare_finalize.activation_format
120 - == FusedMoEActivationFormat.BatchedExperts
121 - ):
122 - logger.debug("BatchedTritonExperts %s", self.moe)
123 - return BatchedTritonExperts(
124 - moe_config=self.moe,
125 - quant_config=self.moe_quant_config,
126 - max_num_tokens=self.moe.max_num_tokens,
127 - num_dispatchers=prepare_finalize.num_dispatchers(),
128 - )
129 - elif (
130 - self.unquantized_backend == UnquantizedMoeBackend.AITER
131 - and rocm_aiter_ops.is_fused_moe_enabled()
132 - ):
133 - from .rocm_aiter_fused_moe import AiterExperts
134 -
135 - logger.debug("AiterExperts %s", self.moe)
136 - return AiterExperts(
137 - moe_config=self.moe,
138 - quant_config=self.moe_quant_config,
139 - )
140 - else:
141 - logger.debug("TritonExperts %s", self.moe)
142 - return TritonExperts(
143 - moe_config=self.moe,
144 - quant_config=self.moe_quant_config,
145 - )
78 + raise ValueError(
79 + f"{self.__class__.__name__} uses the new modular kernel initialization "
80 + "logic. This function should not be called."
81 + )
82
83 def create_weights(
84 self,
@@ -227,14 +163,17 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
163 replace_parameter(layer, "w13_weight", w13)
164 replace_parameter(layer, "w2_weight", w2)
165
230 - # Setup Modular Kernel for TP Case
166 + # Setup moe kernel.
167 self.moe_quant_config = self.get_fused_moe_quant_config(layer)
168 assert self.moe_quant_config is not None
233 -
234 - self.kernel = make_unquantized_moe_kernel(
235 - backend=self.unquantized_backend,
169 + assert self.experts_cls is not None
170 + self.moe_kernel = make_unquantized_moe_kernel(
171 quant_config=self.moe_quant_config,
172 moe_config=self.moe,
173 + backend=self.unquantized_backend,
174 + experts_cls=self.experts_cls,
175 + routing_tables=layer._maybe_init_expert_routing_tables(),
176 + shared_experts=layer.shared_experts,
177 )
178
179 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
@@ -244,22 +183,15 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
183 layer.w13_weight.data = self._maybe_pad_weight(layer.w13_weight.data)
184 layer.w2_weight.data = self._maybe_pad_weight(layer.w2_weight.data)
185
247 - if self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM:
248 - _cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
249 - # Swap halves to arrange as [w3; w1] (kernel expectation)
250 - w1_w, w3_w = torch.chunk(layer.w13_weight.data, 2, dim=1)
251 - w13_weight_swapped = torch.cat([w3_w, w1_w], dim=1)
252 - layer.w13_weight.data = w13_weight_swapped.contiguous()
253 - w13_weights_shuffled, w2_weights_shuffled = (
254 - convert_moe_weights_to_flashinfer_trtllm_block_layout(
255 - _cache_permute_indices,
256 - layer.w13_weight.data,
257 - layer.w2_weight.data,
258 - )
259 - )
260 - layer.w13_weight = Parameter(w13_weights_shuffled, requires_grad=False)
261 - layer.w2_weight = Parameter(w2_weights_shuffled, requires_grad=False)
186 + if self.unquantized_backend in [
187 + UnquantizedMoeBackend.TPU,
188 + UnquantizedMoeBackend.OOT,
189 + ]:
190 + # OOT handles internally.
191 + return
192 +
193 elif self.unquantized_backend == UnquantizedMoeBackend.CPU:
194 + # CPU stays on the old path — no oracle, no moe_kernel.
195 from vllm.model_executor.layers.fused_moe import cpu_fused_moe
196
197 if current_platform.get_cpu_architecture() == CpuArchEnum.X86:
@@ -290,13 +222,22 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
222 self.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
223 else:
224 self.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
293 - elif current_platform.is_cuda_alike() or current_platform.is_xpu():
225 + else:
226 self._setup_kernel(
227 layer=layer,
228 w13=layer.w13_weight,
229 w2=layer.w2_weight,
230 )
231
232 + def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
233 + if self.moe.has_bias:
234 + return biased_moe_quant_config(
235 + layer.w13_bias,
236 + layer.w2_bias,
237 + )
238 + else:
239 + return FUSED_MOE_UNQUANTIZED_CONFIG
240 +
241 def apply(
242 self,
243 layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
@@ -313,16 +254,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
254 shared_experts_input=shared_experts_input,
255 )
256
316 - def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
317 - if self.moe.has_bias:
318 - return biased_moe_quant_config(
319 - layer.w13_bias,
320 - layer.w2_bias,
321 - )
322 - else:
323 - return FUSED_MOE_UNQUANTIZED_CONFIG
324 -
325 - def forward_cuda(
257 + def forward_native(
258 self,
259 layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
260 x: torch.Tensor,
@@ -330,9 +262,8 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
262 topk_ids: torch.Tensor,
263 shared_experts_input: torch.Tensor | None,
264 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
333 - assert self.kernel is not None
334 -
335 - return self.kernel.apply(
265 + assert self.moe_kernel is not None
266 + return self.moe_kernel.apply(
267 hidden_states=x,
268 w1=layer.w13_weight,
269 w2=layer.w2_weight,
@@ -345,53 +276,58 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
276 shared_experts_input=shared_experts_input,
277 )
278
348 - def forward_monolithic_cuda(
279 + def forward_cuda(
280 self,
281 layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
282 x: torch.Tensor,
352 - router_logits: torch.Tensor,
283 + topk_weights: torch.Tensor,
284 + topk_ids: torch.Tensor,
285 + shared_experts_input: torch.Tensor | None,
286 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
354 - import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: F401
355 -
356 - assert self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM
357 -
358 - return torch.ops.vllm.flashinfer_fused_moe_bf16(
359 - routing_logits=router_logits,
360 - routing_bias=layer.e_score_correction_bias,
361 - hidden_states=x,
362 - gemm1_weights=layer.w13_weight,
363 - gemm2_weights=layer.w2_weight,
364 - num_experts=layer.global_num_experts,
365 - top_k=layer.top_k,
366 - n_group=layer.num_expert_group,
367 - topk_group=layer.topk_group,
368 - intermediate_size=layer.intermediate_size_per_partition,
369 - local_expert_offset=layer.ep_rank * layer.local_num_experts,
370 - local_num_experts=layer.local_num_experts,
371 - routing_method_type=layer.routing_method_type,
287 + return self.forward_native(
288 + layer, x, topk_weights, topk_ids, shared_experts_input
289 )
290
374 - def forward_monolithic_cpu(
291 + def apply_monolithic(
292 self,
293 layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
294 x: torch.Tensor,
295 router_logits: torch.Tensor,
296 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
380 - return self.cpu_fused_moe(
381 - layer,
382 - x,
383 - layer.use_grouped_topk,
384 - layer.top_k,
385 - router_logits,
386 - layer.renormalize,
387 - layer.topk_group,
388 - layer.num_expert_group,
389 - layer.global_num_experts,
390 - layer.expert_map,
391 - layer.custom_routing_function,
392 - layer.scoring_func,
393 - layer.routed_scaling_factor,
394 - layer.e_score_correction_bias,
395 - layer.apply_router_weight_on_input,
396 - layer.activation,
397 - )
297 + assert self.is_monolithic
298 + if self.unquantized_backend == UnquantizedMoeBackend.CPU:
299 + assert self.moe_kernel is None
300 + return self.cpu_fused_moe(
301 + layer,
302 + x,
303 + layer.use_grouped_topk,
304 + layer.top_k,
305 + router_logits,
306 + layer.renormalize,
307 + layer.topk_group,
308 + layer.num_expert_group,
309 + layer.global_num_experts,
310 + layer.expert_map,
311 + layer.custom_routing_function,
312 + layer.scoring_func,
313 + layer.routed_scaling_factor,
314 + layer.e_score_correction_bias,
315 + layer.apply_router_weight_on_input,
316 + layer.activation,
317 + )
318 + else:
319 + assert self.moe_kernel is not None
320 + return self.moe_kernel.apply_monolithic(
321 + x,
322 + layer.w13_weight,
323 + layer.w2_weight,
324 + router_logits,
325 + activation=layer.activation,
326 + global_num_experts=layer.global_num_experts,
327 + expert_map=layer.expert_map,
328 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
329 + num_expert_group=layer.num_expert_group,
330 + topk_group=layer.topk_group,
331 + e_score_correction_bias=layer.e_score_correction_bias,
332 + routed_scaling_factor=layer.routed_scaling_factor,
333 + )
vllm/utils/flashinfer.py
+1
@@ -202,6 +202,7 @@ def has_flashinfer_trtllm_fused_moe() -> bool:
202 ("flashinfer.fused_moe", "trtllm_fp8_per_tensor_scale_moe"),
203 ("flashinfer.fused_moe", "trtllm_fp4_block_scale_moe"),
204 ("flashinfer.fused_moe", "trtllm_mxint4_block_scale_moe"),
205 + ("flashinfer.fused_moe", "trtllm_bf16_moe"),
206 ]
207 for module_name, attr_name in required_functions:
208 mod = _get_submodule(module_name)