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 12, 09:45 UTC
Revert "[LoRA] Support FusedMoE LoRA Triton kernel for mxfp4 (#28971)" (#29697)
Signed-off-by: Huamin Li <3ericli@gmail.com>
Huamin Li committed
Nov 28, 2025 at 15:26 UTC
3fd1fb0b6016d1471853f5114fd97c74f1a8d29c
4 files changed
+11
-440
tests/kernels/moe/test_modular_oai_triton_moe.py
deleted
-250
@@ -1,250 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""
4
-Test modular OAI Triton MoE
5
-"""
6
-
7
-import pytest
8
-import torch
9
-
10
-from vllm.utils.import_utils import has_triton_kernels
11
-
12
-if not has_triton_kernels():
13
- pytest.skip(
14
- "triton_kernels not found, skipping all related tests",
15
- allow_module_level=True,
16
- )
17
-
18
-from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
19
-from triton_kernels.numerics import InFlexData
20
-from triton_kernels.numerics_details.mxfp import downcast_to_mxfp, upcast_from_mxfp
21
-from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor
22
-from triton_kernels.tensor_details import layout
23
-from triton_kernels.testing import assert_close
24
-
25
-from vllm.config import VllmConfig, set_current_vllm_config
26
-from vllm.model_executor.layers.fused_moe.config import mxfp4_w4a16_moe_quant_config
27
-from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
28
- OAITritonExperts,
29
- UnfusedOAITritonExperts,
30
-)
31
-from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel
32
-from vllm.model_executor.layers.fused_moe.prepare_finalize import (
33
- MoEPrepareAndFinalizeNoEP,
34
-)
35
-from vllm.model_executor.layers.utils import shuffle_weight
36
-from vllm.platforms import current_platform
37
-
38
-MNK = [
39
- (1, 512, 384),
40
- (1, 2880, 2880),
41
- (2, 512, 384),
42
- (2, 2880, 2880),
43
- (32, 2880, 2880),
44
- (64, 2880, 2880),
45
-]
46
-
47
-
48
-def unshuffle_weight(w: torch.Tensor):
49
- first = w[..., ::2]
50
- second = w[..., 1::2]
51
- return torch.concat((first, second), dim=-1)
52
-
53
-
54
-def make_weights(dtype, k, n, e):
55
- w1 = torch.randn((e, k, 2 * n), dtype=dtype, device="cuda")
56
- w1_bias = torch.randn((e, 2 * n), dtype=dtype, device="cuda")
57
-
58
- w2 = torch.randn((e, n, k), dtype=dtype, device="cuda")
59
- w2_bias = torch.randn((e, k), dtype=dtype, device="cuda")
60
-
61
- w1_tri = w1.clone()
62
- w2_tri = w2.clone()
63
-
64
- w1_bias_tri = w1_bias.clone()
65
- w2_bias_tri = w2_bias.clone()
66
- w1_bias_tri = w1_bias_tri.to(torch.float32)
67
- w2_bias_tri = w2_bias_tri.to(torch.float32)
68
-
69
- # shuffle weights
70
- w1_tri = shuffle_weight(w1_tri)
71
- w1_bias_tri = shuffle_weight(w1_bias_tri)
72
-
73
- # quant triton_weights
74
- w1_tri, w1_scale_tri = downcast_to_mxfp(w1_tri, torch.uint8, axis=1)
75
- w1 = upcast_from_mxfp(w1_tri, w1_scale_tri, dtype, axis=1)
76
- w1 = unshuffle_weight(w1)
77
-
78
- w2_tri, w2_scale_tri = downcast_to_mxfp(w2_tri, torch.uint8, axis=1)
79
- w2 = upcast_from_mxfp(w2_tri, w2_scale_tri, dtype, axis=1)
80
-
81
- num_warps = 8
82
- w_layout, w_layout_opts = layout.make_default_matmul_mxfp4_w_layout(mx_axis=1)
83
- w_scale_layout, w_scale_layout_opts = (
84
- layout.make_default_matmul_mxfp4_w_scale_layout(mx_axis=1, num_warps=num_warps)
85
- )
86
-
87
- w1_tri = convert_layout(wrap_torch_tensor(w1_tri, FP4), w_layout, **w_layout_opts)
88
- w1_scale_tri = convert_layout(
89
- wrap_torch_tensor(w1_scale_tri),
90
- w_scale_layout,
91
- **w_scale_layout_opts,
92
- )
93
-
94
- w2_tri = convert_layout(wrap_torch_tensor(w2_tri, FP4), w_layout, **w_layout_opts)
95
- w2_scale_tri = convert_layout(
96
- wrap_torch_tensor(w2_scale_tri),
97
- w_scale_layout,
98
- **w_scale_layout_opts,
99
- )
100
-
101
- w1_precision_config = PrecisionConfig(
102
- weight_scale=w1_scale_tri, flex_ctx=FlexCtx(rhs_data=InFlexData())
103
- )
104
- w2_precision_config = PrecisionConfig(
105
- weight_scale=w2_scale_tri, flex_ctx=FlexCtx(rhs_data=InFlexData())
106
- )
107
-
108
- return (
109
- w1,
110
- w2,
111
- w1_bias,
112
- w2_bias,
113
- w1_tri,
114
- w2_tri,
115
- w1_bias_tri,
116
- w2_bias_tri,
117
- w1_precision_config,
118
- w2_precision_config,
119
- )
120
-
121
-
122
-def swiglu(x, alpha: float = 1.702, limit: float = 1.0):
123
- # Note we add an extra bias of 1 to the linear layer
124
- x_glu, x_linear = torch.chunk(x, 2, dim=-1)
125
- if limit is not None:
126
- x_glu = x_glu.clamp(max=limit)
127
- out_glu = x_glu * torch.sigmoid(alpha * x_glu)
128
- if limit is not None:
129
- x_linear = x_linear.clamp(min=-limit, max=limit)
130
- return out_glu * (x_linear + 1)
131
-
132
-
133
-def torch_moe_impl(
134
- hidden_states: torch.Tensor, # (M, K)
135
- w1: torch.Tensor, # (E, K, 2N)
136
- w2: torch.Tensor, # (E, N, K)
137
- w1_bias: torch.Tensor, # (E, 2N)
138
- w2_bias: torch.Tensor, # (E, K)
139
- topk_weights: torch.Tensor, # (M, topk)
140
- topk_ids: torch.Tensor, # (M, topk)
141
-):
142
- w1 = w1[topk_ids, ...]
143
- w1_bias = w1_bias[topk_ids, ...]
144
- hidden_states = torch.einsum("bekc,bk->bec", w1, hidden_states) + w1_bias
145
- hidden_states = swiglu(hidden_states, limit=7)
146
-
147
- w2 = w2[topk_ids, ...]
148
- w2_bias = w2_bias[topk_ids, ...]
149
- hidden_states = torch.einsum("bekc,bek->bec", w2, hidden_states) + w2_bias
150
-
151
- # Weighted sum of experts
152
- hidden_states = torch.einsum("bec,be->bc", hidden_states, topk_weights)
153
- return hidden_states
154
-
155
-
156
-def oai_triton_moe_impl(
157
- x: torch.Tensor,
158
- w1: torch.Tensor,
159
- w2: torch.Tensor,
160
- w1_scale: "PrecisionConfig",
161
- w2_scale: "PrecisionConfig",
162
- w1_bias: torch.Tensor | None,
163
- w2_bias: torch.Tensor | None,
164
- num_experts: int,
165
- topk_weights: torch.Tensor,
166
- topk_ids: torch.Tensor,
167
- unfused: bool = False,
168
-) -> torch.Tensor:
169
- quant_config = mxfp4_w4a16_moe_quant_config(
170
- w1_bias=w1_bias,
171
- w2_bias=w2_bias,
172
- w1_scale=w1_scale,
173
- w2_scale=w2_scale,
174
- )
175
-
176
- if unfused:
177
- fused_experts = UnfusedOAITritonExperts(quant_config)
178
- else:
179
- fused_experts = OAITritonExperts(quant_config)
180
-
181
- mk = FusedMoEModularKernel(MoEPrepareAndFinalizeNoEP(), fused_experts)
182
-
183
- return mk.forward(
184
- hidden_states=x,
185
- w1=w1,
186
- w2=w2,
187
- topk_weights=topk_weights,
188
- topk_ids=topk_ids,
189
- inplace=True,
190
- activation="swigluoai",
191
- global_num_experts=num_experts,
192
- expert_map=None,
193
- apply_router_weight_on_input=False,
194
- )
195
-
196
-
197
-@pytest.mark.skipif(
198
- not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
199
-)
200
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
201
-@pytest.mark.parametrize("m,n,k", MNK)
202
-@pytest.mark.parametrize("num_experts", [32, 128])
203
-@pytest.mark.parametrize("topk", [4])
204
-@pytest.mark.parametrize("unfused", [True, False])
205
-def test_oai_triton_moe(
206
- dtype: torch.dtype,
207
- m: int,
208
- n: int,
209
- k: int,
210
- num_experts: int,
211
- topk: int,
212
- unfused: bool,
213
-):
214
- current_platform.seed_everything(0)
215
- (
216
- w1,
217
- w2,
218
- w1_bias,
219
- w2_bias,
220
- w1_tri,
221
- w2_tri,
222
- w1_bias_tri,
223
- w2_bias_tri,
224
- w1_precision_config,
225
- w2_precision_config,
226
- ) = make_weights(dtype, k, n, num_experts)
227
-
228
- x = torch.randn((m, k), dtype=dtype, device="cuda")
229
- router_logits = torch.randn(m, num_experts, device="cuda", dtype=dtype)
230
- topk_weights, topk_ids = torch.topk(router_logits, k=topk, dim=-1, sorted=True)
231
- topk_weights = torch.nn.functional.softmax(topk_weights, dim=-1)
232
-
233
- with set_current_vllm_config(VllmConfig()):
234
- out_ref = torch_moe_impl(x, w1, w2, w1_bias, w2_bias, topk_weights, topk_ids)
235
-
236
- out = oai_triton_moe_impl(
237
- x,
238
- w1_tri,
239
- w2_tri,
240
- w1_precision_config,
241
- w2_precision_config,
242
- w1_bias_tri,
243
- w2_bias_tri,
244
- num_experts,
245
- topk_weights,
246
- topk_ids,
247
- unfused,
248
- )
249
-
250
- assert_close(ref=out_ref, tri=out, maxtol=0.025, rmstol=0.005)
vllm/lora/layers/fused_moe.py
+9
-26
@@ -20,24 +20,15 @@ from vllm.model_executor.layers.fused_moe.config import (
20
_get_config_dtype_str,
21
)
22
from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
23
- MarlinExperts,
23
+ modular_marlin_fused_moe,
24
)
25
from vllm.model_executor.layers.fused_moe.fused_moe import (
26
- TritonExperts,
26
+ modular_triton_fused_moe,
27
try_get_optimal_moe_config,
28
)
29
from vllm.model_executor.layers.fused_moe.fused_moe_modular_method import (
30
FusedMoEModularMethod,
31
)
32
-from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
33
- UnfusedOAITritonExperts,
34
-)
35
-from vllm.model_executor.layers.fused_moe.modular_kernel import (
36
- FusedMoEModularKernel,
37
-)
38
-from vllm.model_executor.layers.fused_moe.prepare_finalize import (
39
- MoEPrepareAndFinalizeNoEP,
40
-)
32
33
from .utils import _get_lora_device
34
@@ -123,23 +114,15 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
114
self.base_layer.ensure_moe_quant_config_init()
115
quant_config = self.base_layer.quant_method.moe_quant_config
116
126
- prepare_finalize = MoEPrepareAndFinalizeNoEP()
127
- m_fused_moe_fn = FusedMoEModularKernel(
128
- prepare_finalize,
129
- self.base_layer.quant_method.select_gemm_impl(
130
- prepare_finalize, self.base_layer
131
- ),
132
- self.base_layer.shared_experts,
133
- getattr(self.base_layer, "shared_experts_stream", None),
134
- )
135
- if quant_config.use_mxfp4_w4a16:
136
- assert isinstance(
137
- m_fused_moe_fn.fused_experts, (MarlinExperts, UnfusedOAITritonExperts)
117
+ m_fused_moe_fn = (
118
+ modular_triton_fused_moe(
119
+ quant_config, shared_experts=self.base_layer.shared_experts
120
)
139
- else:
140
- assert isinstance(
141
- m_fused_moe_fn.fused_experts, (MarlinExperts, TritonExperts)
121
+ if not quant_config.use_mxfp4_w4a16
122
+ else modular_marlin_fused_moe(
123
+ quant_config, shared_experts=self.base_layer.shared_experts
124
)
125
+ )
126
127
def fwd_decorator(layer, func):
128
def wrapper(*args, **kwargs):
vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py
-146
@@ -5,7 +5,6 @@
5
import torch
6
7
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8
-from vllm import _custom_ops as ops
8
from vllm.logger import init_logger
9
from vllm.model_executor.layers.fused_moe.config import (
10
FUSED_MOE_UNQUANTIZED_CONFIG,
@@ -377,148 +376,3 @@ class OAITritonExperts(BaseOAITritonExperts):
376
intermediate_cache=workspace2,
377
a1q_scale=a1q_scale,
378
)
380
-
381
-
382
-class UnfusedOAITritonExperts(BaseOAITritonExperts):
383
- """
384
- A Triton based MoE expert class that operates on expert standard
385
- format and explicitly keeps the activation and reduction (moe_sum) steps
386
- unfused from the matmul_ogs kernel. This exposes injection points
387
- for activation and moe_sum.
388
-
389
- One use case for it is to inject LoRA modules on the activation and moe_sum.
390
- """
391
-
392
- def __init__(self, quant_config: FusedMoEQuantConfig):
393
- # TODO (varun) : Enable activation quantization
394
- assert quant_config.use_mxfp4_w4a16, "Supports only mxfp4_w4a16"
395
- super().__init__(quant_config)
396
-
397
- @property
398
- def activation_formats(
399
- self,
400
- ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
401
- return (
402
- mk.FusedMoEActivationFormat.Standard,
403
- mk.FusedMoEActivationFormat.Standard,
404
- )
405
-
406
- def supports_chunking(self) -> bool:
407
- return True
408
-
409
- def workspace_shapes(
410
- self,
411
- M: int,
412
- N: int,
413
- K: int,
414
- topk: int,
415
- global_num_experts: int,
416
- local_num_experts: int,
417
- expert_tokens_meta: mk.ExpertTokensMetadata | None,
418
- ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
419
- # workspace are allocated inside the kernel
420
- workspace1 = (M * topk, N // 2)
421
- workspace2 = (M * topk, max(N, K))
422
- output = (M, K)
423
- return (workspace1, workspace2, output)
424
-
425
- def moe_sum(self, input: torch.Tensor, output: torch.Tensor):
426
- ops.moe_sum(input, output)
427
-
428
- def apply(
429
- self,
430
- output: torch.Tensor,
431
- hidden_states: torch.Tensor,
432
- w1: torch.Tensor,
433
- w2: torch.Tensor,
434
- topk_weights: torch.Tensor,
435
- topk_ids: torch.Tensor,
436
- activation: str,
437
- global_num_experts: int,
438
- expert_map: torch.Tensor | None,
439
- a1q_scale: torch.Tensor | None,
440
- a2_scale: torch.Tensor | None,
441
- workspace13: torch.Tensor,
442
- workspace2: torch.Tensor,
443
- expert_tokens_meta: mk.ExpertTokensMetadata | None,
444
- apply_router_weight_on_input: bool,
445
- ):
446
- if self.quant_config is None:
447
- self.quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
448
-
449
- if expert_map is not None:
450
- topk_ids = expert_map[topk_ids]
451
-
452
- local_num_experts = w1.size(0)
453
- if global_num_experts == -1:
454
- global_num_experts = local_num_experts
455
-
456
- routing_data, gather_indx, scatter_indx = self._make_routing_data(
457
- topk_ids, topk_weights, local_num_experts
458
- )
459
-
460
- topk = topk_ids.size(1)
461
-
462
- # type check, uint8 means mxfp4
463
- assert hidden_states.dtype == torch.bfloat16
464
- assert (
465
- self.quant_config.w1_bias is None
466
- or self.quant_config.w1_bias.dtype == torch.float32
467
- )
468
- assert (
469
- self.quant_config.w2_bias is None
470
- or self.quant_config.w2_bias.dtype == torch.float32
471
- )
472
-
473
- # Shape check, only check non-mxfp4
474
- assert hidden_states.ndim == 2
475
- assert hidden_states.shape[-1] == w1.shape[-2]
476
- assert w2.shape[-1] == w1.shape[1]
477
-
478
- batch_dim = 1
479
- M, K = hidden_states.shape
480
- E, _, N = w1.shape
481
-
482
- if global_num_experts == -1:
483
- global_num_experts = E
484
-
485
- # Note that the output tensor might be in workspace13
486
- intermediate_cache1 = _resize_cache(workspace2, (batch_dim, M * topk, N))
487
- intermediate_cache3 = _resize_cache(workspace2, (batch_dim, M * topk, K))
488
- intermediate_cache2 = _resize_cache(workspace13, (M * topk, N // 2))
489
-
490
- gammas = routing_data.gate_scal if routing_data else None
491
-
492
- matmul_ogs(
493
- hidden_states,
494
- w1,
495
- self.quant_config.w1_bias,
496
- routing_data,
497
- gather_indx=gather_indx,
498
- precision_config=self.quant_config.w1_precision,
499
- gammas=gammas if apply_router_weight_on_input else None,
500
- fused_activation=None,
501
- y=intermediate_cache1,
502
- )
503
-
504
- self.activation(
505
- activation, intermediate_cache2, intermediate_cache1.view(-1, N)
506
- )
507
-
508
- # matmul_ogs grouped reduction fuse sum across multiple experts:
509
- # y[dst_ind // n_expts_act, :] += x[src_ind, :]
510
- # Need to set n_expts_act to 1 to unfuse moe_sum
511
- routing_data.n_expts_act = 1
512
-
513
- matmul_ogs(
514
- intermediate_cache2,
515
- w2,
516
- self.quant_config.w2_bias,
517
- routing_data,
518
- scatter_indx=scatter_indx,
519
- precision_config=self.quant_config.w2_precision,
520
- gammas=None if apply_router_weight_on_input else gammas,
521
- y=intermediate_cache3,
522
- )
523
-
524
- self.moe_sum(intermediate_cache3.view(-1, topk, K), output)
vllm/model_executor/layers/quantization/mxfp4.py
+2
-18
@@ -30,7 +30,6 @@ from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
30
)
31
from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
32
OAITritonExperts,
33
- UnfusedOAITritonExperts,
33
)
34
from vllm.model_executor.layers.fused_moe.trtllm_moe import TrtLlmGenExperts
35
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
@@ -84,21 +83,8 @@ def get_mxfp4_backend_with_lora() -> Mxfp4Backend:
83
if not current_platform.is_cuda():
84
return Mxfp4Backend.NONE
85
87
- # If FlashInfer is not available, try either Marlin or Triton
88
- triton_kernels_supported = (
89
- has_triton_kernels()
90
- and is_torch_equal_or_newer("2.8.0")
91
- # NOTE: triton_kernels are only confirmed to work on SM90 and SM100
92
- # SM110 fails with this error: https://github.com/vllm-project/vllm/issues/29317
93
- # SM120 needs this fix: https://github.com/triton-lang/triton/pull/8498
94
- and (9, 0) <= current_platform.get_device_capability() < (11, 0)
95
- )
96
- if envs.VLLM_MXFP4_USE_MARLIN or not triton_kernels_supported:
97
- logger.info_once("[get_mxfp4_backend_with_lora] Using Marlin backend")
98
- return Mxfp4Backend.MARLIN
99
-
100
- logger.info_once("[get_mxfp4_backend_with_lora] Using Triton backend")
101
- return Mxfp4Backend.TRITON
86
+ logger.info_once("[get_mxfp4_backend_with_lora] Using Marlin backend")
87
+ return Mxfp4Backend.MARLIN
88
89
90
def get_mxfp4_backend(with_lora_support: bool) -> Mxfp4Backend:
@@ -868,8 +854,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
854
elif self.mxfp4_backend == Mxfp4Backend.MARLIN:
855
return MarlinExperts(self.moe_quant_config)
856
elif self.mxfp4_backend == Mxfp4Backend.TRITON:
871
- if self.moe.is_lora_enabled:
872
- return UnfusedOAITritonExperts(self.moe_quant_config)
857
return OAITritonExperts(self.moe_quant_config)
858
else:
859
raise NotImplementedError(