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 6, 07:00 UTC
[MoE Refactor] Mxfp4 oracle rebased (#37128)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Yongye Zhu committed
Mar 20, 2026 at 22:37 UTC
87bd91892f8c63ed9aeb2ad2e701472ead8be84c
18 files changed
+1695
-1369
docs/design/moe_kernel_features.md
+2
-2
@@ -88,8 +88,8 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k
88
| flashinfer | standard | nvfp4,</br>fp8 | T | <sup>5</sup> | N | Y | [`FlashInferExperts`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts] |
89
| gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.OAITritonExperts] |
90
| marlin | standard,</br>batched | <sup>3</sup> / N/A | <sup>3</sup> / N/A | silu,</br>swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe],</br>[`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],</br>[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] |
91
-| trtllm | standard | mxfp4,</br>nvfp4 | G(16),G(32) | <sup>5</sup> | N | Y | [`TrtLlmGenExperts`][vllm.model_executor.layers.fused_moe.trtllm_moe.TrtLlmGenExperts] |
92
-| rocm aiter moe | standard | fp8 | G(128),A,T | silu, gelu | Y | N | [`rocm_aiter_fused_experts`][vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe.rocm_aiter_fused_experts] |
91
+| trtllm | standard | mxfp4,</br>nvfp4 | G(16),G(32) | <sup>5</sup> | N | Y | [`TrtLlmMxfp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsMonolithic],</br>[`TrtLlmMxfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsModular],</br>[`TrtLlmNvFp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsMonolithic],</br>[`TrtLlmNvfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsModular] |
92
+| rocm aiter moe | standard | mxfp4,</br>fp8 | G(32),G(128),A,T | silu, gelu,</br>swigluoai | Y | N | `rocm_aiter_fused_experts`,</br>`AiterExperts` |
93
| cpu_fused_moe | standard | N/A | N/A | silu | N | N | [`CPUFusedMOE`][vllm.model_executor.layers.fused_moe.cpu_fused_moe.CPUFusedMOE] |
94
| naive batched<sup>4</sup> | batched | int8,</br>fp8 | G,A,T | silu, gelu | <sup>6</sup> | Y | [`NaiveBatchedExperts`][vllm.model_executor.layers.fused_moe.fused_batched_moe.NaiveBatchedExperts] |
95
tests/compile/fusions_e2e/conftest.py
+4
-1
@@ -84,7 +84,10 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
84
85
# TODO: remove this after finishing migration from envs to model kwargs
86
if model_name == "openai/gpt-oss-20b":
87
- monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8", "1")
87
+ from .common import is_blackwell
88
+
89
+ if is_blackwell():
90
+ monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8", "1")
91
92
# Disable, compile cache to make sure custom passes run.
93
# Otherwise, we can't verify fusion happened through the logs.
tests/kernels/moe/test_gpt_oss_triton_kernels.py
+8
@@ -6,6 +6,7 @@ import pytest
6
import torch
7
import torch.nn.functional as F
8
9
+from vllm.platforms import current_platform
10
from vllm.utils.import_utils import has_triton_kernels
11
12
if not has_triton_kernels():
@@ -14,6 +15,7 @@ if not has_triton_kernels():
15
allow_module_level=True,
16
)
17
18
+import triton_kernels.matmul_ogs_details.opt_flags as opt_flags
19
import triton_kernels.swiglu
20
from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
21
from triton_kernels.numerics import InFlexData
@@ -303,6 +305,12 @@ def test_equiv(num_token, a_dtype, w_dtype, tp, workspace_init):
305
pc2,
306
) = init_compute_data(M, K, N, E, a_dtype, w_dtype, num_warps=8)
307
308
+ if current_platform.is_device_capability_family(100):
309
+ constraints = {
310
+ "is_persistent": True,
311
+ }
312
+ opt_flags.update_opt_flags_constraints(constraints)
313
+
314
if a_dtype == "bf16" and w_dtype == "mx4":
315
quant_config = mxfp4_w4a16_moe_quant_config(
316
w1_scale=pc1,
tests/kernels/moe/test_ocp_mx_moe.py
+1
-1
@@ -82,7 +82,7 @@ def test_mxfp4_loading_and_execution_moe(vllm_runner, model_case: ModelCase):
82
model_case.model_id,
83
tensor_parallel_size=model_case.tp,
84
load_format="dummy",
85
- cudagraph_capture_sizes=[16],
85
+ compilation_config={"cudagraph_capture_sizes": [16]},
86
) as llm:
87
# Disabled as check_model is broken: https://github.com/vllm-project/vllm/pull/18465#issuecomment-3329880562
88
# def check_model(model):
tests/kernels/quantization/test_mxfp4_triton_ep.py
-83
@@ -17,89 +17,6 @@ from unittest.mock import MagicMock, patch
17
import pytest
18
import torch
19
20
-from vllm.model_executor.layers.quantization.mxfp4 import (
21
- Mxfp4Backend,
22
- Mxfp4MoEMethod,
23
-)
24
-
25
-
26
-def _make_mock_moe_config(ep_size: int = 1) -> MagicMock:
27
- """Create a mock FusedMoEConfig with the given EP size."""
28
- parallel_config = MagicMock()
29
- parallel_config.ep_size = ep_size
30
-
31
- moe_config = MagicMock()
32
- moe_config.ep_size = ep_size
33
- moe_config.is_lora_enabled = False
34
- moe_config.moe_parallel_config = parallel_config
35
- return moe_config
36
-
37
-
38
-class TestMxfp4TritonIsMonolithic:
39
- """Verify that is_monolithic is always True for the TRITON backend,
40
- regardless of EP size, since triton_kernel_moe_forward now handles
41
- expert_map remapping internally."""
42
-
43
- @pytest.mark.parametrize(
44
- "backend,ep_size,expected_monolithic",
45
- [
46
- # TRITON is always monolithic (handles EP via expert_map remapping)
47
- (Mxfp4Backend.TRITON, 1, True),
48
- (Mxfp4Backend.TRITON, 2, True),
49
- (Mxfp4Backend.TRITON, 4, True),
50
- # SM100 backends are always monolithic
51
- (Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM, 1, True),
52
- (Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM, 2, True),
53
- (Mxfp4Backend.SM100_FI_MXFP4_BF16, 1, True),
54
- (Mxfp4Backend.SM100_FI_MXFP4_BF16, 2, True),
55
- # MARLIN is never monolithic
56
- (Mxfp4Backend.MARLIN, 1, False),
57
- (Mxfp4Backend.MARLIN, 2, False),
58
- ],
59
- ids=[
60
- "triton-no-ep",
61
- "triton-ep2",
62
- "triton-ep4",
63
- "sm100-trtllm-no-ep",
64
- "sm100-trtllm-ep2",
65
- "sm100-bf16-no-ep",
66
- "sm100-bf16-ep2",
67
- "marlin-no-ep",
68
- "marlin-ep2",
69
- ],
70
- )
71
- @patch(
72
- "vllm.model_executor.layers.quantization.mxfp4.get_mxfp4_backend",
73
- )
74
- @patch(
75
- "vllm.model_executor.layers.quantization.mxfp4.get_current_vllm_config",
76
- )
77
- def test_is_monolithic(
78
- self,
79
- mock_get_config,
80
- mock_get_backend,
81
- backend,
82
- ep_size,
83
- expected_monolithic,
84
- ):
85
- """is_monolithic should be True for TRITON regardless of EP size."""
86
- mock_get_backend.return_value = backend
87
-
88
- mock_compilation_config = MagicMock()
89
- mock_compilation_config.max_cudagraph_capture_size = 1024
90
- mock_vllm_config = MagicMock()
91
- mock_vllm_config.compilation_config = mock_compilation_config
92
- mock_get_config.return_value = mock_vllm_config
93
-
94
- moe_config = _make_mock_moe_config(ep_size=ep_size)
95
- method = Mxfp4MoEMethod(moe_config)
96
-
97
- assert method.is_monolithic == expected_monolithic, (
98
- f"Expected is_monolithic={expected_monolithic} for "
99
- f"backend={backend.name}, ep_size={ep_size}, "
100
- f"but got {method.is_monolithic}."
101
- )
102
-
20
21
class TestTritonMoeForwardExpertMap:
22
"""Test that triton_kernel_moe_forward applies expert_map remapping
vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py
new
+352
@@ -0,0 +1,352 @@
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.fused_moe.topk_weight_and_reduce import (
15
+ TopKWeightAndReduceNoOP,
16
+)
17
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
18
+ QuantKey,
19
+ kMxfp4Static,
20
+ kMxfp8Dynamic,
21
+)
22
+from vllm.platforms import current_platform
23
+from vllm.utils.flashinfer import has_flashinfer
24
+
25
+
26
+class TrtLlmMxfp4ExpertsBase:
27
+ """
28
+ MXFP4 TRTLLM-Gen MoE kernels. Shared base for modular and monolithic.
29
+ """
30
+
31
+ def __init__(
32
+ self,
33
+ moe_config: FusedMoEConfig,
34
+ quant_config: FusedMoEQuantConfig,
35
+ ):
36
+ # NOTE: FusedMoEExperts.__init__ is called by the concrete subclass
37
+ # (Monolithic/Modular) via MRO, not here, to avoid mypy issues with
38
+ # multiple inheritance. This matches the NvFP4 expert pattern.
39
+ self.moe_config = moe_config
40
+ self.quant_config = quant_config
41
+
42
+ self.routing_method_type = moe_config.routing_method
43
+ self.topk = moe_config.experts_per_token
44
+ self.intermediate_size_per_partition = (
45
+ moe_config.intermediate_size_per_partition
46
+ )
47
+ self.hidden_dim = moe_config.hidden_dim
48
+ self.local_num_experts = moe_config.num_local_experts
49
+ self.ep_rank = moe_config.moe_parallel_config.ep_rank
50
+
51
+ # MXFP4-specific TRTLLM parameters
52
+ device = torch.accelerator.current_device_index()
53
+ self.gemm1_alpha = torch.tensor(
54
+ [1.702] * self.local_num_experts,
55
+ dtype=torch.float32,
56
+ device=device,
57
+ )
58
+ self.gemm1_beta = torch.tensor(
59
+ [1.0] * self.local_num_experts,
60
+ dtype=torch.float32,
61
+ device=device,
62
+ )
63
+ self.gemm1_clamp_limit = torch.tensor(
64
+ [7.0] * self.local_num_experts,
65
+ dtype=torch.float32,
66
+ device=device,
67
+ )
68
+
69
+ from vllm.config import get_current_vllm_config
70
+
71
+ self.max_capture_size = (
72
+ get_current_vllm_config().compilation_config.max_cudagraph_capture_size
73
+ )
74
+
75
+ # P1-5 fix: use public quant_dtype property instead of private _a1
76
+ self.use_mxfp8_input = quant_config.quant_dtype == "mxfp8"
77
+
78
+ @staticmethod
79
+ def _supports_current_device() -> bool:
80
+ p = current_platform
81
+ return p.is_cuda() and p.is_device_capability_family(100) and has_flashinfer()
82
+
83
+ @staticmethod
84
+ def _supports_no_act_and_mul() -> bool:
85
+ return False
86
+
87
+ @staticmethod
88
+ def _supports_quant_scheme(
89
+ weight_key: QuantKey | None,
90
+ activation_key: QuantKey | None,
91
+ ) -> bool:
92
+ SUPPORTED_W_A = [
93
+ (kMxfp4Static, None),
94
+ (kMxfp4Static, kMxfp8Dynamic),
95
+ ]
96
+ return (weight_key, activation_key) in SUPPORTED_W_A
97
+
98
+ @staticmethod
99
+ def _supports_activation(activation: MoEActivation) -> bool:
100
+ return activation == MoEActivation.SWIGLUOAI
101
+
102
+ @staticmethod
103
+ def activation_format() -> mk.FusedMoEActivationFormat:
104
+ return mk.FusedMoEActivationFormat.Standard
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
+ # Expert handles MXFP8 quantization internally if needed
115
+ return True
116
+
117
+
118
+class TrtLlmMxfp4ExpertsMonolithic(
119
+ TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsMonolithic
120
+):
121
+ """
122
+ Monolithic version of the MXFP4 TRTLLM kernel (router + experts).
123
+ Wraps flashinfer.trtllm_fp4_block_scale_moe().
124
+ """
125
+
126
+ @staticmethod
127
+ def _supports_parallel_config(
128
+ moe_parallel_config: FusedMoEParallelConfig,
129
+ ) -> bool:
130
+ return (
131
+ not moe_parallel_config.use_all2all_kernels
132
+ and not moe_parallel_config.enable_eplb
133
+ and moe_parallel_config.dp_size <= 1
134
+ )
135
+
136
+ @staticmethod
137
+ def _supports_routing_method(
138
+ routing_method: RoutingMethodType,
139
+ weight_key: QuantKey | None,
140
+ activation_key: QuantKey | None,
141
+ ) -> bool:
142
+ return routing_method in [
143
+ RoutingMethodType.Renormalize,
144
+ RoutingMethodType.RenormalizeNaive,
145
+ ]
146
+
147
+ @staticmethod
148
+ def _supports_router_logits_dtype(
149
+ router_logits_dtype: torch.dtype | None,
150
+ routing_method: RoutingMethodType,
151
+ ) -> bool:
152
+ # Kernel converts to bfloat16 internally
153
+ return True
154
+
155
+ def apply(
156
+ self,
157
+ hidden_states: torch.Tensor,
158
+ w1: torch.Tensor,
159
+ w2: torch.Tensor,
160
+ router_logits: torch.Tensor,
161
+ activation: MoEActivation,
162
+ global_num_experts: int,
163
+ expert_map: torch.Tensor | None,
164
+ a1q_scale: torch.Tensor | None,
165
+ apply_router_weight_on_input: bool,
166
+ # grouped topk + fused topk bias parameters
167
+ num_expert_group: int | None = None,
168
+ e_score_correction_bias: torch.Tensor | None = None,
169
+ routed_scaling_factor: float | None = None,
170
+ topk_group: int | None = None,
171
+ ) -> torch.Tensor:
172
+ from flashinfer import trtllm_fp4_block_scale_moe
173
+
174
+ # Handle input quantization
175
+ if self.use_mxfp8_input:
176
+ from flashinfer import mxfp8_quantize
177
+
178
+ x_quant, x_scale = mxfp8_quantize(
179
+ hidden_states,
180
+ is_sf_swizzled_layout=False,
181
+ alignment=256,
182
+ )
183
+ x_scale = x_scale.view(torch.float8_e4m3fn).reshape(
184
+ *hidden_states.shape[:-1], -1
185
+ )
186
+ else:
187
+ assert hidden_states.dtype == torch.bfloat16
188
+ x_quant = hidden_states
189
+ x_scale = None
190
+
191
+ output = torch.empty_like(hidden_states)
192
+
193
+ return trtllm_fp4_block_scale_moe(
194
+ routing_logits=router_logits.to(torch.bfloat16),
195
+ routing_bias=None,
196
+ hidden_states=x_quant,
197
+ hidden_states_scale=x_scale,
198
+ gemm1_weights=w1,
199
+ gemm1_weights_scale=self.w1_scale,
200
+ gemm1_bias=self.w1_bias,
201
+ gemm1_alpha=self.gemm1_alpha,
202
+ gemm1_beta=self.gemm1_beta,
203
+ gemm1_clamp_limit=self.gemm1_clamp_limit,
204
+ gemm2_weights=w2,
205
+ gemm2_weights_scale=self.w2_scale,
206
+ gemm2_bias=self.w2_bias,
207
+ output1_scale_scalar=None,
208
+ output1_scale_gate_scalar=None,
209
+ output2_scale_scalar=None,
210
+ num_experts=global_num_experts,
211
+ top_k=self.topk,
212
+ n_group=None,
213
+ topk_group=None,
214
+ intermediate_size=self.intermediate_size_per_partition,
215
+ local_expert_offset=self.ep_rank * self.local_num_experts,
216
+ local_num_experts=self.local_num_experts,
217
+ routed_scaling_factor=None,
218
+ routing_method_type=self.routing_method_type,
219
+ do_finalize=True,
220
+ tune_max_num_tokens=max(self.max_capture_size, 1),
221
+ output=output,
222
+ )[0]
223
+
224
+
225
+class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModular):
226
+ """
227
+ Modular version of the MXFP4 TRTLLM kernel (just the experts).
228
+ Wraps flashinfer.trtllm_fp4_block_scale_routed_moe().
229
+ Moved from trtllm_moe.py.
230
+ """
231
+
232
+ @property
233
+ def expects_unquantized_inputs(self) -> bool:
234
+ return True
235
+
236
+ @staticmethod
237
+ def _supports_parallel_config(
238
+ moe_parallel_config: FusedMoEParallelConfig,
239
+ ) -> bool:
240
+ return True
241
+
242
+ def supports_expert_map(self) -> bool:
243
+ return True
244
+
245
+ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
246
+ return TopKWeightAndReduceNoOP()
247
+
248
+ def workspace_shapes(
249
+ self,
250
+ M: int,
251
+ N: int,
252
+ K: int,
253
+ topk: int,
254
+ global_num_experts: int,
255
+ local_num_experts: int,
256
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
257
+ activation: MoEActivation,
258
+ ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
259
+ # The workspaces for this implementation are managed by flashinfer.
260
+ workspace1 = (0,)
261
+ workspace2 = (0,)
262
+ output = (M, K)
263
+ return (workspace1, workspace2, output)
264
+
265
+ def apply(
266
+ self,
267
+ output: torch.Tensor,
268
+ hidden_states: torch.Tensor,
269
+ w1: torch.Tensor,
270
+ w2: torch.Tensor,
271
+ topk_weights: torch.Tensor,
272
+ topk_ids: torch.Tensor,
273
+ activation: MoEActivation,
274
+ global_num_experts: int,
275
+ expert_map: torch.Tensor | None,
276
+ a1q_scale: torch.Tensor | None,
277
+ a2_scale: torch.Tensor | None,
278
+ workspace13: torch.Tensor,
279
+ workspace2: torch.Tensor,
280
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
281
+ apply_router_weight_on_input: bool,
282
+ ):
283
+ topk = topk_ids.size(-1)
284
+ local_num_experts = w1.size(0)
285
+ intermediate_size = w2.size(1)
286
+ local_expert_offset = self.moe_config.ep_rank * local_num_experts
287
+
288
+ # Handle input quantization
289
+ if self.use_mxfp8_input:
290
+ from flashinfer import mxfp8_quantize
291
+
292
+ x_quant, x_scale = mxfp8_quantize(
293
+ hidden_states,
294
+ is_sf_swizzled_layout=False,
295
+ alignment=256,
296
+ )
297
+ x_scale = x_scale.view(torch.float8_e4m3fn).reshape(
298
+ *hidden_states.shape[:-1], -1
299
+ )
300
+ else:
301
+ assert hidden_states.dtype == torch.bfloat16
302
+ x_quant = hidden_states
303
+ x_scale = None
304
+
305
+ packed_tensor = (topk_ids.to(torch.int32) << 16) | topk_weights.to(
306
+ torch.bfloat16
307
+ ).view(torch.int16)
308
+
309
+ assert self.w1_scale is not None
310
+ assert self.w2_scale is not None
311
+ kwargs = {
312
+ "topk_ids": packed_tensor,
313
+ "routing_bias": None,
314
+ "hidden_states": x_quant,
315
+ "hidden_states_scale": x_scale,
316
+ "gemm1_weights": w1,
317
+ "gemm1_weights_scale": self.w1_scale,
318
+ "gemm1_bias": self.w1_bias,
319
+ "gemm1_alpha": self.gemm1_alpha,
320
+ "gemm1_beta": self.gemm1_beta,
321
+ "gemm1_clamp_limit": self.gemm1_clamp_limit,
322
+ "gemm2_weights": w2,
323
+ "gemm2_weights_scale": self.w2_scale,
324
+ "gemm2_bias": self.w2_bias,
325
+ "output1_scale_scalar": None,
326
+ "output1_scale_gate_scalar": None,
327
+ "output2_scale_scalar": None,
328
+ "num_experts": global_num_experts,
329
+ "top_k": topk,
330
+ "n_group": None,
331
+ "topk_group": None,
332
+ "intermediate_size": intermediate_size,
333
+ "local_expert_offset": local_expert_offset,
334
+ "local_num_experts": local_num_experts,
335
+ "routed_scaling_factor": None,
336
+ "routing_method_type": self.routing_method_type,
337
+ "do_finalize": True,
338
+ "output": output,
339
+ "tune_max_num_tokens": max(self.max_capture_size, 1),
340
+ }
341
+
342
+ from flashinfer import trtllm_fp4_block_scale_routed_moe
343
+
344
+ from vllm.utils.flashinfer import autotune
345
+
346
+ with autotune(False):
347
+ # Enable autotune when,
348
+ # https://github.com/flashinfer-ai/flashinfer/issues/2023 is
349
+ # resolved.
350
+ trtllm_fp4_block_scale_routed_moe(**kwargs)
351
+
352
+ return output
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
+3
-1
@@ -40,6 +40,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
40
kFp8Static128BlockSym,
41
kFp8StaticChannelSym,
42
kFp8StaticTensorSym,
43
+ kMxfp4Static,
44
kNvfp4Static,
45
)
46
from vllm.platforms import current_platform
@@ -574,12 +575,13 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
575
weight_key: QuantKey | None,
576
activation_key: QuantKey | None,
577
) -> bool:
577
- # TODO(rob): add int4, mxfp4, int8 as integrations
578
+ # TODO(rob): add int4, int8 as integrations
579
# are migrated to use the oracle one-by-one.
580
SUPPORTED_W = [
581
kFp8Static128BlockSym,
582
kFp8StaticChannelSym,
583
kFp8StaticTensorSym,
584
+ kMxfp4Static,
585
kNvfp4Static,
586
]
587
return weight_key in SUPPORTED_W
vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py
+151
-20
@@ -11,8 +11,10 @@ from vllm.logger import init_logger
11
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
12
from vllm.model_executor.layers.fused_moe.config import (
13
FUSED_MOE_UNQUANTIZED_CONFIG,
14
+ FusedMoEConfig,
15
FusedMoEParallelConfig,
16
FusedMoEQuantConfig,
17
+ RoutingMethodType,
18
)
19
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
20
TopKWeightAndReduceNoOP,
@@ -20,6 +22,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
22
from vllm.model_executor.layers.fused_moe.utils import _resize_cache
23
from vllm.model_executor.layers.quantization.utils.quant_utils import (
24
QuantKey,
25
+ kMxfp4Static,
26
)
27
from vllm.platforms import current_platform
28
from vllm.triton_utils import tl, triton
@@ -537,43 +540,43 @@ def make_routing_data(
540
541
542
class BaseOAITritonExperts(mk.FusedMoEExpertsModular):
543
+ @property
544
+ def expects_unquantized_inputs(self) -> bool:
545
+ return True
546
+
547
@staticmethod
548
def _supports_current_device() -> bool:
542
- raise NotImplementedError(
543
- "OAITritonExperts is not yet used by an Oracle. "
544
- "This method should not be called."
545
- )
549
+ p = current_platform
550
+ if not p.is_cuda_alike():
551
+ return False
552
+ cap = p.get_device_capability()
553
+ if cap is None:
554
+ return False
555
+ # (9,0) <= cap < (11,0) covers CUDA SM90 (Hopper), SM100+ (Blackwell)
556
+ # and ROCm gfx942/gfx950 (which map to 9.4/9.5).
557
+ return (9, 0) <= (cap.major, cap.minor) < (11, 0)
558
559
@staticmethod
560
def _supports_no_act_and_mul() -> bool:
549
- raise NotImplementedError(
550
- "OAITritonExperts is not yet used by an Oracle. "
551
- "This method should not be called."
552
- )
561
+ return False
562
563
@staticmethod
564
def _supports_quant_scheme(
565
weight_key: QuantKey | None,
566
activation_key: QuantKey | None,
567
) -> bool:
559
- raise NotImplementedError(
560
- "OAITritonExperts is not yet used by an Oracle. "
561
- "This method should not be called."
562
- )
568
+ SUPPORTED_W_A = [
569
+ (kMxfp4Static, None),
570
+ ]
571
+ return (weight_key, activation_key) in SUPPORTED_W_A
572
573
@staticmethod
574
def _supports_activation(activation: MoEActivation) -> bool:
566
- raise NotImplementedError(
567
- "OAITritonExperts is not yet used by an Oracle. "
568
- "This method should not be called."
569
- )
575
+ raise NotImplementedError
576
577
@staticmethod
578
def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
573
- raise NotImplementedError(
574
- "OAITritonExperts is not yet used by an Oracle. "
575
- "This method should not be called."
576
- )
579
+ return True
580
581
def supports_expert_map(self) -> bool:
582
return True
@@ -630,6 +633,10 @@ class BaseOAITritonExperts(mk.FusedMoEExpertsModular):
633
class OAITritonExperts(BaseOAITritonExperts):
634
"""OAI Triton-based fused MoE expert implementation."""
635
636
+ @staticmethod
637
+ def _supports_activation(activation: MoEActivation) -> bool:
638
+ return activation == MoEActivation.SWIGLUOAI
639
+
640
@staticmethod
641
def activation_format() -> mk.FusedMoEActivationFormat:
642
return mk.FusedMoEActivationFormat.Standard
@@ -714,6 +721,15 @@ class UnfusedOAITritonExperts(BaseOAITritonExperts):
721
One use case for it is to inject LoRA modules on the activation and moe_sum.
722
"""
723
724
+ @staticmethod
725
+ def _supports_activation(activation: MoEActivation) -> bool:
726
+ return activation in [
727
+ MoEActivation.SILU,
728
+ MoEActivation.GELU,
729
+ MoEActivation.SWIGLUOAI,
730
+ MoEActivation.SWIGLUSTEP,
731
+ ]
732
+
733
@staticmethod
734
def activation_format() -> mk.FusedMoEActivationFormat:
735
return mk.FusedMoEActivationFormat.Standard
@@ -839,3 +855,118 @@ class UnfusedOAITritonExperts(BaseOAITritonExperts):
855
)
856
857
self.moe_sum(intermediate_cache3.view(-1, topk, K), output)
858
+
859
+
860
+class OAITritonMxfp4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
861
+ """Monolithic Triton MXFP4 expert. Wraps triton_kernel_moe_forward()."""
862
+
863
+ def __init__(
864
+ self,
865
+ moe_config: FusedMoEConfig,
866
+ quant_config: FusedMoEQuantConfig,
867
+ ):
868
+ super().__init__(moe_config, quant_config)
869
+ self.topk = moe_config.experts_per_token
870
+ self.renormalize = moe_config.routing_method in (
871
+ RoutingMethodType.Renormalize,
872
+ RoutingMethodType.RenormalizeNaive,
873
+ )
874
+
875
+ @staticmethod
876
+ def activation_format() -> mk.FusedMoEActivationFormat:
877
+ return mk.FusedMoEActivationFormat.Standard
878
+
879
+ @staticmethod
880
+ def _supports_current_device() -> bool:
881
+ p = current_platform
882
+ if not p.is_cuda_alike():
883
+ return False
884
+ cap = p.get_device_capability()
885
+ if cap is None:
886
+ return False
887
+ # (9,0) <= cap < (11,0) covers CUDA SM90 (Hopper), SM100+ (Blackwell)
888
+ # and ROCm gfx942/gfx950 (which map to 9.4/9.5).
889
+ return (9, 0) <= (cap.major, cap.minor) < (11, 0)
890
+
891
+ @staticmethod
892
+ def _supports_no_act_and_mul() -> bool:
893
+ return False
894
+
895
+ @staticmethod
896
+ def _supports_quant_scheme(
897
+ weight_key: QuantKey | None,
898
+ activation_key: QuantKey | None,
899
+ ) -> bool:
900
+ SUPPORTED_W_A = [
901
+ (kMxfp4Static, None),
902
+ ]
903
+ return (weight_key, activation_key) in SUPPORTED_W_A
904
+
905
+ @staticmethod
906
+ def _supports_activation(activation: MoEActivation) -> bool:
907
+ return activation == MoEActivation.SWIGLUOAI
908
+
909
+ @staticmethod
910
+ def _supports_parallel_config(
911
+ moe_parallel_config: FusedMoEParallelConfig,
912
+ ) -> bool:
913
+ return (
914
+ not moe_parallel_config.use_all2all_kernels
915
+ and not moe_parallel_config.enable_eplb
916
+ and moe_parallel_config.dp_size <= 1
917
+ )
918
+
919
+ @staticmethod
920
+ def _supports_routing_method(
921
+ routing_method: RoutingMethodType,
922
+ weight_key: QuantKey | None,
923
+ activation_key: QuantKey | None,
924
+ ) -> bool:
925
+ return routing_method in [
926
+ RoutingMethodType.Renormalize,
927
+ RoutingMethodType.RenormalizeNaive,
928
+ ]
929
+
930
+ @staticmethod
931
+ def _supports_router_logits_dtype(
932
+ router_logits_dtype: torch.dtype | None,
933
+ routing_method: RoutingMethodType,
934
+ ) -> bool:
935
+ return True
936
+
937
+ def supports_expert_map(self) -> bool:
938
+ return True
939
+
940
+ @property
941
+ def expects_unquantized_inputs(self) -> bool:
942
+ return True
943
+
944
+ def apply(
945
+ self,
946
+ hidden_states: torch.Tensor,
947
+ w1: torch.Tensor,
948
+ w2: torch.Tensor,
949
+ router_logits: torch.Tensor,
950
+ activation: MoEActivation,
951
+ global_num_experts: int,
952
+ expert_map: torch.Tensor | None,
953
+ a1q_scale: torch.Tensor | None,
954
+ apply_router_weight_on_input: bool,
955
+ # grouped topk + fused topk bias parameters
956
+ num_expert_group: int | None = None,
957
+ e_score_correction_bias: torch.Tensor | None = None,
958
+ routed_scaling_factor: float | None = None,
959
+ topk_group: int | None = None,
960
+ ) -> torch.Tensor:
961
+ return triton_kernel_moe_forward(
962
+ hidden_states=hidden_states,
963
+ w1=w1,
964
+ w2=w2,
965
+ gating_output=router_logits,
966
+ topk=self.topk,
967
+ renormalize=self.renormalize,
968
+ global_num_experts=global_num_experts,
969
+ expert_map=expert_map,
970
+ quant_config=self.quant_config,
971
+ apply_router_weight_on_input=apply_router_weight_on_input,
972
+ )
vllm/model_executor/layers/fused_moe/layer.py
-28
@@ -52,7 +52,6 @@ from vllm.model_executor.layers.quantization.base_config import (
52
QuantizationConfig,
53
)
54
from vllm.platforms import current_platform
55
-from vllm.utils.math_utils import round_up
55
56
logger = init_logger(__name__)
57
@@ -218,7 +217,6 @@ def maybe_roundup_hidden_size(
217
moe_parallel_config: FusedMoEParallelConfig,
218
is_lora_enabled: bool,
219
model_type: str | None,
221
- is_mxfp4_quant: bool,
220
) -> int:
221
"""
222
Given layer hidden size and MoE configurations, round up hidden_size
@@ -232,7 +230,6 @@ def maybe_roundup_hidden_size(
230
is used in the case of mxfp4 quantization in selecting the
231
MxFP4Backend.
232
model_type: for checking if gpt-oss
235
- is_mxfp4_quant: whether the layer is quantized with mxfp4
233
234
Return:
235
Rounded up hidden_size if rounding up is required based on the configs.
@@ -246,28 +243,6 @@ def maybe_roundup_hidden_size(
243
hidden_size, act_dtype, moe_parallel_config
244
)
245
249
- # we are padding globally so EP buffer allocation works
250
- if model_type == "gpt_oss" and is_mxfp4_quant:
251
- from vllm.model_executor.layers.quantization.mxfp4 import (
252
- Mxfp4Backend,
253
- get_mxfp4_backend,
254
- )
255
-
256
- current_mxfp4_backend = get_mxfp4_backend(is_lora_enabled)
257
-
258
- if (
259
- current_mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16
260
- or current_mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS
261
- ):
262
- hidden_size = round_up(hidden_size, 128)
263
- elif (
264
- current_platform.is_rocm()
265
- or current_mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
266
- or current_mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16
267
- or current_mxfp4_backend == Mxfp4Backend.MARLIN
268
- ):
269
- hidden_size = round_up(hidden_size, 256)
270
-
246
return hidden_size
247
248
@@ -540,9 +515,6 @@ class FusedMoE(CustomOp):
515
moe_parallel_config=self.moe_parallel_config,
516
is_lora_enabled=vllm_config.lora_config is not None,
517
model_type=self.model_type,
543
- is_mxfp4_quant=(
544
- quant_config is not None and quant_config.is_mxfp4_quant(prefix, self)
545
- ),
518
)
519
self.hidden_size = hidden_size
520
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
new
+847
@@ -0,0 +1,847 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from enum import Enum
4
+from typing import Union
5
+
6
+import torch
7
+
8
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9
+from vllm import envs
10
+from vllm.logger import init_logger
11
+from vllm.model_executor.layers.fused_moe import (
12
+ FusedMoEConfig,
13
+)
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
+ FusedMoEQuantConfig,
19
+ mxfp4_mxfp8_moe_quant_config,
20
+ mxfp4_w4a16_moe_quant_config,
21
+ ocp_mx_moe_quant_config,
22
+)
23
+from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
24
+ _swizzle_mxfp4,
25
+ get_padding_alignment,
26
+)
27
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
28
+ QuantKey,
29
+ kMxfp4Static,
30
+ kMxfp8Dynamic,
31
+)
32
+from vllm.platforms import current_platform
33
+from vllm.utils.import_utils import has_triton_kernels
34
+from vllm.utils.math_utils import round_up
35
+
36
+logger = init_logger(__name__)
37
+
38
+if has_triton_kernels():
39
+ try:
40
+ from triton_kernels.matmul_ogs import PrecisionConfig
41
+ except (ImportError, AttributeError) as e:
42
+ logger.error(
43
+ "Failed to import Triton kernels. Please make sure your triton "
44
+ "version is compatible. Error: %s",
45
+ e,
46
+ )
47
+
48
+
49
+class Mxfp4MoeBackend(Enum):
50
+ NONE = "None"
51
+ # FlashInfer TRTLLM backends
52
+ FLASHINFER_TRTLLM_MXFP4_MXFP8 = "FLASHINFER_TRTLLM_MXFP4_MXFP8"
53
+ FLASHINFER_TRTLLM_MXFP4_BF16 = "FLASHINFER_TRTLLM_MXFP4_BF16"
54
+ # FlashInfer CUTLASS backends
55
+ FLASHINFER_CUTLASS_MXFP4_MXFP8 = "FLASHINFER_CUTLASS_MXFP4_MXFP8"
56
+ FLASHINFER_CUTLASS_MXFP4_BF16 = "FLASHINFER_CUTLASS_MXFP4_BF16"
57
+ # Marlin
58
+ BATCHED_MARLIN = "BATCHED_MARLIN"
59
+ MARLIN = "MARLIN"
60
+ # ROCm AITER (CK)
61
+ CK = "CK"
62
+ # Triton
63
+ TRITON = "TRITON"
64
+ TRITON_UNFUSED = "TRITON_UNFUSED"
65
+ # XPU
66
+ XPU = "XPU"
67
+
68
+
69
+# Backends that share the same TRTLLM weight format
70
+TRTLLM_BACKENDS = (
71
+ Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
72
+ Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
73
+)
74
+
75
+TRITON_BACKENDS = (
76
+ Mxfp4MoeBackend.TRITON,
77
+ Mxfp4MoeBackend.TRITON_UNFUSED,
78
+)
79
+
80
+
81
+def backend_to_kernel_cls(
82
+ backend: Mxfp4MoeBackend,
83
+) -> list[type[mk.FusedMoEExperts]]:
84
+ if backend in (
85
+ Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
86
+ Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
87
+ ):
88
+ from vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe import (
89
+ TrtLlmMxfp4ExpertsModular,
90
+ TrtLlmMxfp4ExpertsMonolithic,
91
+ )
92
+
93
+ # NOTE: prefer Monolithic > Modular, so return Monolithic first.
94
+ return [TrtLlmMxfp4ExpertsMonolithic, TrtLlmMxfp4ExpertsModular]
95
+
96
+ elif backend in (
97
+ Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
98
+ Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
99
+ ):
100
+ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
101
+ FlashInferExperts,
102
+ )
103
+
104
+ return [FlashInferExperts]
105
+
106
+ elif backend == Mxfp4MoeBackend.TRITON:
107
+ from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
108
+ OAITritonExperts,
109
+ OAITritonMxfp4ExpertsMonolithic,
110
+ )
111
+
112
+ # NOTE: prefer Monolithic > Modular, so return Monolithic first.
113
+ return [OAITritonMxfp4ExpertsMonolithic, OAITritonExperts]
114
+
115
+ elif backend == Mxfp4MoeBackend.TRITON_UNFUSED:
116
+ from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
117
+ UnfusedOAITritonExperts,
118
+ )
119
+
120
+ return [UnfusedOAITritonExperts]
121
+
122
+ elif backend == Mxfp4MoeBackend.MARLIN:
123
+ from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
124
+ MarlinExperts,
125
+ )
126
+
127
+ return [MarlinExperts]
128
+
129
+ elif backend == Mxfp4MoeBackend.BATCHED_MARLIN:
130
+ from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
131
+ BatchedMarlinExperts,
132
+ )
133
+
134
+ return [BatchedMarlinExperts]
135
+
136
+ elif backend == Mxfp4MoeBackend.CK:
137
+ from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
138
+ AiterExperts,
139
+ )
140
+
141
+ return [AiterExperts]
142
+
143
+ elif backend == Mxfp4MoeBackend.XPU:
144
+ raise NotImplementedError("XPU backend uses XpuMxfp4MoEMethod directly.")
145
+ else:
146
+ raise ValueError(f"Unknown MXFP4 MoE backend: {backend.value}")
147
+
148
+
149
+def map_mxfp4_backend(runner_backend: str) -> Mxfp4MoeBackend:
150
+ """Map user's moe_backend string to Mxfp4MoeBackend."""
151
+ mapping: dict[str, Mxfp4MoeBackend] = {
152
+ "flashinfer_trtllm": Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
153
+ "flashinfer_trtllm_afp8": Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
154
+ "flashinfer_cutlass": Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
155
+ "flashinfer_cutlass_afp8": Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
156
+ "triton": Mxfp4MoeBackend.TRITON,
157
+ "marlin": Mxfp4MoeBackend.MARLIN,
158
+ "ck": Mxfp4MoeBackend.CK,
159
+ }
160
+ if backend := mapping.get(runner_backend):
161
+ return backend
162
+ raise ValueError(
163
+ f"moe_backend='{runner_backend}' is not supported for MXFP4 MoE. "
164
+ f"Expected one of {list(mapping.keys())}."
165
+ )
166
+
167
+
168
+def _get_priority_backends() -> list[Mxfp4MoeBackend]:
169
+ """
170
+ Get available backends in priority order based on platform and config.
171
+ Only includes BF16 backends. MXFP8 backends are selected via env vars.
172
+ """
173
+ _AVAILABLE_BACKENDS = [
174
+ Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
175
+ Mxfp4MoeBackend.CK,
176
+ Mxfp4MoeBackend.TRITON,
177
+ Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
178
+ Mxfp4MoeBackend.TRITON_UNFUSED,
179
+ Mxfp4MoeBackend.MARLIN,
180
+ Mxfp4MoeBackend.BATCHED_MARLIN,
181
+ ]
182
+ return _AVAILABLE_BACKENDS
183
+
184
+
185
+def _backend_activation_key(backend: Mxfp4MoeBackend) -> QuantKey | None:
186
+ """Map backend to its activation key (MXFP8 or None for BF16)."""
187
+ if backend in (
188
+ Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
189
+ Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
190
+ ):
191
+ return kMxfp8Dynamic
192
+ return None
193
+
194
+
195
+def select_mxfp4_moe_backend(
196
+ config: FusedMoEConfig,
197
+) -> tuple[Mxfp4MoeBackend, type[mk.FusedMoEExperts] | None]:
198
+ """
199
+ Select the primary MXFP4 MoE backend.
200
+ Note: Shape-specific fallbacks may still occur at runtime.
201
+ """
202
+ triton_kernels_supported = has_triton_kernels() and (
203
+ 9,
204
+ 0,
205
+ ) <= current_platform.get_device_capability() < (11, 0)
206
+
207
+ # LoRA: separate experts backend path
208
+ if config.is_lora_enabled:
209
+ if not current_platform.is_cuda():
210
+ raise NotImplementedError("Mxfp4 LoRA only supported on CUDA Platform.")
211
+ if envs.VLLM_MXFP4_USE_MARLIN is False and triton_kernels_supported:
212
+ logger.info_once("Using Triton backend for mxfp4 lora")
213
+ return Mxfp4MoeBackend.TRITON_UNFUSED, backend_to_kernel_cls(
214
+ Mxfp4MoeBackend.TRITON_UNFUSED
215
+ )[0]
216
+ logger.info_once("Using Marlin backend for mxfp4 lora")
217
+ return Mxfp4MoeBackend.MARLIN, backend_to_kernel_cls(Mxfp4MoeBackend.MARLIN)[0]
218
+
219
+ activation_format = (
220
+ mk.FusedMoEActivationFormat.BatchedExperts
221
+ if config.moe_parallel_config.use_batched_activation_format
222
+ else mk.FusedMoEActivationFormat.Standard
223
+ )
224
+
225
+ def _make_log_backend(backend: Mxfp4MoeBackend):
226
+ return f"Using '{backend.value}' Mxfp4 MoE backend."
227
+
228
+ def _make_log_unsupported(backend: Mxfp4MoeBackend, reason: str | None) -> str:
229
+ if reason:
230
+ return (
231
+ f"Mxfp4 MoE backend '{backend.value}' does not support the "
232
+ f"deployment configuration since {reason}."
233
+ )
234
+ return (
235
+ f"Mxfp4 MoE backend '{backend.value}' does not support the "
236
+ "deployment configuration."
237
+ )
238
+
239
+ def _return_or_raise(
240
+ backend: Mxfp4MoeBackend,
241
+ config: FusedMoEConfig,
242
+ weight_key: QuantKey | None,
243
+ activation_key: QuantKey | None,
244
+ activation_format: mk.FusedMoEActivationFormat,
245
+ ) -> tuple[Mxfp4MoeBackend, type[mk.FusedMoEExperts]]:
246
+ reason: str | None = None
247
+ for k_cls in backend_to_kernel_cls(backend):
248
+ supported, reason = k_cls.is_supported_config(
249
+ k_cls, config, weight_key, activation_key, activation_format
250
+ )
251
+ if supported:
252
+ logger.info_once(_make_log_backend(backend), scope="local")
253
+ return backend, k_cls
254
+ raise ValueError(_make_log_unsupported(backend, reason))
255
+
256
+ runner_backend = config.moe_backend
257
+ if runner_backend != "auto":
258
+ requested_backend = map_mxfp4_backend(runner_backend)
259
+ if (
260
+ activation_format == mk.FusedMoEActivationFormat.BatchedExperts
261
+ and requested_backend == Mxfp4MoeBackend.MARLIN
262
+ ):
263
+ requested_backend = Mxfp4MoeBackend.BATCHED_MARLIN
264
+ return _return_or_raise(
265
+ requested_backend,
266
+ config,
267
+ kMxfp4Static,
268
+ _backend_activation_key(requested_backend),
269
+ activation_format,
270
+ )
271
+
272
+ # Select kernels in order of backend.
273
+ AVAILABLE_BACKENDS = _get_priority_backends()
274
+
275
+ # Handle explicit FlashInfer MXFP4 BF16 configuration.
276
+ if envs.is_set("VLLM_USE_FLASHINFER_MOE_MXFP4_BF16"):
277
+ if not envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16:
278
+ AVAILABLE_BACKENDS.remove(Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16)
279
+ AVAILABLE_BACKENDS.remove(Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16)
280
+ else:
281
+ if current_platform.is_device_capability(90):
282
+ return _return_or_raise(
283
+ Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
284
+ config,
285
+ kMxfp4Static,
286
+ None,
287
+ activation_format,
288
+ )
289
+ if current_platform.is_device_capability_family(100):
290
+ return _return_or_raise(
291
+ Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
292
+ config,
293
+ kMxfp4Static,
294
+ None,
295
+ activation_format,
296
+ )
297
+ raise ValueError(
298
+ "VLLM_USE_FLASHINFER_MOE_MXFP4_BF16=1 is set but the "
299
+ "current device capability is not supported. "
300
+ "Only SM90 (CUTLASS) and SM100+ (TRTLLM) are supported."
301
+ )
302
+
303
+ # Handle explicit FlashInfer MXFP4 MXFP8 TRTLLM configuration.
304
+ if (
305
+ envs.is_set("VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8")
306
+ and envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
307
+ ):
308
+ return _return_or_raise(
309
+ Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
310
+ config,
311
+ kMxfp4Static,
312
+ kMxfp8Dynamic,
313
+ activation_format,
314
+ )
315
+
316
+ # Handle explicit FlashInfer MXFP4 MXFP8 CUTLASS configuration.
317
+ if (
318
+ envs.is_set("VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS")
319
+ and envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS
320
+ ):
321
+ return _return_or_raise(
322
+ Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
323
+ config,
324
+ kMxfp4Static,
325
+ kMxfp8Dynamic,
326
+ activation_format,
327
+ )
328
+
329
+ # Handle explicit Marlin MXFP4 configuration.
330
+ if envs.is_set("VLLM_MXFP4_USE_MARLIN") and envs.VLLM_MXFP4_USE_MARLIN:
331
+ return _return_or_raise(
332
+ Mxfp4MoeBackend.MARLIN,
333
+ config,
334
+ kMxfp4Static,
335
+ None,
336
+ activation_format,
337
+ )
338
+
339
+ for backend in AVAILABLE_BACKENDS:
340
+ activation_key = _backend_activation_key(backend)
341
+ for k_cls in backend_to_kernel_cls(backend):
342
+ supported, reason = k_cls.is_supported_config(
343
+ k_cls, config, kMxfp4Static, activation_key, activation_format
344
+ )
345
+ if supported:
346
+ logger.info_once(_make_log_backend(backend), scope="local")
347
+ return backend, k_cls
348
+ else:
349
+ logger.debug_once(_make_log_unsupported(backend, reason), scope="local")
350
+
351
+ if current_platform.is_xpu():
352
+ backend = Mxfp4MoeBackend.XPU
353
+ logger.info_once(_make_log_backend(backend))
354
+ return backend, None
355
+
356
+ if current_platform.is_cuda() or current_platform.is_rocm():
357
+ raise NotImplementedError(
358
+ "No MXFP4 MoE backend supports the deployment configuration."
359
+ )
360
+
361
+ return Mxfp4MoeBackend.NONE, None
362
+
363
+
364
+def mxfp4_round_up_hidden_size_and_intermediate_size(
365
+ backend: Mxfp4MoeBackend, hidden_size: int, intermediate_size: int
366
+) -> tuple[int, int]:
367
+ """Round up hidden_size and intermediate_size based on backend requirements."""
368
+ if backend in (Mxfp4MoeBackend.MARLIN, Mxfp4MoeBackend.BATCHED_MARLIN):
369
+ intermediate_size = round_up(intermediate_size, 128)
370
+ if current_platform.is_xpu():
371
+ hidden_size = round_up(hidden_size, 128)
372
+ else:
373
+ hidden_size = round_up(hidden_size, 256)
374
+ elif backend in TRTLLM_BACKENDS:
375
+ intermediate_size = round_up(intermediate_size, 256)
376
+ hidden_size = round_up(hidden_size, 256)
377
+ elif backend in (
378
+ Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
379
+ Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
380
+ ):
381
+ intermediate_size = round_up(intermediate_size, 128)
382
+ hidden_size = round_up(hidden_size, 128)
383
+ elif current_platform.is_rocm():
384
+ pad_align = get_padding_alignment()
385
+ intermediate_size = round_up(intermediate_size, pad_align)
386
+ hidden_size = round_up(hidden_size, pad_align)
387
+ else:
388
+ intermediate_size = round_up(intermediate_size, 64)
389
+ return hidden_size, intermediate_size
390
+
391
+
392
+def convert_to_mxfp4_moe_kernel_format(
393
+ mxfp4_backend: Mxfp4MoeBackend,
394
+ layer: torch.nn.Module,
395
+ w13_weight: torch.Tensor,
396
+ w2_weight: torch.Tensor,
397
+ w13_weight_scale: torch.Tensor,
398
+ w2_weight_scale: torch.Tensor,
399
+ w13_bias: torch.Tensor | None = None,
400
+ w2_bias: torch.Tensor | None = None,
401
+ _cache_permute_indices: dict[torch.Size, torch.Tensor] | None = None,
402
+) -> tuple[
403
+ torch.Tensor,
404
+ torch.Tensor,
405
+ Union[torch.Tensor, "PrecisionConfig"],
406
+ Union[torch.Tensor, "PrecisionConfig"],
407
+ torch.Tensor | None,
408
+ torch.Tensor | None,
409
+]:
410
+ """Convert loaded weights into backend-specific kernel format."""
411
+
412
+ num_experts = w13_weight.shape[0]
413
+ intermediate_size = w13_weight.shape[1] // 2
414
+ hidden_size = w13_weight.shape[2] * 2
415
+
416
+ sf_block_size = 32 # mxfp4 block size
417
+
418
+ if mxfp4_backend in (Mxfp4MoeBackend.MARLIN, Mxfp4MoeBackend.BATCHED_MARLIN):
419
+ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
420
+ prepare_moe_mxfp4_layer_for_marlin,
421
+ )
422
+
423
+ return prepare_moe_mxfp4_layer_for_marlin(
424
+ layer,
425
+ w13_weight,
426
+ w2_weight,
427
+ w13_weight_scale,
428
+ w2_weight_scale,
429
+ w13_bias,
430
+ w2_bias,
431
+ )
432
+
433
+ elif mxfp4_backend in TRTLLM_BACKENDS:
434
+ assert _cache_permute_indices is not None
435
+ from flashinfer.fp4_quantization import nvfp4_block_scale_interleave
436
+ from flashinfer.fused_moe.core import get_w2_permute_indices_with_cache
437
+
438
+ # gemm1_alpha/beta/clamp_limit are created by the expert class
439
+ # (TrtLlmMxfp4ExpertsBase), not on the layer.
440
+
441
+ w13_weight = w13_weight.data
442
+ w2_weight = w2_weight.data
443
+ w13_weight_scale = w13_weight_scale.data
444
+ w2_weight_scale = w2_weight_scale.data
445
+ assert w13_bias is not None and w2_bias is not None
446
+ w13_bias = w13_bias.data.to(torch.float32)
447
+ w2_bias = w2_bias.data.to(torch.float32)
448
+
449
+ # Swap w1 and w3 as the definition of swiglu is different in trtllm-gen
450
+ def swap_every_two_rows(x, axis=-1):
451
+ shape = x.shape
452
+ if axis < 0:
453
+ axis = len(shape) + axis
454
+ new_shape = list(shape)
455
+ new_shape[axis] = shape[axis] // 2
456
+ new_shape.insert(axis + 1, 2)
457
+ x = x.reshape(*new_shape)
458
+ x = x.flip(axis + 1)
459
+ new_shape = list(shape)
460
+ return x.reshape(*new_shape)
461
+
462
+ w13_weight_scale = swap_every_two_rows(w13_weight_scale, -2)
463
+ w13_weight = swap_every_two_rows(w13_weight, -2)
464
+ w13_bias = swap_every_two_rows(w13_bias, -1)
465
+
466
+ # Shuffle weights and scaling factors for transposed mma output
467
+ gemm1_weights_shuffled = []
468
+ gemm1_scales_shuffled = []
469
+ gemm2_weights_shuffled = []
470
+ gemm2_scales_shuffled = []
471
+ gemm1_bias_shuffled = []
472
+ gemm2_bias_shuffled = []
473
+ epilogue_tile_m = 128
474
+ for i in range(num_experts):
475
+ # w13 weight
476
+ permute_indices = get_w2_permute_indices_with_cache(
477
+ _cache_permute_indices,
478
+ w13_weight[i].view(torch.uint8),
479
+ epilogue_tile_m,
480
+ )
481
+ gemm1_weights_shuffled.append(
482
+ w13_weight[i]
483
+ .view(torch.uint8)[permute_indices.to(w13_weight.device)]
484
+ .contiguous()
485
+ )
486
+ # w13 scale
487
+ permute_sf_indices = get_w2_permute_indices_with_cache(
488
+ _cache_permute_indices,
489
+ w13_weight_scale[i].view(torch.uint8),
490
+ epilogue_tile_m,
491
+ num_elts_per_sf=16,
492
+ )
493
+ gemm1_scales_shuffled.append(
494
+ nvfp4_block_scale_interleave(
495
+ w13_weight_scale[i]
496
+ .view(torch.uint8)[permute_sf_indices.to(w13_weight_scale.device)]
497
+ .contiguous()
498
+ )
499
+ )
500
+ # w13 bias
501
+ permute_bias_indices = get_w2_permute_indices_with_cache(
502
+ _cache_permute_indices,
503
+ w13_bias[i].clone().reshape(-1, 1),
504
+ epilogue_tile_m,
505
+ )
506
+ gemm1_bias_shuffled.append(
507
+ w13_bias[i]
508
+ .clone()
509
+ .reshape(-1, 1)[permute_bias_indices.to(w13_bias.device)]
510
+ .contiguous()
511
+ )
512
+ # w2 weight
513
+ permute_indices = get_w2_permute_indices_with_cache(
514
+ _cache_permute_indices,
515
+ w2_weight[i].view(torch.uint8),
516
+ epilogue_tile_m,
517
+ )
518
+ gemm2_weights_shuffled.append(
519
+ w2_weight[i]
520
+ .view(torch.uint8)[permute_indices.to(w2_weight.device)]
521
+ .contiguous()
522
+ )
523
+ # w2 scale
524
+ permute_sf_indices = get_w2_permute_indices_with_cache(
525
+ _cache_permute_indices,
526
+ w2_weight_scale[i].view(torch.uint8),
527
+ epilogue_tile_m,
528
+ num_elts_per_sf=16,
529
+ )
530
+ gemm2_scales_shuffled.append(
531
+ nvfp4_block_scale_interleave(
532
+ w2_weight_scale[i]
533
+ .view(torch.uint8)[permute_sf_indices.to(w2_weight_scale.device)]
534
+ .contiguous()
535
+ )
536
+ )
537
+ # w2 bias
538
+ permute_indices = get_w2_permute_indices_with_cache(
539
+ _cache_permute_indices,
540
+ w2_bias[i].clone().reshape(-1, 1),
541
+ epilogue_tile_m,
542
+ )
543
+ gemm2_bias_shuffled.append(
544
+ w2_bias[i]
545
+ .clone()
546
+ .reshape(-1, 1)[permute_indices.to(w2_bias.device)]
547
+ .contiguous()
548
+ )
549
+
550
+ w13_weight = torch.stack(gemm1_weights_shuffled)
551
+ w13_weight_scale = (
552
+ torch.stack(gemm1_scales_shuffled)
553
+ .reshape(num_experts, 2 * intermediate_size, hidden_size // sf_block_size)
554
+ .view(torch.float8_e4m3fn)
555
+ )
556
+ w2_weight = torch.stack(gemm2_weights_shuffled)
557
+ w2_weight_scale = (
558
+ torch.stack(gemm2_scales_shuffled)
559
+ .reshape(num_experts, hidden_size, intermediate_size // sf_block_size)
560
+ .view(torch.float8_e4m3fn)
561
+ )
562
+ w13_bias = torch.stack(gemm1_bias_shuffled).reshape(num_experts, -1)
563
+ w2_bias = torch.stack(gemm2_bias_shuffled).reshape(num_experts, -1)
564
+
565
+ return (
566
+ w13_weight,
567
+ w2_weight,
568
+ w13_weight_scale,
569
+ w2_weight_scale,
570
+ w13_bias,
571
+ w2_bias,
572
+ )
573
+
574
+ elif mxfp4_backend in (
575
+ Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
576
+ Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
577
+ ):
578
+ # De-interleave and swap for w13 weight, bias, and scales
579
+ w13_w = w13_weight.data
580
+ gate_w, up_w = w13_w[:, ::2, :], w13_w[:, 1::2, :]
581
+ deinterleaved_w13_w = torch.cat([gate_w, up_w], dim=1)
582
+ w1_w, w3_w = torch.chunk(deinterleaved_w13_w, 2, dim=1)
583
+ w13_weight_swapped = torch.cat([w3_w, w1_w], dim=1)
584
+
585
+ assert w13_bias is not None and w2_bias is not None
586
+ w13_b = w13_bias.data.to(torch.float32)
587
+ gate_b, up_b = w13_b[:, ::2], w13_b[:, 1::2]
588
+ deinterleaved_w13_b = torch.cat([gate_b, up_b], dim=1)
589
+ b1, b3 = torch.chunk(deinterleaved_w13_b, 2, dim=-1)
590
+ w13_bias_swapped = torch.cat([b3, b1], dim=-1).to(torch.bfloat16)
591
+
592
+ w13_s = w13_weight_scale.data
593
+ gate_s, up_s = w13_s[:, ::2, :], w13_s[:, 1::2, :]
594
+ deinterleaved_w13_s = torch.cat([gate_s, up_s], dim=1)
595
+ s1, s3 = torch.chunk(deinterleaved_w13_s, 2, dim=1)
596
+ w13_scale_swapped = torch.cat([s3, s1], dim=1)
597
+
598
+ if mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8:
599
+ from flashinfer import block_scale_interleave
600
+
601
+ orig_shape = w13_scale_swapped.shape
602
+ w13_scale_interleaved = block_scale_interleave(
603
+ w13_scale_swapped.view(torch.uint8)
604
+ ).reshape(orig_shape)
605
+
606
+ w2_s = w2_weight_scale.data
607
+ orig_shape = w2_s.shape
608
+ w2_scale_interleaved = block_scale_interleave(
609
+ w2_s.view(torch.uint8)
610
+ ).reshape(orig_shape)
611
+
612
+ return (
613
+ w13_weight_swapped,
614
+ w2_weight,
615
+ w13_scale_interleaved,
616
+ w2_scale_interleaved,
617
+ w13_bias_swapped,
618
+ w2_bias,
619
+ )
620
+
621
+ else:
622
+ assert mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16
623
+
624
+ def _interleave_mxfp4_cutlass_sm90(w):
625
+ w_shape = w.shape
626
+ w_interleaved = w.reshape(w_shape[0], w_shape[1], (w_shape[2] // 4), 4)
627
+ w_interleaved = w_interleaved.permute(0, 2, 1, 3)
628
+ w_interleaved = w_interleaved.reshape(
629
+ w_shape[0], w_shape[2] // 4, w_shape[1] * 4
630
+ )
631
+ return w_interleaved
632
+
633
+ w31_scales = w13_scale_swapped.to(torch.uint8)
634
+ w31_scales_interleaved = _interleave_mxfp4_cutlass_sm90(w31_scales)
635
+
636
+ w2_scale = w2_weight_scale.data.to(torch.uint8)
637
+ w2_scale_interleaved = _interleave_mxfp4_cutlass_sm90(w2_scale)
638
+
639
+ return (
640
+ w13_weight_swapped,
641
+ w2_weight,
642
+ w31_scales_interleaved,
643
+ w2_scale_interleaved,
644
+ w13_bias_swapped,
645
+ w2_bias,
646
+ )
647
+
648
+ elif mxfp4_backend == Mxfp4MoeBackend.CK:
649
+ from vllm._aiter_ops import rocm_aiter_ops
650
+
651
+ if w13_bias is not None:
652
+ w13_bias = w13_bias.data.to(torch.float32)
653
+ if w2_bias is not None:
654
+ w2_bias = w2_bias.data.to(torch.float32)
655
+
656
+ e, n, k = w13_weight.shape
657
+
658
+ # De-interleave w13 rows: gate/up pairs -> contiguous gate, up blocks
659
+ w13_weight.view(torch.uint8).copy_(
660
+ w13_weight.data.view(torch.uint8)
661
+ .view(e, n // 2, 2, k)
662
+ .permute(0, 2, 1, 3)
663
+ .contiguous()
664
+ .view(e, n, k)
665
+ )
666
+ w13_weight_scale.data = (
667
+ w13_weight_scale.data.view(e, n // 2, 2, -1)
668
+ .permute(0, 2, 1, 3)
669
+ .contiguous()
670
+ .view(e, n, -1)
671
+ )
672
+
673
+ # View as native FP4 dtype for AITER shuffle
674
+ w13_weight.data = w13_weight.data.view(torch.float4_e2m1fn_x2)
675
+ w2_weight.data = w2_weight.data.view(torch.float4_e2m1fn_x2)
676
+
677
+ # Shuffle weights and scales for AITER CK kernel layout
678
+ w13_weight.data = rocm_aiter_ops.shuffle_weight_a16w4(w13_weight, 16, True)
679
+ shuffled_w13_scale = rocm_aiter_ops.shuffle_scale_a16w4(
680
+ w13_weight_scale.view(-1, w13_weight_scale.shape[-1]),
681
+ num_experts,
682
+ True,
683
+ )
684
+
685
+ w2_weight.data = rocm_aiter_ops.shuffle_weight_a16w4(w2_weight, 16, False)
686
+ shuffled_w2_scale = rocm_aiter_ops.shuffle_scale_a16w4(
687
+ w2_weight_scale.view(-1, w2_weight_scale.shape[-1]),
688
+ num_experts,
689
+ False,
690
+ )
691
+
692
+ # Permute bias to match de-interleaved weight layout
693
+ if w13_bias is not None:
694
+ w13_bias = (
695
+ w13_bias.data.view(-1, n // 2, 2)
696
+ .permute(0, 2, 1)
697
+ .contiguous()
698
+ .view(-1, n)
699
+ )
700
+
701
+ return (
702
+ w13_weight,
703
+ w2_weight,
704
+ shuffled_w13_scale,
705
+ shuffled_w2_scale,
706
+ w13_bias,
707
+ w2_bias,
708
+ )
709
+
710
+ elif mxfp4_backend in TRITON_BACKENDS:
711
+ from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
712
+
713
+ assert w13_bias is not None and w2_bias is not None
714
+ w13_bias = w13_bias.to(torch.float32)
715
+ w2_bias = w2_bias.to(torch.float32)
716
+
717
+ w13_weight, w13_flex, w13_scale = _swizzle_mxfp4(
718
+ w13_weight,
719
+ w13_weight_scale,
720
+ )
721
+ w2_weight, w2_flex, w2_scale = _swizzle_mxfp4(
722
+ w2_weight,
723
+ w2_weight_scale,
724
+ )
725
+
726
+ w13_precision_config = PrecisionConfig(
727
+ weight_scale=w13_scale, flex_ctx=FlexCtx(rhs_data=w13_flex)
728
+ )
729
+ w2_precision_config = PrecisionConfig(
730
+ weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex)
731
+ )
732
+
733
+ del layer.w13_weight
734
+ del layer.w2_weight
735
+
736
+ return (
737
+ w13_weight,
738
+ w2_weight,
739
+ w13_precision_config,
740
+ w2_precision_config,
741
+ w13_bias,
742
+ w2_bias,
743
+ )
744
+ else:
745
+ raise ValueError(
746
+ f"Unsupported mxfp4_backend: {mxfp4_backend}: "
747
+ f"should be one of: {list(Mxfp4MoeBackend)}."
748
+ )
749
+
750
+
751
+def make_mxfp4_moe_quant_config(
752
+ mxfp4_backend: Mxfp4MoeBackend,
753
+ w1_scale: Union[torch.Tensor, "PrecisionConfig"],
754
+ w2_scale: Union[torch.Tensor, "PrecisionConfig"],
755
+ w1_bias: torch.Tensor | None = None,
756
+ w2_bias: torch.Tensor | None = None,
757
+) -> FusedMoEQuantConfig | None:
758
+ """Create a FusedMoEQuantConfig for the given MXFP4 backend."""
759
+ if mxfp4_backend in (
760
+ Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
761
+ Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
762
+ ):
763
+ return mxfp4_mxfp8_moe_quant_config(
764
+ w1_bias=w1_bias,
765
+ w2_bias=w2_bias,
766
+ w1_scale=w1_scale,
767
+ w2_scale=w2_scale,
768
+ )
769
+ elif mxfp4_backend in (
770
+ Mxfp4MoeBackend.MARLIN,
771
+ Mxfp4MoeBackend.BATCHED_MARLIN,
772
+ Mxfp4MoeBackend.TRITON,
773
+ Mxfp4MoeBackend.TRITON_UNFUSED,
774
+ Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
775
+ Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
776
+ Mxfp4MoeBackend.CK,
777
+ ):
778
+ return mxfp4_w4a16_moe_quant_config(
779
+ w1_bias=w1_bias,
780
+ w2_bias=w2_bias,
781
+ w1_scale=w1_scale,
782
+ w2_scale=w2_scale,
783
+ )
784
+ else:
785
+ return ocp_mx_moe_quant_config(
786
+ quant_dtype="mxfp4",
787
+ w1_bias=w1_bias,
788
+ w2_bias=w2_bias,
789
+ w1_scale=w1_scale,
790
+ w2_scale=w2_scale,
791
+ )
792
+
793
+
794
+def make_mxfp4_moe_kernel(
795
+ moe_quant_config: FusedMoEQuantConfig,
796
+ moe_config: FusedMoEConfig,
797
+ experts_cls: type[mk.FusedMoEExperts],
798
+ mxfp4_backend: Mxfp4MoeBackend,
799
+ routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
800
+ shared_experts: torch.nn.Module | None = None,
801
+) -> mk.FusedMoEKernel:
802
+ """Create a FusedMoEKernel for the given MXFP4 backend."""
803
+ is_monolithic = issubclass(experts_cls, mk.FusedMoEExpertsMonolithic)
804
+
805
+ # Create Prepare/Finalize.
806
+ prepare_finalize = maybe_make_prepare_finalize(
807
+ moe=moe_config,
808
+ quant_config=moe_quant_config,
809
+ routing_tables=routing_tables,
810
+ allow_new_interface=True,
811
+ use_monolithic=is_monolithic,
812
+ )
813
+ assert prepare_finalize is not None
814
+
815
+ logger.info_once("Using %s", prepare_finalize.__class__.__name__, scope="local")
816
+
817
+ # Create Experts.
818
+ if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
819
+ max_num_tokens = prepare_finalize.max_num_tokens_per_rank()
820
+ assert max_num_tokens is not None
821
+ experts = experts_cls(
822
+ moe_config=moe_config,
823
+ quant_config=moe_quant_config,
824
+ max_num_tokens=max_num_tokens,
825
+ num_dispatchers=prepare_finalize.num_dispatchers(),
826
+ )
827
+ else:
828
+ experts = experts_cls(
829
+ moe_config=moe_config,
830
+ quant_config=moe_quant_config,
831
+ )
832
+
833
+ kernel = mk.FusedMoEKernel(
834
+ prepare_finalize,
835
+ experts,
836
+ shared_experts=(
837
+ shared_experts
838
+ if moe_config.moe_parallel_config.use_deepep_ll_kernels
839
+ else None
840
+ ),
841
+ moe_parallel_config=moe_config.moe_parallel_config,
842
+ inplace=(
843
+ not moe_config.disable_inplace and mxfp4_backend not in TRTLLM_BACKENDS
844
+ ),
845
+ )
846
+
847
+ return kernel
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
-11
@@ -14,7 +14,6 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
14
from vllm.model_executor.layers.fused_moe.config import (
15
FusedMoEConfig,
16
FusedMoEQuantConfig,
17
- mxfp4_w4a16_moe_quant_config,
17
nvfp4_moe_quant_config,
18
nvfp4_w4a16_moe_quant_config,
19
)
@@ -347,16 +346,6 @@ def convert_to_nvfp4_moe_kernel_format(
346
)
347
348
350
-def make_mxfp4_moe_quant_config(
351
- w13_scale: torch.Tensor,
352
- w2_scale: torch.Tensor,
353
-) -> FusedMoEQuantConfig:
354
- return mxfp4_w4a16_moe_quant_config(
355
- w1_scale=w13_scale,
356
- w2_scale=w2_scale,
357
- )
358
-
359
-
349
def make_nvfp4_moe_quant_config(
350
backend: NvFp4MoeBackend,
351
w13_scale: torch.Tensor,
vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
+13
-6
@@ -24,6 +24,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
24
kFp8Static128BlockSym,
25
kFp8StaticChannelSym,
26
kFp8StaticTensorSym,
27
+ kMxfp4Static,
28
)
29
30
@@ -201,6 +202,8 @@ def rocm_aiter_fused_experts(
202
activation_method = ActivationMethod.SILU
203
elif activation == MoEActivation.GELU:
204
activation_method = ActivationMethod.GELU
205
+ elif activation == MoEActivation.SWIGLUOAI:
206
+ activation_method = rocm_aiter_ops.get_aiter_activation_type("swiglu")
207
else:
208
raise ValueError(f"Unsupported activation: {activation}")
209
@@ -247,8 +250,8 @@ def rocm_aiter_fused_experts(
250
251
else:
252
quant_method = QuantMethod.NO.value
250
- # quark moe for mxfp4 w_dtype mxfp4 a_dtype
251
- if quant_config.use_mxfp4_w4a4:
253
+ # mxfp4: both w4a4 (quark) and w4a16 (oracle CK) use BLOCK_1X32
254
+ if quant_config.use_mxfp4_w4a4 or quant_config.use_mxfp4_w4a16:
255
quant_method = QuantMethod.BLOCK_1X32.value
256
# w8a8 block-scaled
257
if quant_config.block_shape is not None and quant_config.use_fp8_w8a8:
@@ -289,6 +292,8 @@ def rocm_aiter_fused_experts(
292
doweight_stage1=apply_router_weight_on_input,
293
num_local_tokens=num_local_tokens,
294
output_dtype=output_dtype,
295
+ bias1=quant_config.w1_bias if quant_config.use_mxfp4_w4a16 else None,
296
+ bias2=quant_config.w2_bias if quant_config.use_mxfp4_w4a16 else None,
297
)
298
299
@@ -319,21 +324,23 @@ class AiterExperts(mk.FusedMoEExpertsModular):
324
weight_key: QuantKey | None,
325
activation_key: QuantKey | None,
326
) -> bool:
322
- # TODO(rob): AITER also supports MXFP4, which is not
323
- # yet supported via an Oracle. Once it is, we will add
324
- # MXFP4 to this list.
327
SUPPORTED_W_A = [
328
(None, None),
329
(kFp8Static128BlockSym, kFp8Dynamic128Sym),
330
(kFp8StaticTensorSym, kFp8StaticTensorSym),
331
(kFp8StaticTensorSym, kFp8DynamicTensorSym),
332
(kFp8StaticChannelSym, kFp8DynamicTokenSym),
333
+ (kMxfp4Static, None),
334
]
335
return (weight_key, activation_key) in SUPPORTED_W_A
336
337
@staticmethod
338
def _supports_activation(activation: MoEActivation) -> bool:
336
- return activation in [MoEActivation.SILU, MoEActivation.GELU]
339
+ return activation in [
340
+ MoEActivation.SILU,
341
+ MoEActivation.GELU,
342
+ MoEActivation.SWIGLUOAI,
343
+ ]
344
345
@staticmethod
346
def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
vllm/model_executor/layers/fused_moe/trtllm_moe.py
deleted
-184
@@ -1,184 +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
- FusedMoEQuantConfig,
12
-)
13
-from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
14
- TopKWeightAndReduceNoOP,
15
-)
16
-from vllm.model_executor.layers.quantization.utils.quant_utils import (
17
- QuantKey,
18
-)
19
-
20
-
21
-class TrtLlmGenExperts(mk.FusedMoEExpertsModular):
22
- """TensorRT-LLM-based fused MoE expert implementation."""
23
-
24
- def __init__(
25
- self,
26
- moe_config: FusedMoEConfig,
27
- quant_config: FusedMoEQuantConfig,
28
- max_capture_size,
29
- ):
30
- super().__init__(moe_config, quant_config)
31
- self.device = torch.accelerator.current_device_index()
32
- self.num_experts = moe_config.num_local_experts
33
- self.gemm1_alpha = torch.tensor(
34
- [1.702] * self.num_experts, dtype=torch.float32, device=self.device
35
- )
36
- self.gemm1_beta = torch.tensor(
37
- [1.0] * self.num_experts, dtype=torch.float32, device=self.device
38
- )
39
- self.gemm1_clamp_limit = torch.tensor(
40
- [7.0] * self.num_experts, dtype=torch.float32, device=self.device
41
- )
42
- self.max_capture_size = max_capture_size
43
-
44
- @staticmethod
45
- def activation_format() -> mk.FusedMoEActivationFormat:
46
- return mk.FusedMoEActivationFormat.Standard
47
-
48
- @staticmethod
49
- def _supports_current_device() -> bool:
50
- raise NotImplementedError(
51
- "TrtLlmGenExperts is not yet used by an Oracle. "
52
- "This method should not be called."
53
- )
54
-
55
- @staticmethod
56
- def _supports_no_act_and_mul() -> bool:
57
- raise NotImplementedError(
58
- "TrtLlmGenExperts is not yet used by an Oracle. "
59
- "This method should not be called."
60
- )
61
-
62
- @staticmethod
63
- def _supports_quant_scheme(
64
- weight_key: QuantKey | None,
65
- activation_key: QuantKey | None,
66
- ) -> bool:
67
- raise NotImplementedError(
68
- "TrtLlmGenExperts is not yet used by an Oracle. "
69
- "This method should not be called."
70
- )
71
-
72
- @staticmethod
73
- def _supports_activation(activation: MoEActivation) -> bool:
74
- raise NotImplementedError(
75
- "TrtLlmGenExperts is not yet used by an Oracle. "
76
- "This method should not be called."
77
- )
78
-
79
- @staticmethod
80
- def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
81
- raise NotImplementedError(
82
- "TrtLlmGenExperts is not yet used by an Oracle. "
83
- "This method should not be called."
84
- )
85
-
86
- def supports_expert_map(self) -> bool:
87
- return True
88
-
89
- def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
90
- return TopKWeightAndReduceNoOP()
91
-
92
- def workspace_shapes(
93
- self,
94
- M: int,
95
- N: int,
96
- K: int,
97
- topk: int,
98
- global_num_experts: int,
99
- local_num_experts: int,
100
- expert_tokens_meta: mk.ExpertTokensMetadata | None,
101
- activation: MoEActivation,
102
- ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
103
- # The workspaces for this implementation are managed by flashinfer.
104
- workspace1 = (0,)
105
- workspace2 = (0,)
106
- output = (M, K)
107
- return (workspace1, workspace2, output)
108
-
109
- def apply(
110
- self,
111
- output: torch.Tensor,
112
- hidden_states: torch.Tensor,
113
- w1: torch.Tensor,
114
- w2: torch.Tensor,
115
- topk_weights: torch.Tensor,
116
- topk_ids: torch.Tensor,
117
- activation: MoEActivation,
118
- global_num_experts: int,
119
- expert_map: torch.Tensor | None,
120
- a1q_scale: torch.Tensor | None,
121
- a2_scale: torch.Tensor | None,
122
- workspace13: torch.Tensor,
123
- workspace2: torch.Tensor,
124
- expert_tokens_meta: mk.ExpertTokensMetadata | None,
125
- apply_router_weight_on_input: bool,
126
- ):
127
- topk = topk_ids.size(-1)
128
- local_num_experts = w1.size(0)
129
- intermediate_size = w2.size(1)
130
- local_expert_offset = self.moe_config.ep_rank * local_num_experts
131
-
132
- x_quant = hidden_states
133
- x_scale = a1q_scale
134
- if x_scale is not None:
135
- x_scale = x_scale.view(torch.float8_e4m3fn).reshape(*x_quant.shape[:-1], -1)
136
-
137
- packed_tensor = (topk_ids.to(torch.int32) << 16) | topk_weights.to(
138
- torch.bfloat16
139
- ).view(torch.int16)
140
-
141
- assert self.w1_scale is not None
142
- assert self.w2_scale is not None
143
- kwargs = {
144
- "topk_ids": packed_tensor,
145
- "routing_bias": None,
146
- "hidden_states": x_quant,
147
- "hidden_states_scale": x_scale,
148
- "gemm1_weights": w1,
149
- "gemm1_weights_scale": self.w1_scale,
150
- "gemm1_bias": self.w1_bias,
151
- "gemm1_alpha": self.gemm1_alpha,
152
- "gemm1_beta": self.gemm1_beta,
153
- "gemm1_clamp_limit": self.gemm1_clamp_limit,
154
- "gemm2_weights": w2,
155
- "gemm2_weights_scale": self.w2_scale,
156
- "gemm2_bias": self.w2_bias,
157
- "output1_scale_scalar": None,
158
- "output1_scale_gate_scalar": None,
159
- "output2_scale_scalar": None,
160
- "num_experts": global_num_experts,
161
- "top_k": topk,
162
- "n_group": None,
163
- "topk_group": None,
164
- "intermediate_size": intermediate_size,
165
- "local_expert_offset": local_expert_offset,
166
- "local_num_experts": local_num_experts,
167
- "routed_scaling_factor": None,
168
- "routing_method_type": 1,
169
- "do_finalize": True,
170
- "output": output,
171
- "tune_max_num_tokens": max(self.max_capture_size, 1),
172
- }
173
-
174
- from flashinfer import trtllm_fp4_block_scale_routed_moe
175
-
176
- from vllm.utils.flashinfer import autotune
177
-
178
- with autotune(False):
179
- # Enable autotune when,
180
- # https://github.com/flashinfer-ai/flashinfer/issues/2023 is
181
- # resolved.
182
- trtllm_fp4_block_scale_routed_moe(**kwargs)
183
-
184
- return output
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+11
-5
@@ -45,11 +45,14 @@ from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
45
make_fp8_moe_quant_config,
46
select_fp8_moe_backend,
47
)
48
+from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
49
+ Mxfp4MoeBackend,
50
+ make_mxfp4_moe_kernel,
51
+ make_mxfp4_moe_quant_config,
52
+)
53
from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
49
- NvFp4MoeBackend,
54
convert_to_nvfp4_moe_kernel_format,
55
is_global_sf_supported_for_nvfp4_backend,
52
- make_mxfp4_moe_quant_config,
56
make_nvfp4_moe_kernel,
57
make_nvfp4_moe_quant_config,
58
select_nvfp4_moe_backend,
@@ -235,7 +238,7 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
238
def __init__(self, moe):
239
super().__init__(moe)
240
self.group_size = 32
238
- self.mxfp4_backend = NvFp4MoeBackend.MARLIN
241
+ self.mxfp4_backend = Mxfp4MoeBackend.MARLIN
242
self.experts_cls = MarlinExperts
243
244
def create_weights(
@@ -310,7 +313,9 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
313
self, layer: torch.nn.Module
314
) -> FusedMoEQuantConfig | None:
315
return make_mxfp4_moe_quant_config(
313
- w13_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale
316
+ mxfp4_backend=self.mxfp4_backend,
317
+ w1_scale=layer.w13_weight_scale,
318
+ w2_scale=layer.w2_weight_scale,
319
)
320
321
def process_weights_after_loading(self, layer: FusedMoE) -> None:
@@ -334,10 +339,11 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
339
340
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
341
if self.moe_quant_config is not None:
337
- self.moe_kernel = make_nvfp4_moe_kernel(
342
+ self.moe_kernel = make_mxfp4_moe_kernel(
343
moe_quant_config=self.moe_quant_config,
344
moe_config=self.moe,
345
experts_cls=self.experts_cls,
346
+ mxfp4_backend=self.mxfp4_backend,
347
shared_experts=layer.shared_experts,
348
routing_tables=layer._maybe_init_expert_routing_tables(),
349
)
vllm/model_executor/layers/quantization/mxfp4.py
+180
-987
@@ -1,12 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from enum import Enum
3
4
import torch
6
-from torch.nn.parameter import Parameter
5
8
-from vllm import envs
9
-from vllm._aiter_ops import rocm_aiter_ops
6
from vllm.config import get_current_vllm_config
7
from vllm.logger import init_logger
8
from vllm.model_executor.layers.attention import Attention
@@ -17,173 +13,31 @@ from vllm.model_executor.layers.fused_moe import (
13
MoEActivation,
14
)
15
from vllm.model_executor.layers.fused_moe import modular_kernel as mk
20
-from vllm.model_executor.layers.fused_moe.all2all_utils import (
21
- maybe_make_prepare_finalize,
22
-)
16
from vllm.model_executor.layers.fused_moe.config import (
17
FusedMoEQuantConfig,
25
- mxfp4_mxfp8_moe_quant_config,
26
- mxfp4_w4a16_moe_quant_config,
27
- ocp_mx_moe_quant_config,
28
-)
29
-from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
30
- BatchedMarlinExperts,
31
- MarlinExperts,
18
)
33
-from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
34
- OAITritonExperts,
35
- UnfusedOAITritonExperts,
19
+from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
20
+ TRITON_BACKENDS,
21
+ Mxfp4MoeBackend,
22
+ convert_to_mxfp4_moe_kernel_format,
23
+ make_mxfp4_moe_kernel,
24
+ make_mxfp4_moe_quant_config,
25
+ mxfp4_round_up_hidden_size_and_intermediate_size,
26
+ select_mxfp4_moe_backend,
27
)
37
-from vllm.model_executor.layers.fused_moe.trtllm_moe import TrtLlmGenExperts
28
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
29
from vllm.model_executor.layers.quantization import QuantizationMethods
30
from vllm.model_executor.layers.quantization.base_config import (
31
QuantizationConfig,
32
QuantizeMethodBase,
33
)
44
-from vllm.model_executor.layers.quantization.utils.marlin_utils import (
45
- get_marlin_input_dtype,
46
-)
47
-from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
48
- prepare_moe_fp4_layer_for_marlin,
49
-)
50
-from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
51
- CK_MXFP4_MOE_DIM_ALIGNMENT,
52
- _can_support_mxfp4,
53
- _swizzle_mxfp4,
54
- get_padding_alignment,
55
-)
34
from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped
57
-from vllm.model_executor.utils import set_weight_attrs
35
+from vllm.model_executor.utils import replace_parameter, set_weight_attrs
36
from vllm.platforms import current_platform
59
-from vllm.utils.flashinfer import has_flashinfer
60
-from vllm.utils.import_utils import has_triton_kernels
61
-from vllm.utils.math_utils import round_up
37
38
logger = init_logger(__name__)
39
40
66
-# enum for mxfp4 backend
67
-class Mxfp4Backend(Enum):
68
- NONE = 0
69
-
70
- # FlashInfer Backend
71
- SM100_FI_MXFP4_MXFP8_TRTLLM = 1
72
- SM100_FI_MXFP4_MXFP8_CUTLASS = 2
73
- SM100_FI_MXFP4_BF16 = 3
74
- SM90_FI_MXFP4_BF16 = 4
75
-
76
- # Marlin Backend
77
- MARLIN = 5
78
-
79
- # Triton Backend
80
- TRITON = 6
81
-
82
- CK = 7
83
-
84
-
85
-def get_mxfp4_backend_with_lora() -> Mxfp4Backend:
86
- """
87
- Not all MXFP4 backends support LoRA. Select backends that are known to
88
- have LoRA support.
89
- """
90
- if not current_platform.is_cuda():
91
- return Mxfp4Backend.NONE
92
-
93
- # If FlashInfer is not available, try either Marlin or Triton
94
- triton_kernels_supported = (
95
- has_triton_kernels()
96
- # NOTE: triton_kernels are only confirmed to work on SM90 and SM100
97
- # SM110 fails with this error: https://github.com/vllm-project/vllm/issues/29317
98
- # SM120 needs this fix: https://github.com/triton-lang/triton/pull/8498
99
- and (9, 0) <= current_platform.get_device_capability() < (11, 0)
100
- )
101
- if envs.VLLM_MXFP4_USE_MARLIN is False and triton_kernels_supported:
102
- logger.info_once("[get_mxfp4_backend_with_lora] Using Triton backend")
103
- return Mxfp4Backend.TRITON
104
-
105
- logger.info_once("[get_mxfp4_backend_with_lora] Using Marlin backend")
106
- return Mxfp4Backend.MARLIN
107
-
108
-
109
-def get_mxfp4_backend(with_lora_support: bool) -> Mxfp4Backend:
110
- # Backend Selection
111
-
112
- if with_lora_support:
113
- return get_mxfp4_backend_with_lora()
114
-
115
- if current_platform.is_cuda():
116
- if (
117
- current_platform.is_device_capability(90)
118
- and has_flashinfer()
119
- and envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16
120
- ):
121
- logger.info_once("Using FlashInfer MXFP4 BF16 backend for SM90")
122
- return Mxfp4Backend.SM90_FI_MXFP4_BF16
123
- elif (
124
- current_platform.is_device_capability_family(100)
125
- and has_flashinfer()
126
- and envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS
127
- ):
128
- logger.info_once("Using FlashInfer MXFP4 MXFP8 CUTLASS backend for SM100")
129
- return Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS
130
- elif (
131
- current_platform.is_device_capability_family(100)
132
- and has_flashinfer()
133
- and envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
134
- ):
135
- logger.info_once(
136
- "Using FlashInfer MXFP4 MXFP8 TRTLLM backend for SM100", scope="local"
137
- )
138
- return Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
139
- elif current_platform.is_device_capability_family(100) and has_flashinfer():
140
- logger.info_once(
141
- "Using FlashInfer MXFP4 BF16 backend for SM100, "
142
- "For faster performance on SM100, consider setting "
143
- "VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8=1, though this may impact "
144
- "accuracy."
145
- )
146
- return Mxfp4Backend.SM100_FI_MXFP4_BF16
147
- elif (
148
- current_platform.is_device_capability_family(100)
149
- or current_platform.is_device_capability(90)
150
- ) and not has_flashinfer():
151
- logger.warning_once(
152
- "MXFP4 MoE is enabled on Hopper/Blackwell but FlashInfer "
153
- "is not available. This may result in degraded performance. "
154
- "Please `pip install vllm[flashinfer]` for best results."
155
- )
156
-
157
- # If FlashInfer is not available, try either Marlin or Triton
158
- triton_kernels_supported = (
159
- has_triton_kernels()
160
- # NOTE: triton_kernels are only confirmed to work on SM90 and SM100
161
- # SM110 fails with this error: https://github.com/vllm-project/vllm/issues/29317
162
- # SM120 needs this fix: https://github.com/triton-lang/triton/pull/8498
163
- and (9, 0) <= current_platform.get_device_capability() < (11, 0)
164
- )
165
- if envs.VLLM_MXFP4_USE_MARLIN or not triton_kernels_supported:
166
- logger.info_once("Using Marlin backend")
167
- return Mxfp4Backend.MARLIN
168
- else:
169
- logger.info_once("Using Triton backend")
170
- return Mxfp4Backend.TRITON
171
- elif current_platform.is_xpu():
172
- logger.info_once("Using xpu backend on XPU")
173
- return Mxfp4Backend.MARLIN
174
- elif current_platform.is_rocm():
175
- from vllm.platforms.rocm import on_gfx950
176
-
177
- if rocm_aiter_ops.is_enabled() and on_gfx950():
178
- logger.info_once("Using CK MXFP4 MoE backend (Aiter ROCm)")
179
- return Mxfp4Backend.CK
180
- elif has_triton_kernels():
181
- logger.info_once("Using Triton backend")
182
- return Mxfp4Backend.TRITON
183
-
184
- return Mxfp4Backend.NONE
185
-
186
-
41
class Mxfp4Config(QuantizationConfig):
42
def __init__(self, ignored_layers: list[str] | None = None):
43
super().__init__()
@@ -219,9 +73,6 @@ class Mxfp4Config(QuantizationConfig):
73
fused_mapping=self.packed_modules_mapping,
74
):
75
return UnquantizedLinearMethod()
222
- # TODO: Add support for MXFP4 Linear Method.
223
- # MXFP4 LinearMethod is available in AMD-Quark, refer to that implementation
224
- # if you are interested in enabling MXFP4 here.
76
logger.debug_once(
77
"MXFP4 linear layer is not implemented - falling back to "
78
"UnquantizedLinearMethod.",
@@ -232,10 +83,8 @@ class Mxfp4Config(QuantizationConfig):
83
if current_platform.is_xpu():
84
return XpuMxfp4MoEMethod(layer.moe_config)
85
else:
235
- quant_method = Mxfp4MoEMethod(layer.moe_config)
236
- return quant_method
86
+ return Mxfp4MoEMethod(layer.moe_config)
87
elif isinstance(layer, Attention):
238
- # TODO: Add support for MXFP4 Attention.
88
logger.debug_once(
89
"MXFP4 attention layer is not implemented. "
90
"Skipping quantization for this layer.",
@@ -254,51 +103,36 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
103
def __init__(self, moe: FusedMoEConfig):
104
super().__init__(moe)
105
self.weight_dtype = "mxfp4"
257
- self.mxfp4_backend = get_mxfp4_backend(moe.is_lora_enabled)
106
+ self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(moe)
107
108
self.max_capture_size = (
109
get_current_vllm_config().compilation_config.max_cudagraph_capture_size
110
)
111
263
- # CK's pre-compiled MXFP4 MoE GEMM kernel instances have dimension
264
- # alignment requirements. Fall back to Triton when not met.
265
- if (
266
- self.mxfp4_backend == Mxfp4Backend.CK
267
- and moe.intermediate_size_per_partition % CK_MXFP4_MOE_DIM_ALIGNMENT != 0
268
- ):
269
- if has_triton_kernels():
270
- logger.warning_once(
271
- "CK MXFP4 MoE GEMM does not support "
272
- "intermediate_size_per_partition=%d (not a multiple of "
273
- "%d). Falling back to Triton backend.",
274
- moe.intermediate_size_per_partition,
275
- CK_MXFP4_MOE_DIM_ALIGNMENT,
276
- )
277
- self.mxfp4_backend = Mxfp4Backend.TRITON
278
- else:
279
- raise ValueError(
280
- f"CK MXFP4 MoE GEMM does not support "
281
- f"intermediate_size_per_partition="
282
- f"{moe.intermediate_size_per_partition} (not a multiple "
283
- f"of {CK_MXFP4_MOE_DIM_ALIGNMENT}) and no Triton "
284
- f"fallback is available. Use a compatible "
285
- f"tensor_parallel_size."
286
- )
287
-
288
- assert self.mxfp4_backend != Mxfp4Backend.NONE, (
289
- f"get_mxfp4_backend(with_lora_support={moe.is_lora_enabled}) found"
290
- "no compatible MXFP4 MoE backend (FlashInfer/Marlin/Triton)."
291
- "Please check your environment and try again."
292
- )
112
self._cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
294
- # Initialized in process_weights_after_loading for CUTLASS/SM90 backends
113
self.moe_kernel: mk.FusedMoEKernel | None = None
114
115
+ # Round up dims once based on backend. This mutates the shared
116
+ # FusedMoEConfig in-place so that create_weights() and all
117
+ # downstream code see the padded dimensions. This must happen
118
+ # before create_weights() is called.
119
+ self.moe.hidden_dim, self.moe.intermediate_size_per_partition = (
120
+ mxfp4_round_up_hidden_size_and_intermediate_size(
121
+ self.mxfp4_backend,
122
+ self.moe.hidden_dim,
123
+ self.moe.intermediate_size_per_partition,
124
+ )
125
+ )
126
+
127
+ # Used for triton kernel precision configs
128
+ self.w13_precision_config = None
129
+ self.w2_precision_config = None
130
+
131
@property
132
def skip_forward_padding(self) -> bool:
133
# SM100_FI_MXFP4_MXFP8_TRTLLM supports padding with mxfp8 quant
134
# so can skip the padding in the forward before applying the moe method
301
- return self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
135
+ return self.mxfp4_backend == Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8
136
137
def create_weights(
138
self,
@@ -312,77 +146,14 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
146
self.num_experts = num_experts
147
weight_dtype = torch.uint8
148
scale_dtype = torch.uint8
315
-
316
- # FIXME (zyongye): ship after torch and safetensors support mxfp4
317
- # is_torch_mxfp4_available = (
318
- # hasattr(torch, "float4_e2m1fn_x2") and
319
- # hasattr(torch, "float8_e8m0fnu"))
320
- # if is_torch_mxfp4_available:
321
- # weight_dtype = torch.float4_e2m1fn_x2
322
- # scale_dtype = torch.float8_e8m0fnu
323
-
149
mxfp4_block = 32
150
326
- intermediate_size_per_partition_after_pad = intermediate_size_per_partition
327
- if self.mxfp4_backend == Mxfp4Backend.MARLIN:
328
- # The moe marlin kernel requires that for each linear
329
- # n % 256 == 0 and k % 128 == 0.
330
- # In gate_up_proj:
331
- # n = 2 * intermediate_size_per_partition_after_pad
332
- # k = hidden_size
333
- # In down_proj
334
- # n = hidden_size
335
- # k = intermediate_size_per_partition_after_pad
336
- intermediate_size_per_partition_after_pad = round_up(
337
- intermediate_size_per_partition, 128
338
- )
339
- if current_platform.is_xpu():
340
- hidden_size = round_up(hidden_size, 128)
341
- else:
342
- hidden_size = round_up(hidden_size, 256)
343
-
344
- layer.params_dtype = params_dtype
345
- layer.num_experts = num_experts
346
- layer.hidden_size = hidden_size
347
- layer.intermediate_size_per_partition = (
348
- intermediate_size_per_partition_after_pad
349
- )
350
- elif (
351
- self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
352
- or self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16
353
- ):
354
- # pad the intermediate size to be a multiple of 2 * mxfp4_block
355
- # for to hold non-uniform sharded tensor as well as swizzling
356
- # other padding to increase performance
357
- intermediate_size_per_partition_after_pad = round_up(
358
- intermediate_size_per_partition, 256
359
- )
360
- hidden_size = round_up(hidden_size, 256)
361
- elif (
362
- self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS
363
- or self.mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16
364
- ):
365
- intermediate_size_per_partition_after_pad = round_up(
366
- intermediate_size_per_partition, 128
367
- )
368
- hidden_size = round_up(hidden_size, 128)
369
- elif current_platform.is_rocm():
370
- pad_align = get_padding_alignment()
371
- intermediate_size_per_partition_after_pad = round_up(
372
- intermediate_size_per_partition, pad_align
373
- )
374
- hidden_size = round_up(hidden_size, pad_align)
375
- else:
376
- intermediate_size_per_partition_after_pad = round_up(
377
- intermediate_size_per_partition, 64
378
- )
379
-
380
- self.intermediate_size = intermediate_size_per_partition_after_pad
381
- self.hidden_size = hidden_size
382
- self.hidden_pad = extra_weight_attrs.get("hidden_pad", 0)
383
- self.intermediate_pad = (
384
- intermediate_size_per_partition_after_pad - intermediate_size_per_partition
151
+ # Use pre-rounded sizes from config
152
+ self.intermediate_size = intermediate_size_per_partition_after_pad = (
153
+ self.moe.intermediate_size_per_partition
154
)
155
+ self.hidden_size = hidden_size = self.moe.hidden_dim
156
+
157
# Fused gate_up_proj (column parallel)
158
w13_weight = torch.nn.Parameter(
159
torch.zeros(
@@ -408,17 +179,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
179
layer.register_parameter("w13_weight_scale", w13_weight_scale)
180
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
181
411
- w13_bias = torch.nn.Parameter(
412
- torch.zeros(
413
- num_experts,
414
- 2 * intermediate_size_per_partition_after_pad,
415
- dtype=torch.bfloat16,
416
- ),
417
- requires_grad=False,
418
- )
419
- layer.register_parameter("w13_bias", w13_bias)
420
- set_weight_attrs(w13_bias, extra_weight_attrs)
421
-
182
# down_proj (row parallel)
183
w2_weight = torch.nn.Parameter(
184
torch.zeros(
@@ -444,604 +204,170 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
204
layer.register_parameter("w2_weight_scale", w2_weight_scale)
205
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
206
447
- w2_bias = torch.nn.Parameter(
448
- torch.zeros(
449
- num_experts,
450
- hidden_size,
451
- dtype=torch.bfloat16,
452
- ),
453
- requires_grad=False,
454
- )
455
- layer.register_parameter("w2_bias", w2_bias)
456
- set_weight_attrs(w2_bias, extra_weight_attrs)
457
-
458
- def process_weights_after_loading(self, layer):
459
- if self.mxfp4_backend == Mxfp4Backend.MARLIN:
460
- prepare_moe_fp4_layer_for_marlin(
461
- layer, input_dtype=get_marlin_input_dtype()
462
- )
463
-
464
- self.moe_quant_config = self.get_fused_moe_quant_config(layer)
465
- assert self.moe_quant_config is not None
466
-
467
- prepare_finalize = maybe_make_prepare_finalize(
468
- moe=self.moe,
469
- quant_config=self.moe_quant_config,
470
- routing_tables=layer._maybe_init_expert_routing_tables(),
471
- allow_new_interface=True,
472
- )
473
- assert prepare_finalize is not None
474
-
475
- self.moe_kernel = mk.FusedMoEKernel(
476
- prepare_finalize,
477
- MarlinExperts(
478
- self.moe,
479
- self.moe_quant_config,
207
+ if self.moe.has_bias:
208
+ w13_bias = torch.nn.Parameter(
209
+ torch.zeros(
210
+ num_experts,
211
+ 2 * intermediate_size_per_partition_after_pad,
212
+ dtype=torch.bfloat16,
213
),
481
- inplace=not self.moe.disable_inplace,
482
- shared_experts=None,
483
- )
484
- elif (
485
- self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
486
- or self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16
487
- ):
488
- from flashinfer.fp4_quantization import nvfp4_block_scale_interleave
489
- from flashinfer.fused_moe.core import get_w2_permute_indices_with_cache
490
-
491
- layer.gemm1_alpha = Parameter(
492
- torch.tensor([1.702] * self.num_experts, dtype=torch.float32).cuda(),
214
requires_grad=False,
215
)
495
- layer.gemm1_beta = Parameter(
496
- torch.tensor([1.0] * self.num_experts, dtype=torch.float32).cuda(),
497
- requires_grad=False,
498
- )
499
- layer.gemm1_clamp_limit = Parameter(
500
- torch.tensor([7.0] * self.num_experts, dtype=torch.float32).cuda(),
501
- requires_grad=False,
502
- )
503
- sf_block_size = 32 # mxfp4 block size
504
-
505
- assert (
506
- layer.w13_weight.dim() == 3
507
- and layer.w13_weight.shape[0] == self.num_experts
508
- and layer.w13_weight.shape[1] == self.intermediate_size * 2
509
- and layer.w13_weight.shape[2] == self.hidden_size // 2
510
- )
511
- assert (
512
- layer.w13_weight_scale.dim() == 3
513
- and layer.w13_weight_scale.shape[0] == self.num_experts
514
- and layer.w13_weight_scale.shape[1] == self.intermediate_size * 2
515
- and layer.w13_weight_scale.shape[2] == self.hidden_size // sf_block_size
516
- )
517
- assert (
518
- layer.w2_weight.dim() == 3
519
- and layer.w2_weight.shape[0] == self.num_experts
520
- and layer.w2_weight.shape[1] == self.hidden_size
521
- and layer.w2_weight.shape[2] == self.intermediate_size // 2
522
- )
523
- assert (
524
- layer.w2_weight_scale.dim() == 3
525
- and layer.w2_weight_scale.shape[1] == self.hidden_size
526
- and layer.w2_weight_scale.shape[2]
527
- == self.intermediate_size // sf_block_size
528
- )
529
- assert (
530
- layer.w13_bias.dim() == 2
531
- and layer.w13_bias.shape[0] == self.num_experts
532
- and layer.w13_bias.shape[1] == self.intermediate_size * 2
533
- )
534
- assert (
535
- layer.w2_bias.dim() == 2
536
- and layer.w2_bias.shape[0] == self.num_experts
537
- and layer.w2_bias.shape[1] == self.hidden_size
538
- )
539
-
540
- w13_weight_scale = layer.w13_weight_scale.data
541
- w2_weight_scale = layer.w2_weight_scale.data
542
- w13_weight = layer.w13_weight.data
543
- w2_weight = layer.w2_weight.data
544
- w13_bias = layer.w13_bias.data.to(torch.float32)
545
- w2_bias = layer.w2_bias.data.to(torch.float32)
546
-
547
- # Swap w1 and w3 as the definition of
548
- # swiglu is different in the trtllm-gen
549
- def swap_every_two_rows(x, axis=-1):
550
- shape = x.shape
551
- if axis < 0:
552
- axis = len(shape) + axis
553
-
554
- # Create a new shape with pairs swapped along specified axis
555
- new_shape = list(shape)
556
- new_shape[axis] = shape[axis] // 2
557
- new_shape.insert(axis + 1, 2)
216
+ layer.register_parameter("w13_bias", w13_bias)
217
+ set_weight_attrs(w13_bias, extra_weight_attrs)
218
559
- # Reshape to expose pairs, swap them, and reshape back
560
- x = x.reshape(*new_shape)
561
- x = x.flip(axis + 1)
562
- new_shape = list(shape)
563
- return x.reshape(*new_shape)
564
-
565
- w13_weight_scale = swap_every_two_rows(w13_weight_scale, -2)
566
- w13_weight = swap_every_two_rows(w13_weight, -2)
567
- w13_bias = swap_every_two_rows(w13_bias, -1)
568
-
569
- # Do not interleave as the checkpoint is already interleaved
570
-
571
- # Shuffle weights and scaling factors for transposed mma output
572
- gemm1_weights_mxfp4_shuffled = []
573
- gemm1_scales_mxfp4_shuffled = []
574
- gemm2_weights_mxfp4_shuffled = []
575
- gemm2_scales_mxfp4_shuffled = []
576
- gemm1_bias_shuffled = []
577
- gemm2_bias_shuffled = []
578
- epilogue_tile_m = 128 # FIXME: this depends on the kernel internals
579
- for i in range(self.num_experts):
580
- # w13 weight shuffling
581
- permute_indices = get_w2_permute_indices_with_cache(
582
- self._cache_permute_indices,
583
- w13_weight[i].view(torch.uint8),
584
- epilogue_tile_m,
585
- )
586
- gemm1_weights_mxfp4_shuffled.append(
587
- w13_weight[i]
588
- .view(torch.uint8)[permute_indices.to(w13_weight.device)]
589
- .contiguous()
590
- )
591
- # w13 scale shuffling
592
- permute_sf_indices = get_w2_permute_indices_with_cache(
593
- self._cache_permute_indices,
594
- w13_weight_scale[i].view(torch.uint8),
595
- epilogue_tile_m,
596
- num_elts_per_sf=16,
597
- )
598
- gemm1_scales_mxfp4_shuffled.append(
599
- nvfp4_block_scale_interleave(
600
- w13_weight_scale[i]
601
- .view(torch.uint8)[
602
- permute_sf_indices.to(w13_weight_scale.device)
603
- ]
604
- .contiguous()
605
- )
606
- )
607
- # w13 bias shuffling
608
- permute_bias_indices = get_w2_permute_indices_with_cache(
609
- self._cache_permute_indices,
610
- w13_bias[i].clone().reshape(-1, 1),
611
- epilogue_tile_m,
612
- )
613
- gemm1_bias_shuffled.append(
614
- w13_bias[i]
615
- .clone()
616
- .reshape(-1, 1)[permute_bias_indices.to(w13_bias.device)]
617
- .contiguous()
618
- )
619
- # w2 weight shuffling
620
- permute_indices = get_w2_permute_indices_with_cache(
621
- self._cache_permute_indices,
622
- w2_weight[i].view(torch.uint8),
623
- epilogue_tile_m,
624
- )
625
- gemm2_weights_mxfp4_shuffled.append(
626
- w2_weight[i]
627
- .view(torch.uint8)[permute_indices.to(w2_weight.device)]
628
- .contiguous()
629
- )
630
- # w2 scale shuffling
631
- permute_sf_indices = get_w2_permute_indices_with_cache(
632
- self._cache_permute_indices,
633
- w2_weight_scale[i].view(torch.uint8),
634
- epilogue_tile_m,
635
- num_elts_per_sf=16,
636
- )
637
- gemm2_scales_mxfp4_shuffled.append(
638
- nvfp4_block_scale_interleave(
639
- w2_weight_scale[i]
640
- .view(torch.uint8)[
641
- permute_sf_indices.to(w2_weight_scale.device)
642
- ]
643
- .contiguous()
644
- )
645
- )
646
- # w2 bias shuffling
647
- permute_indices = get_w2_permute_indices_with_cache(
648
- self._cache_permute_indices,
649
- w2_bias[i].clone().reshape(-1, 1),
650
- epilogue_tile_m,
651
- )
652
- gemm2_bias_shuffled.append(
653
- w2_bias[i]
654
- .clone()
655
- .reshape(-1, 1)[permute_indices.to(w2_bias.device)]
656
- .contiguous()
657
- )
658
-
659
- w13_weight = torch.stack(gemm1_weights_mxfp4_shuffled)
660
- w13_weight_scale = (
661
- torch.stack(gemm1_scales_mxfp4_shuffled)
662
- .reshape(
663
- self.num_experts,
664
- 2 * self.intermediate_size,
665
- self.hidden_size // sf_block_size,
666
- )
667
- .view(torch.float8_e4m3fn)
668
- )
669
-
670
- w2_weight = torch.stack(gemm2_weights_mxfp4_shuffled)
671
- w2_weight_scale = (
672
- torch.stack(gemm2_scales_mxfp4_shuffled)
673
- .reshape(
674
- self.num_experts,
675
- self.hidden_size,
676
- self.intermediate_size // sf_block_size,
677
- )
678
- .view(torch.float8_e4m3fn)
679
- )
680
-
681
- layer.w13_weight = Parameter(w13_weight, requires_grad=False)
682
- layer.w13_weight_scale = Parameter(w13_weight_scale, requires_grad=False)
683
- layer.w2_weight = Parameter(w2_weight, requires_grad=False)
684
- layer.w2_weight_scale = Parameter(w2_weight_scale, requires_grad=False)
685
- layer.w13_bias = Parameter(
686
- torch.stack(gemm1_bias_shuffled).reshape(self.num_experts, -1),
687
- requires_grad=False,
688
- )
689
- layer.w2_bias = Parameter(
690
- torch.stack(gemm2_bias_shuffled).reshape(self.num_experts, -1),
219
+ w2_bias = torch.nn.Parameter(
220
+ torch.zeros(
221
+ num_experts,
222
+ hidden_size,
223
+ dtype=torch.bfloat16,
224
+ ),
225
requires_grad=False,
226
)
693
- elif (
694
- self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS
695
- or self.mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16
696
- ):
697
- sf_block_size = 32 # mxfp4 block size
227
+ layer.register_parameter("w2_bias", w2_bias)
228
+ set_weight_attrs(w2_bias, extra_weight_attrs)
229
699
- # Common shape assertions
700
- assert (
701
- layer.w13_weight.dim() == 3
702
- and layer.w13_weight.shape[0] == self.num_experts
703
- and layer.w13_weight.shape[1] == self.intermediate_size * 2
704
- and layer.w13_weight.shape[2] == self.hidden_size // 2
705
- )
706
- assert (
707
- layer.w13_weight_scale.dim() == 3
708
- and layer.w13_weight_scale.shape[0] == self.num_experts
709
- and layer.w13_weight_scale.shape[1] == self.intermediate_size * 2
710
- and layer.w13_weight_scale.shape[2] == self.hidden_size // sf_block_size
711
- )
712
- assert (
713
- layer.w2_weight.dim() == 3
714
- and layer.w2_weight.shape[0] == self.num_experts
715
- and layer.w2_weight.shape[1] == self.hidden_size
716
- and layer.w2_weight.shape[2] == self.intermediate_size // 2
717
- )
718
- assert (
719
- layer.w2_weight_scale.dim() == 3
720
- and layer.w2_weight_scale.shape[1] == self.hidden_size
721
- and layer.w2_weight_scale.shape[2]
722
- == self.intermediate_size // sf_block_size
723
- )
230
+ def _setup_kernel(
231
+ self,
232
+ layer: FusedMoE,
233
+ w13: torch.Tensor,
234
+ w2: torch.Tensor,
235
+ w13_scale: torch.Tensor,
236
+ w2_scale: torch.Tensor,
237
+ w13_bias: torch.Tensor | None = None,
238
+ w2_bias: torch.Tensor | None = None,
239
+ ) -> None:
240
+ num_experts = self.num_experts
241
+ intermediate_size = self.intermediate_size
242
+ hidden_size = self.hidden_size
243
+ sf_block_size = 32
244
+
245
+ # Shape assertions
246
+ assert (
247
+ w13.dim() == 3
248
+ and w13.shape[0] == num_experts
249
+ and w13.shape[1] == intermediate_size * 2
250
+ and w13.shape[2] == hidden_size // 2
251
+ )
252
+ assert (
253
+ w13_scale.dim() == 3
254
+ and w13_scale.shape[0] == num_experts
255
+ and w13_scale.shape[1] == intermediate_size * 2
256
+ and w13_scale.shape[2] == hidden_size // sf_block_size
257
+ )
258
+ assert (
259
+ w2.dim() == 3
260
+ and w2.shape[0] == num_experts
261
+ and w2.shape[1] == hidden_size
262
+ and w2.shape[2] == intermediate_size // 2
263
+ )
264
+ assert (
265
+ w2_scale.dim() == 3
266
+ and w2_scale.shape[1] == hidden_size
267
+ and w2_scale.shape[2] == intermediate_size // sf_block_size
268
+ )
269
+ if w13_bias is not None:
270
assert (
725
- layer.w13_bias.dim() == 2
726
- and layer.w13_bias.shape[0] == self.num_experts
727
- and layer.w13_bias.shape[1] == self.intermediate_size * 2
271
+ w13_bias.dim() == 2
272
+ and w13_bias.shape[0] == num_experts
273
+ and w13_bias.shape[1] == intermediate_size * 2
274
)
275
+ if w2_bias is not None:
276
assert (
730
- layer.w2_bias.dim() == 2
731
- and layer.w2_bias.shape[0] == self.num_experts
732
- and layer.w2_bias.shape[1] == self.hidden_size
733
- )
734
-
735
- # De-interleave and swap for w13 weight, bias, and scales
736
- w13_w = layer.w13_weight.data
737
- gate_w, up_w = w13_w[:, ::2, :], w13_w[:, 1::2, :]
738
- deinterleaved_w13_w = torch.cat([gate_w, up_w], dim=1)
739
- w1_w, w3_w = torch.chunk(deinterleaved_w13_w, 2, dim=1)
740
- w13_weight_swapped = torch.cat([w3_w, w1_w], dim=1)
741
-
742
- w13_b = layer.w13_bias.data.to(torch.float32)
743
- gate_b, up_b = w13_b[:, ::2], w13_b[:, 1::2]
744
- deinterleaved_w13_b = torch.cat([gate_b, up_b], dim=1)
745
- b1, b3 = torch.chunk(deinterleaved_w13_b, 2, dim=-1)
746
- w13_bias_swapped = torch.cat([b3, b1], dim=-1).to(torch.bfloat16)
747
-
748
- w13_s = layer.w13_weight_scale.data
749
- gate_s, up_s = w13_s[:, ::2, :], w13_s[:, 1::2, :]
750
- deinterleaved_w13_s = torch.cat([gate_s, up_s], dim=1)
751
- s1, s3 = torch.chunk(deinterleaved_w13_s, 2, dim=1)
752
- w13_scale_swapped = torch.cat([s3, s1], dim=1)
753
-
754
- if self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS:
755
- from flashinfer import block_scale_interleave
756
-
757
- orig_shape = w13_scale_swapped.shape
758
- w13_scale_interleaved = block_scale_interleave(
759
- w13_scale_swapped.view(torch.uint8)
760
- ).reshape(orig_shape)
761
-
762
- w2_s = layer.w2_weight_scale.data
763
- orig_shape = w2_s.shape
764
- w2_scale_interleaved = block_scale_interleave(
765
- w2_s.view(torch.uint8)
766
- ).reshape(orig_shape)
767
-
768
- layer.w13_weight = Parameter(w13_weight_swapped, requires_grad=False)
769
- layer.w13_weight_scale = Parameter(
770
- w13_scale_interleaved, requires_grad=False
771
- )
772
- layer.w13_bias = Parameter(w13_bias_swapped, requires_grad=False)
773
- layer.w2_weight_scale = Parameter(
774
- w2_scale_interleaved, requires_grad=False
775
- )
776
- elif self.mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16:
777
-
778
- def _interleave_mxfp4_cutlass_sm90(w):
779
- w_shape = w.shape
780
- w_interleaved = w.reshape(
781
- w_shape[0], w_shape[1], (w_shape[2] // 4), 4
782
- )
783
- w_interleaved = w_interleaved.permute(0, 2, 1, 3)
784
- w_interleaved = w_interleaved.reshape(
785
- w_shape[0], w_shape[2] // 4, w_shape[1] * 4
786
- )
787
- return w_interleaved
788
-
789
- w31_scales = w13_scale_swapped.to(torch.uint8).view(torch.uint8)
790
- w31_scales_interleaved = _interleave_mxfp4_cutlass_sm90(w31_scales)
791
-
792
- w2_weight_scale = layer.w2_weight_scale.data
793
- w2_scales = w2_weight_scale.to(torch.uint8).view(torch.uint8)
794
- w2_scales_interleaved = _interleave_mxfp4_cutlass_sm90(w2_scales)
795
-
796
- layer.w13_weight = torch.nn.Parameter(
797
- torch.cat([w3_w, w1_w], dim=1), requires_grad=False
798
- )
799
- layer.w13_bias = torch.nn.Parameter(
800
- w13_bias_swapped, requires_grad=False
801
- )
802
- layer.w13_weight_scale = torch.nn.Parameter(
803
- w31_scales_interleaved, requires_grad=False
804
- )
805
- layer.w2_weight_scale = torch.nn.Parameter(
806
- w2_scales_interleaved, requires_grad=False
807
- )
808
-
809
- # theses two kernels go through the `flashinfer_cutlass_fused_moe` path
810
- from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
811
- FlashInferExperts,
277
+ w2_bias.dim() == 2
278
+ and w2_bias.shape[0] == num_experts
279
+ and w2_bias.shape[1] == hidden_size
280
+ )
281
+
282
+ # Convert weights to kernel format
283
+ w13, w2, w13_scale, w2_scale, w13_bias, w2_bias = (
284
+ convert_to_mxfp4_moe_kernel_format(
285
+ mxfp4_backend=self.mxfp4_backend,
286
+ layer=layer,
287
+ w13_weight=w13,
288
+ w2_weight=w2,
289
+ w13_weight_scale=w13_scale,
290
+ w2_weight_scale=w2_scale,
291
+ w13_bias=w13_bias,
292
+ w2_bias=w2_bias,
293
+ _cache_permute_indices=self._cache_permute_indices,
294
)
295
+ )
296
814
- self.moe_quant_config = self.get_fused_moe_quant_config(layer)
815
- assert self.moe_quant_config is not None
816
- prepare_finalize = maybe_make_prepare_finalize(
817
- moe=self.moe,
818
- quant_config=self.moe_quant_config,
297
+ # For TRITON backends, weights are wrapped tensors from triton_kernels
298
+ # that don't support .detach(). Manually assign parameters.
299
+ if self.mxfp4_backend not in TRITON_BACKENDS:
300
+ replace_parameter(layer, "w13_weight", w13)
301
+ replace_parameter(layer, "w2_weight", w2)
302
+ replace_parameter(layer, "w13_weight_scale", w13_scale)
303
+ replace_parameter(layer, "w2_weight_scale", w2_scale)
304
+ else:
305
+ layer.w13_weight = w13
306
+ layer.w2_weight = w2
307
+ self.w13_precision_config = w13_scale
308
+ self.w2_precision_config = w2_scale
309
+
310
+ if w13_bias is not None and w2_bias is not None:
311
+ replace_parameter(layer, "w13_bias", w13_bias)
312
+ replace_parameter(layer, "w2_bias", w2_bias)
313
+
314
+ # Build quant config
315
+ self.moe_quant_config = self.get_fused_moe_quant_config(layer)
316
+
317
+ # Build kernel (modular or monolithic)
318
+ if self.moe_quant_config is not None and self.experts_cls is not None:
319
+ self.moe_kernel = make_mxfp4_moe_kernel(
320
+ moe_quant_config=self.moe_quant_config,
321
+ moe_config=self.moe,
322
+ mxfp4_backend=self.mxfp4_backend,
323
+ experts_cls=self.experts_cls,
324
routing_tables=layer._maybe_init_expert_routing_tables(),
820
- allow_new_interface=True,
821
- )
822
- assert prepare_finalize is not None
823
-
824
- self.moe_kernel = mk.FusedMoEKernel(
825
- prepare_finalize,
826
- FlashInferExperts(
827
- moe_config=self.moe,
828
- quant_config=self.moe_quant_config,
829
- ),
830
- shared_experts=None,
831
- )
832
- elif self.mxfp4_backend == Mxfp4Backend.CK:
833
- if layer.w13_bias is not None:
834
- layer.w13_bias.data = layer.w13_bias.data.to(torch.float32)
835
- if layer.w2_bias.data is not None:
836
- layer.w2_bias.data = layer.w2_bias.data.to(torch.float32)
837
-
838
- e, n, k = layer.w13_weight.shape
839
- layer.w13_weight.view(torch.uint8).copy_(
840
- layer.w13_weight.data.view(torch.uint8)
841
- .view(e, n // 2, 2, k)
842
- .permute(0, 2, 1, 3)
843
- .contiguous()
844
- .view(e, n, k)
845
- )
846
- layer.w13_weight_scale.data = (
847
- layer.w13_weight_scale.data.view(e, n // 2, 2, -1)
848
- .permute(0, 2, 1, 3)
849
- .contiguous()
850
- .view(e, n, -1)
851
- )
852
- layer.w13_weight.data = layer.w13_weight.data.view(torch.float4_e2m1fn_x2)
853
- layer.w2_weight.data = layer.w2_weight.data.view(torch.float4_e2m1fn_x2)
854
-
855
- layer.w13_weight.data = rocm_aiter_ops.shuffle_weight_a16w4(
856
- layer.w13_weight, 16, True
857
- )
858
- shuffled_w13_scale = rocm_aiter_ops.shuffle_scale_a16w4(
859
- layer.w13_weight_scale.view(-1, layer.w13_weight_scale.shape[-1]),
860
- self.num_experts,
861
- True,
862
- )
863
-
864
- layer.w2_weight.data = rocm_aiter_ops.shuffle_weight_a16w4(
865
- layer.w2_weight, 16, False
866
- )
867
- shuffled_w2_scale = rocm_aiter_ops.shuffle_scale_a16w4(
868
- layer.w2_weight_scale.view(-1, layer.w2_weight_scale.shape[-1]),
869
- self.num_experts,
870
- False,
325
+ shared_experts=layer.shared_experts,
326
)
327
873
- layer.w13_bias.data = (
874
- layer.w13_bias.data.view(-1, n // 2, 2)
875
- .permute(0, 2, 1)
876
- .contiguous()
877
- .view(-1, n)
878
- )
879
-
880
- layer.w13_weight_scale = torch.nn.Parameter(
881
- shuffled_w13_scale, requires_grad=False
882
- )
883
- layer.w2_weight_scale = torch.nn.Parameter(
884
- shuffled_w2_scale, requires_grad=False
885
- )
886
- # replace_parameter(layer, "w13_bias", w13_bias)
887
- # replace_parameter(layer, "w13_weight_scale", w13_weight_scale)
888
- # replace_parameter(layer, "w2_weight_scale", w2_weight_scale)
889
- # replace_parameter(layer, "w13_weight", w13_weight)
890
- # replace_parameter(layer, "w2_weight", w2_weight)
891
-
892
- elif self.mxfp4_backend == Mxfp4Backend.TRITON:
893
- from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
894
-
895
- w13_bias = layer.w13_bias.to(torch.float32)
896
- w2_bias = layer.w2_bias.to(torch.float32)
897
-
898
- layer.w13_bias = Parameter(w13_bias, requires_grad=False)
899
- layer.w2_bias = Parameter(w2_bias, requires_grad=False)
900
- # Ideally we'd use FusedMoEModularKernel.prepare_finalize object
901
- # (stored in self.fused_experts) to determine if the MoE has a
902
- # batched activation format. As self.fused_experts is not
903
- # initialized at this point, we resort to checking the MoE config
904
- # directly.
905
- is_batched_moe = (
906
- self.moe.use_deepep_ll_kernels or self.moe.use_nixl_ep_kernels
907
- )
908
- if is_batched_moe:
909
- num_warps = 4 if envs.VLLM_MOE_DP_CHUNK_SIZE <= 512 else 8
910
- else:
911
- num_warps = 8
912
- w13_weight, w13_flex, w13_scale = _swizzle_mxfp4(
913
- layer.w13_weight, layer.w13_weight_scale, num_warps
914
- )
915
- w2_weight, w2_flex, w2_scale = _swizzle_mxfp4(
916
- layer.w2_weight, layer.w2_weight_scale, num_warps
917
- )
328
+ def process_weights_after_loading(self, layer):
329
+ w13 = layer.w13_weight
330
+ w2 = layer.w2_weight
331
+ w13_scale = layer.w13_weight_scale
332
+ w2_scale = layer.w2_weight_scale
333
+ w13_bias = getattr(layer, "w13_bias", None)
334
+ w2_bias = getattr(layer, "w2_bias", None)
335
919
- self.w13_precision_config = PrecisionConfig(
920
- weight_scale=w13_scale, flex_ctx=FlexCtx(rhs_data=w13_flex)
921
- )
922
- self.w2_precision_config = PrecisionConfig(
923
- weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex)
924
- )
925
- self.w13_weight = w13_weight
926
- self.w2_weight = w2_weight
927
- del layer.w13_weight
928
- del layer.w2_weight
929
- layer.w13_weight = w13_weight
930
- layer.w2_weight = w2_weight
336
+ if self.mxfp4_backend == Mxfp4MoeBackend.NONE:
337
+ return
338
932
- else:
933
- raise ValueError(
934
- f"Unsupported mxfp4_backend: {self.mxfp4_backend}: "
935
- f"should be one of: {list(Mxfp4Backend)}."
936
- )
339
+ self._setup_kernel(layer, w13, w2, w13_scale, w2_scale, w13_bias, w2_bias)
340
341
def get_fused_moe_quant_config(
342
self, layer: torch.nn.Module
343
) -> FusedMoEQuantConfig | None:
941
- if self.mxfp4_backend == Mxfp4Backend.MARLIN:
942
- return mxfp4_w4a16_moe_quant_config(
943
- w1_bias=layer.w13_bias,
944
- w2_bias=layer.w2_bias,
945
- w1_scale=layer.w13_weight_scale,
946
- w2_scale=layer.w2_weight_scale,
947
- )
948
- elif self.mxfp4_backend == Mxfp4Backend.TRITON:
344
+ w1_scale = layer.w13_weight_scale
345
+ w2_scale = layer.w2_weight_scale
346
+ w1_bias = getattr(layer, "w13_bias", None)
347
+ w2_bias = getattr(layer, "w2_bias", None)
348
+
349
+ if self.mxfp4_backend in TRITON_BACKENDS:
350
+ assert self.w13_precision_config is not None
351
+ assert self.w2_precision_config is not None
352
w1_scale = self.w13_precision_config
353
w2_scale = self.w2_precision_config
951
- return mxfp4_w4a16_moe_quant_config(
952
- w1_bias=layer.w13_bias,
953
- w2_bias=layer.w2_bias,
954
- w1_scale=w1_scale,
955
- w2_scale=w2_scale,
956
- )
957
- elif self.mxfp4_backend in [
958
- Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM,
959
- Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS,
960
- ]:
961
- return mxfp4_mxfp8_moe_quant_config(
962
- w1_bias=layer.w13_bias,
963
- w2_bias=layer.w2_bias,
964
- w1_scale=layer.w13_weight_scale,
965
- w2_scale=layer.w2_weight_scale,
966
- )
967
- elif self.mxfp4_backend in [
968
- Mxfp4Backend.SM100_FI_MXFP4_BF16,
969
- Mxfp4Backend.SM90_FI_MXFP4_BF16,
970
- Mxfp4Backend.CK,
971
- ]:
972
- return mxfp4_w4a16_moe_quant_config(
973
- w1_bias=layer.w13_bias,
974
- w2_bias=layer.w2_bias,
975
- w1_scale=layer.w13_weight_scale,
976
- w2_scale=layer.w2_weight_scale,
977
- )
978
- else:
979
- w1_scale = layer.w13_weight_scale
980
- w2_scale = layer.w2_weight_scale
981
- return ocp_mx_moe_quant_config(
982
- quant_dtype="mxfp4",
983
- w1_bias=layer.w13_bias,
984
- w2_bias=layer.w2_bias,
985
- w1_scale=w1_scale,
986
- w2_scale=w2_scale,
987
- )
354
+
355
+ return make_mxfp4_moe_quant_config(
356
+ mxfp4_backend=self.mxfp4_backend,
357
+ w1_scale=w1_scale,
358
+ w2_scale=w2_scale,
359
+ w1_bias=w1_bias,
360
+ w2_bias=w2_bias,
361
+ )
362
363
def select_gemm_impl(
364
self,
991
- prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular,
365
+ prepare_finalize: mk.FusedMoEPrepareAndFinalize,
366
layer: torch.nn.Module,
367
) -> mk.FusedMoEExpertsModular:
994
- if (
995
- prepare_finalize.activation_format
996
- == mk.FusedMoEActivationFormat.BatchedExperts
997
- ):
998
- if self.mxfp4_backend == Mxfp4Backend.MARLIN:
999
- max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank()
1000
- assert max_num_tokens_per_rank is not None
1001
- assert self.moe_quant_config is not None
1002
- return BatchedMarlinExperts(
1003
- max_num_tokens=max_num_tokens_per_rank,
1004
- num_dispatchers=prepare_finalize.num_dispatchers(),
1005
- quant_config=self.moe_quant_config,
1006
- moe_config=self.moe,
1007
- )
1008
- else:
1009
- raise NotImplementedError(
1010
- f"Incompatible Mxfp4 backend ({self.mxfp4_backend}) for "
1011
- "EP batched experts format"
1012
- )
1013
- else:
1014
- assert self.moe_quant_config is not None
1015
- if (
1016
- self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
1017
- or self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16
1018
- ):
1019
- # B200 code-path
1020
- kwargs = {
1021
- # TODO(bnell): part of quant_config
1022
- "max_capture_size": self.max_capture_size,
1023
- }
1024
- return TrtLlmGenExperts(self.moe, self.moe_quant_config, **kwargs)
1025
- elif self.mxfp4_backend == Mxfp4Backend.MARLIN:
1026
- return MarlinExperts(self.moe, self.moe_quant_config)
1027
- elif self.mxfp4_backend == Mxfp4Backend.TRITON:
1028
- if self.moe.is_lora_enabled:
1029
- return UnfusedOAITritonExperts(self.moe, self.moe_quant_config)
1030
- return OAITritonExperts(self.moe, self.moe_quant_config)
1031
- else:
1032
- raise NotImplementedError(
1033
- f"Incompatible Mxfp4 backend ({self.mxfp4_backend}) for EP"
1034
- )
1035
-
1036
- @property
1037
- def is_monolithic(self) -> bool:
1038
- if self.moe.is_lora_enabled:
1039
- return False
1040
- return (
1041
- self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
1042
- or self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16
1043
- or self.mxfp4_backend == Mxfp4Backend.TRITON
1044
- or self.mxfp4_backend == Mxfp4Backend.CK
368
+ raise ValueError(
369
+ f"{self.__class__.__name__} uses the new modular kernel "
370
+ "initialization logic. This function should not be called."
371
)
372
373
def apply(
@@ -1053,30 +379,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
379
shared_experts_input: torch.Tensor | None,
380
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
381
assert not self.is_monolithic
1056
- if layer.enable_eplb:
1057
- raise NotImplementedError("EPLB is not supported for mxfp4")
1058
-
1059
- assert _can_support_mxfp4(
1060
- layer.use_grouped_topk,
1061
- layer.topk_group,
1062
- layer.num_expert_group,
1063
- layer.expert_map,
1064
- layer.custom_routing_function,
1065
- layer.e_score_correction_bias,
1066
- layer.apply_router_weight_on_input,
1067
- layer.scoring_func,
1068
- layer.activation,
1069
- layer.eplb_state.expert_load_view,
1070
- layer.eplb_state.logical_to_physical_map,
1071
- layer.eplb_state.logical_replica_count,
1072
- ), "MXFP4 are not supported with this configuration."
1073
-
1074
- assert (
1075
- self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS
1076
- or self.mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16
1077
- or self.mxfp4_backend == Mxfp4Backend.MARLIN
1078
- )
1079
-
382
assert self.moe_kernel is not None
383
return self.moe_kernel.apply(
384
hidden_states=x,
@@ -1098,126 +400,17 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
400
router_logits: torch.Tensor,
401
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
402
assert self.is_monolithic
1101
-
1102
- if layer.enable_eplb:
1103
- raise NotImplementedError("EPLB is not supported for mxfp4")
1104
-
1105
- assert _can_support_mxfp4(
1106
- layer.use_grouped_topk,
1107
- layer.topk_group,
1108
- layer.num_expert_group,
1109
- layer.expert_map,
1110
- layer.custom_routing_function,
1111
- layer.e_score_correction_bias,
1112
- layer.apply_router_weight_on_input,
1113
- layer.scoring_func,
1114
- layer.activation,
1115
- layer.eplb_state.expert_load_view,
1116
- layer.eplb_state.logical_to_physical_map,
1117
- layer.eplb_state.logical_replica_count,
1118
- ), "MXFP4 are not supported with this configuration."
1119
-
1120
- # Apply routing simulation strategy if specified.
1121
- # This applies to all monolithic backends (SM100_FI and TRITON).
1122
- routing_strategy = envs.VLLM_MOE_ROUTING_SIMULATION_STRATEGY
1123
- if routing_strategy == "uniform_random":
1124
- router_logits = torch.rand_like(router_logits)
1125
-
1126
- if (
1127
- self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
1128
- or self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16
1129
- ):
1130
- from flashinfer import trtllm_fp4_block_scale_moe
1131
-
1132
- if self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16:
1133
- assert x.dtype == torch.bfloat16
1134
- x_quant = x
1135
- x_scale = None
1136
- elif self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM:
1137
- from flashinfer import mxfp8_quantize
1138
-
1139
- # x_quant is padded in hidden dimension with alignment=256
1140
- x_quant, x_scale = mxfp8_quantize(
1141
- x,
1142
- is_sf_swizzled_layout=False,
1143
- alignment=256,
1144
- )
1145
- x_scale = x_scale.view(torch.float8_e4m3fn).reshape(*x.shape[:-1], -1)
1146
-
1147
- # output with original unpadded hidden size
1148
- output = torch.empty_like(x)
1149
-
1150
- trtllm_gen_output = trtllm_fp4_block_scale_moe(
1151
- routing_logits=router_logits.to(torch.bfloat16),
1152
- routing_bias=None,
1153
- hidden_states=x_quant,
1154
- hidden_states_scale=x_scale,
1155
- gemm1_weights=layer.w13_weight, # uint8 (e2m1 x 2)
1156
- gemm1_weights_scale=layer.w13_weight_scale, # uint8 (e4m3 x 2)
1157
- gemm1_bias=layer.w13_bias, # fp32 per expert per channel
1158
- gemm1_alpha=layer.gemm1_alpha, # fp32 per expert
1159
- gemm1_beta=layer.gemm1_beta, # fp32 per expert
1160
- gemm1_clamp_limit=layer.gemm1_clamp_limit, # fp32 per expert
1161
- gemm2_weights=layer.w2_weight, # uint8 (e2m1 x 2)
1162
- gemm2_weights_scale=layer.w2_weight_scale, # ue8m0
1163
- gemm2_bias=layer.w2_bias, # fp32 per expert per channel
1164
- output1_scale_scalar=None,
1165
- output1_scale_gate_scalar=None,
1166
- output2_scale_scalar=None,
1167
- num_experts=layer.global_num_experts,
1168
- top_k=layer.top_k,
1169
- n_group=None,
1170
- topk_group=None,
1171
- intermediate_size=self.intermediate_size, # padded to multiple of 256
1172
- local_expert_offset=layer.ep_rank * layer.local_num_experts,
1173
- local_num_experts=self.num_experts,
1174
- routed_scaling_factor=None,
1175
- routing_method_type=1 if layer.renormalize else 0,
1176
- do_finalize=True,
1177
- tune_max_num_tokens=max(self.max_capture_size, 1),
1178
- output=output,
1179
- )[0]
1180
- return trtllm_gen_output
1181
- elif self.mxfp4_backend == Mxfp4Backend.CK:
1182
- topk_weights, topk_ids = rocm_aiter_ops.fused_topk(
1183
- x, router_logits, layer.top_k, True
1184
- )
1185
- output = rocm_aiter_ops.fused_moe(
1186
- x,
1187
- layer.w13_weight,
1188
- layer.w2_weight,
1189
- topk_weights,
1190
- topk_ids,
1191
- activation_method=rocm_aiter_ops.get_aiter_activation_type("swiglu"),
1192
- quant_method=rocm_aiter_ops.get_aiter_quant_type("per_1x32"),
1193
- w1_scale=layer.w13_weight_scale,
1194
- w2_scale=layer.w2_weight_scale,
1195
- doweight_stage1=False,
1196
- hidden_pad=self.hidden_pad // 128 * 128,
1197
- intermediate_pad=self.intermediate_pad // 64 * 64 * 2,
1198
- bias1=layer.w13_bias,
1199
- bias2=layer.w2_bias,
1200
- )
1201
- return output
1202
- elif self.mxfp4_backend == Mxfp4Backend.TRITON:
1203
- from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import ( # noqa: E501
1204
- triton_kernel_moe_forward,
1205
- )
1206
-
1207
- return triton_kernel_moe_forward(
1208
- hidden_states=x,
1209
- w1=layer.w13_weight,
1210
- w2=layer.w2_weight,
1211
- gating_output=router_logits,
1212
- topk=layer.top_k,
1213
- renormalize=layer.renormalize,
1214
- global_num_experts=layer.global_num_experts,
1215
- expert_map=layer.expert_map,
1216
- quant_config=self.moe_quant_config,
1217
- apply_router_weight_on_input=layer.apply_router_weight_on_input,
1218
- )
1219
- else:
1220
- raise ValueError(f"Unsupported backend: {self.mxfp4_backend}")
403
+ assert self.moe_kernel is not None
404
+ return self.moe_kernel.apply_monolithic(
405
+ hidden_states=x,
406
+ w1=layer.w13_weight,
407
+ w2=layer.w2_weight,
408
+ router_logits=router_logits,
409
+ activation=layer.activation,
410
+ global_num_experts=layer.global_num_experts,
411
+ expert_map=layer.expert_map,
412
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
413
+ )
414
415
416
class XpuMxfp4MoEMethod(Mxfp4MoEMethod):
vllm/model_executor/layers/quantization/quark/quark_moe.py
+5
-5
@@ -25,9 +25,9 @@ from vllm.model_executor.layers.fused_moe.config import (
25
ocp_mx_moe_quant_config,
26
)
27
from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe
28
-from vllm.model_executor.layers.quantization.mxfp4 import (
29
- Mxfp4Backend,
30
- get_mxfp4_backend,
28
+from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
29
+ Mxfp4MoeBackend,
30
+ select_mxfp4_moe_backend,
31
)
32
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
33
prepare_fp8_moe_layer_for_marlin,
@@ -699,9 +699,9 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
699
f"Please check that the combination is supported in OCP_MX_Scheme."
700
)
701
702
- self.mxfp4_backend: Mxfp4Backend | None = None
702
+ self.mxfp4_backend: Mxfp4MoeBackend | None = None
703
if self.ocp_mx_scheme == "w_mxfp4":
704
- self.mxfp4_backend = get_mxfp4_backend(moe.is_lora_enabled)
704
+ self.mxfp4_backend, _ = select_mxfp4_moe_backend(moe)
705
706
if self.input_quant is not None:
707
self.static_input_scales = not self.input_quant.get("is_dynamic")
vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py
+117
-3
@@ -389,9 +389,9 @@ def prepare_moe_fp4_layer_for_marlin(
389
390
group_size = 16 if is_nvfp4 else 32
391
392
- e = layer.num_experts
393
- k = layer.hidden_size
394
- n = layer.intermediate_size_per_partition
392
+ e = layer.moe_config.num_experts
393
+ k = layer.moe_config.hidden_dim
394
+ n = layer.moe_config.intermediate_size_per_partition
395
396
# WORKSPACE
397
device = layer.w13_weight.device
@@ -500,6 +500,120 @@ def prepare_moe_fp4_layer_for_marlin(
500
setattr(layer, name, bias)
501
502
503
+def prepare_moe_mxfp4_layer_for_marlin(
504
+ layer: torch.nn.Module,
505
+ w13: torch.Tensor,
506
+ w2: torch.Tensor,
507
+ w13_scale: torch.Tensor,
508
+ w2_scale: torch.Tensor,
509
+ w13_bias: torch.Tensor | None,
510
+ w2_bias: torch.Tensor | None,
511
+) -> tuple[
512
+ torch.Tensor,
513
+ torch.Tensor,
514
+ torch.Tensor,
515
+ torch.Tensor,
516
+ torch.Tensor | None,
517
+ torch.Tensor | None,
518
+]:
519
+ """Pure-function version of prepare_moe_fp4_layer_for_marlin for MXFP4.
520
+
521
+ Takes weight tensors as inputs and returns transformed tensors.
522
+ Does NOT modify the layer in-place.
523
+ """
524
+ input_dtype = get_marlin_input_dtype()
525
+ if (
526
+ input_dtype is not None
527
+ and input_dtype.itemsize == 1
528
+ and input_dtype != torch.float8_e4m3fn
529
+ ):
530
+ raise RuntimeError("MXFP4 weight + INT8 activation is not supported.")
531
+
532
+ group_size = 32 # MXFP4 block size
533
+
534
+ # Derive dimensions from actual weight shapes to handle rounded/padded
535
+ # sizes correctly (e.g., Mxfp4MoEMethod rounds up hidden_dim).
536
+ # w13 shape: (E, 2*N, K//2)
537
+ e = w13.shape[0]
538
+ n = w13.shape[1] // 2 # intermediate_size_per_partition
539
+ k = w13.shape[2] * 2 # hidden_size
540
+
541
+ device = w13.device
542
+ param_dtype = layer.params_dtype
543
+ is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
544
+ perm = torch.empty(0, dtype=torch.int, device=device)
545
+
546
+ # WEIGHT: Repack weights to marlin format
547
+ def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor:
548
+ tensor_list = []
549
+ if "w13" in name:
550
+ size_n, size_k = n * 2, k
551
+ else:
552
+ size_n, size_k = k, n
553
+
554
+ assert weight.shape == (e, size_n, size_k // 2)
555
+
556
+ for i in range(e):
557
+ qweight = weight[i].view(torch.int32).T.contiguous()
558
+ marlin_qweight = ops.gptq_marlin_repack(
559
+ b_q_weight=qweight,
560
+ perm=perm,
561
+ size_k=size_k,
562
+ size_n=size_n,
563
+ num_bits=4,
564
+ is_a_8bit=is_a_8bit,
565
+ )
566
+ tensor_list.append(marlin_qweight)
567
+ return torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
568
+
569
+ w13 = repack_weight(w13, "w13")
570
+ w2 = repack_weight(w2, "w2")
571
+
572
+ # WEIGHT SCALES: Permute scales
573
+ def permute_scales(scales: torch.Tensor, name: str) -> torch.Tensor:
574
+ scales = scales.view(torch.float8_e8m0fnu)
575
+ scales = scales.to(param_dtype)
576
+
577
+ tensor_list = []
578
+ if "w13" in name:
579
+ size_n, size_k = n * 2, k
580
+ else:
581
+ size_n, size_k = k, n
582
+
583
+ for i in range(e):
584
+ scale = scales[i].T
585
+ marlin_scales = marlin_permute_scales(
586
+ s=scale,
587
+ size_k=size_k,
588
+ size_n=size_n,
589
+ group_size=group_size,
590
+ is_a_8bit=is_a_8bit,
591
+ )
592
+ marlin_scales = mxfp4_marlin_process_scales(
593
+ marlin_scales, input_dtype=input_dtype
594
+ )
595
+ tensor_list.append(marlin_scales)
596
+ return torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
597
+
598
+ w13_scale = permute_scales(w13_scale, "w13")
599
+ w2_scale = permute_scales(w2_scale, "w2")
600
+
601
+ # BIAS: Permute bias
602
+ def permute_bias(bias: torch.Tensor | None) -> torch.Tensor | None:
603
+ if bias is None:
604
+ return None
605
+ bias = bias.to(param_dtype)
606
+ tensor_list = []
607
+ for i in range(e):
608
+ tensor_list.append(marlin_permute_bias(bias[i]))
609
+ return torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
610
+
611
+ w13_bias = permute_bias(w13_bias)
612
+ w2_bias = permute_bias(w2_bias)
613
+
614
+ return w13, w2, w13_scale, w2_scale, w13_bias, w2_bias
615
+
616
+
617
def rand_marlin_weight_nvfp4_like(weight, group_size, input_dtype=None):
618
is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
619
vllm/model_executor/layers/quantization/utils/mxfp4_utils.py
+1
-32
@@ -1,12 +1,10 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from collections.abc import Callable
3
from typing import Any
4
5
import torch
6
7
from vllm.logger import init_logger
9
-from vllm.model_executor.layers.fused_moe.activation import MoEActivation
8
from vllm.platforms import current_platform
9
from vllm.triton_utils import triton
10
from vllm.utils.import_utils import has_triton_kernels
@@ -22,7 +20,7 @@ logger = init_logger(__name__)
20
CK_MXFP4_MOE_DIM_ALIGNMENT = 256
21
22
25
-def _swizzle_mxfp4(quant_tensor, scale, num_warps):
23
+def _swizzle_mxfp4(quant_tensor, scale, num_warps=8):
24
"""weight swizzle for mxfp4 moe, used for OAI mxfp4 kernel"""
25
assert has_triton_kernels()
26
import triton_kernels.matmul_ogs_details.opt_flags as opt_flags
@@ -87,35 +85,6 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps):
85
return quant_tensor, InFlexData(), scale
86
87
90
-def _can_support_mxfp4(
91
- use_grouped_topk: bool = False,
92
- topk_group: int | None = None,
93
- num_expert_group: int | None = None,
94
- expert_map: torch.Tensor | None = None,
95
- custom_routing_function: Callable | None = None,
96
- e_score_correction_bias: torch.Tensor | None = None,
97
- apply_router_weight_on_input: bool = False,
98
- scoring_func: str = "softmax",
99
- activation: MoEActivation = MoEActivation.SWIGLUOAI,
100
- expert_load_view: torch.Tensor | None = None,
101
- logical_to_physical_map: torch.Tensor | None = None,
102
- logical_replica_count: torch.Tensor | None = None,
103
-):
104
- return not (
105
- use_grouped_topk
106
- or topk_group
107
- or num_expert_group
108
- or custom_routing_function
109
- or e_score_correction_bias
110
- or apply_router_weight_on_input
111
- or scoring_func != "softmax"
112
- or activation != MoEActivation.SWIGLUOAI
113
- or expert_load_view
114
- or logical_to_physical_map
115
- or logical_replica_count
116
- )
117
-
118
-
88
def get_padding_alignment():
89
return (
90
256