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 14, 05:30 UTC
[GPTOSS][DP/EP][Marlin] Enable GPTOSS DP/EP using Marlin kernels (#25488)
Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: mgoin <mgoin64@gmail.com>
Varun Sundar Rabindranath committed
Oct 3, 2025 at 20:13 UTC
7ef40bb9832f4a8fca9f9924a35ae77a69ee7076
9 files changed
+264
-101
docs/design/moe_kernel_features.md
+4
-2
@@ -93,6 +93,8 @@ To be used with a particular `FusedMoEPrepareAndFinalize` sub-class, MoE kernels
93
| 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] |
94
| deep gemm+triton<sup>2</sup> | standard,</br>batched | all<sup>1</sup> | G(128),A,T | silu, gelu | <sup>6</sup> | Y | [`TritonOrDeepGemmExperts`][vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe.TritonOrDeepGemmExperts],</br>[`BatchedTritonOrDeepGemmExperts`][vllm.model_executor.layers.fused_moe.batched_triton_or_deep_gemm_moe.BatchedTritonOrDeepGemmExperts] |
95
| marlin | standard | <sup>3</sup> | <sup>3</sup> | silu,</br>swigluoai | Y | N | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe] |
96
+
97
+| marlin experts | standard | N/A | N/A | silu,</br>swigluoai | Y | Y | [`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts] |
98
| trtllm | standard | mxfp4,</br>nvfp4 | G(16),G(32) | <sup>5</sup> | N | Y | [`TrtLlmGenExperts`][vllm.model_executor.layers.fused_moe.trtllm_moe.TrtLlmGenExperts] |
99
| pallas | standard | N/A | N/A | silu | N | N | [`fused_moe`][vllm.model_executor.layers.fused_moe.moe_pallas.fused_moe] |
100
| iterative | standard | N/A | N/A | silu | N | N | [`fused_moe`][vllm.model_executor.layers.fused_moe.moe_torch_iterative.fused_moe] |
@@ -114,6 +116,6 @@ The following table shows "families" of modular kernels that are intended to wor
116
117
| backend | `FusedMoEPrepareAndFinalize` subclasses | `FusedMoEPermuteExpertsUnpermute` subclasses |
118
|----------------------------------|------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------|
117
-| deepep_high_throughput,</br>pplx | `DeepEPHTPrepareAndFinalize`,</br>`PplxPrepareAndFinalize` | `BatchedDeepGemmExperts`,</br>`BatchedTritonExperts`,</br>`BatchedTritonOrDeepGemmExperts`,</br>`CutlassBatchedExpertsFp8` |
118
-| deepep_low_latency | `DeepEPLLPrepareAndFinalize` | `DeepGemmExperts`,</br>`TritonExperts`,</br>`TritonOrDeepGemmExperts`,</br>`CutlassExpertsFp8` |
119
+| deepep_high_throughput | `DeepEPHTPrepareAndFinalize` | `DeepGemmExperts`,</br>`TritonExperts`,</br>`TritonOrDeepGemmExperts`,</br>`CutlassExpertsFp8`, </br>`MarlinExperts` |
120
+| deepep_low_latency,</br>pplx | `DeepEPLLPrepareAndFinalize`,</br>`PplxPrepareAndFinalize` | `BatchedDeepGemmExperts`,</br>`BatchedTritonExperts`,</br>`BatchedTritonOrDeepGemmExperts`,</br>`CutlassBatchedExpertsFp8`|
121
| flashinfer | `FlashInferCutlassMoEPrepareAndFinalize` | `FlashInferExperts` |
vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py
+1
-1
@@ -303,7 +303,7 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
303
304
assert w2.size(1) == K
305
306
- E, max_num_tokens, N, K, top_k_num = mk._moe_problem_size(
306
+ E, max_num_tokens, N, K, top_k_num = self.moe_problem_size(
307
hidden_states, w1, w2, topk_ids)
308
309
workspace1 = _resize_cache(workspace13, (E, max_num_tokens, N))
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+1
-1
@@ -712,7 +712,7 @@ class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):
712
expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
713
apply_router_weight_on_input: bool,
714
):
715
- e, m, n, k, _ = mk._moe_problem_size(hidden_states, w1, w2, topk_ids)
715
+ e, m, n, k, _ = self.moe_problem_size(hidden_states, w1, w2, topk_ids)
716
n = w2.shape[2] * 2
717
718
run_cutlass_moe_fp4(
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
+1
-1
@@ -906,7 +906,7 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
906
907
expert_num_tokens = expert_tokens_meta.expert_num_tokens
908
909
- E, max_num_tokens, N, K, top_k_num = mk._moe_problem_size(
909
+ E, max_num_tokens, N, K, top_k_num = self.moe_problem_size(
910
hidden_states, w1, w2, topk_ids)
911
912
assert w1.size(0) == E
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
+164
-26
@@ -4,11 +4,18 @@
4
from typing import Optional
5
6
import torch
7
+from typing_extensions import override
8
9
import vllm._custom_ops as ops
10
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
12
from vllm.model_executor.layers.fused_moe.fused_moe import moe_align_block_size
13
+from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
14
+ TopKWeightAndReduceNoOP)
15
+from vllm.model_executor.layers.fused_moe.utils import _resize_cache
16
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
11
- marlin_make_workspace_new, maybe_warn_marlin_atomic_add)
17
+ marlin_make_workspace_new, marlin_moe_intermediate_size,
18
+ maybe_warn_marlin_atomic_add)
19
from vllm.scalar_type import ScalarType, scalar_types
20
from vllm.utils import direct_register_custom_op
21
@@ -20,7 +27,7 @@ def fused_marlin_moe(hidden_states: torch.Tensor,
27
bias2: Optional[torch.Tensor],
28
w1_scale: torch.Tensor,
29
w2_scale: torch.Tensor,
23
- gating_output: torch.Tensor,
30
+ gating_output: Optional[torch.Tensor],
31
topk_weights: torch.Tensor,
32
topk_ids: torch.Tensor,
33
quant_type_id: int,
@@ -37,7 +44,10 @@ def fused_marlin_moe(hidden_states: torch.Tensor,
44
w1_zeros: Optional[torch.Tensor] = None,
45
w2_zeros: Optional[torch.Tensor] = None,
46
workspace: Optional[torch.Tensor] = None,
47
+ intermediate_cache13: Optional[torch.Tensor] = None,
48
+ intermediate_cache2: Optional[torch.Tensor] = None,
49
is_k_full: bool = True,
50
+ output: Optional[torch.Tensor] = None,
51
inplace: bool = False) -> torch.Tensor:
52
"""
53
This function computes a Mixture of Experts (MoE) layer using two sets of
@@ -49,8 +59,8 @@ def fused_marlin_moe(hidden_states: torch.Tensor,
59
- w2 (torch.Tensor): The second set of expert weights.
60
- w1_scale (torch.Tensor): Scale to be used for w1.
61
- w2_scale (torch.Tensor): Scale to be used for w2.
52
- - gating_output (torch.Tensor): The output of the gating operation
53
- (before softmax).
62
+ - gating_output (Optional[torch.Tensor]): The output of the gating
63
+ operation (before softmax).
64
- g_idx1 (Optional[torch.Tensor]): The first set of act_order indices.
65
- g_idx2 (Optional[torch.Tensor]): The second set of act_order indices.
66
- sort_indices1 (Optional[torch.Tensor]): The first act_order input
@@ -78,8 +88,9 @@ def fused_marlin_moe(hidden_states: torch.Tensor,
88
num_bits = 4 if quant_type in bit4_scalar_types else 8
89
90
# Check constraints.
81
- assert hidden_states.shape[0] == gating_output.shape[
82
- 0], "Number of tokens mismatch"
91
+ if gating_output is not None:
92
+ assert hidden_states.shape[0] == gating_output.shape[
93
+ 0], "Number of tokens mismatch"
94
assert hidden_states.shape[
95
1] == w1.shape[1] * 16, "Hidden size mismatch w1"
96
assert hidden_states.shape[1] == w2.shape[2] // (
@@ -93,7 +104,7 @@ def fused_marlin_moe(hidden_states: torch.Tensor,
104
105
M, K = hidden_states.shape
106
E = w1.shape[0]
96
- N = w2.shape[1] * 16
107
+ N = marlin_moe_intermediate_size(w1, w2)
108
topk = topk_ids.shape[1]
109
110
# M block size selection logic
@@ -111,20 +122,24 @@ def fused_marlin_moe(hidden_states: torch.Tensor,
122
if workspace is None:
123
workspace = marlin_make_workspace_new(hidden_states.device, 4)
124
114
- intermediate_cache2 = torch.empty(
115
- (M * topk_ids.shape[1], N),
116
- device=hidden_states.device,
117
- dtype=hidden_states.dtype,
118
- )
119
- intermediate_cache13 = torch.empty(
120
- (M * topk_ids.shape[1] * max(2 * N, K), ),
121
- device=hidden_states.device,
122
- dtype=hidden_states.dtype,
123
- )
124
- intermediate_cache1 = intermediate_cache13[:M * topk_ids.shape[1] * 2 * N]
125
- intermediate_cache1 = intermediate_cache1.view(-1, 2 * N)
126
- intermediate_cache3 = intermediate_cache13[:M * topk_ids.shape[1] * K]
127
- intermediate_cache3 = intermediate_cache3.view(-1, K)
125
+ if intermediate_cache2 is None:
126
+ intermediate_cache2 = torch.empty(
127
+ (M * topk, N),
128
+ device=hidden_states.device,
129
+ dtype=hidden_states.dtype,
130
+ )
131
+
132
+ if intermediate_cache13 is None:
133
+ intermediate_cache13 = torch.empty(
134
+ (M * topk * max(2 * N, K), ),
135
+ device=hidden_states.device,
136
+ dtype=hidden_states.dtype,
137
+ )
138
+
139
+ intermediate_cache1 = _resize_cache(intermediate_cache13,
140
+ (M * topk, 2 * N))
141
+ intermediate_cache3 = _resize_cache(intermediate_cache13, (M * topk, K))
142
+ intermediate_cache2 = _resize_cache(intermediate_cache2, (M * topk, N))
143
144
maybe_warn_marlin_atomic_add(hidden_states.device, hidden_states.dtype)
145
use_atomic_add = hidden_states.dtype == torch.half or \
@@ -200,10 +215,9 @@ def fused_marlin_moe(hidden_states: torch.Tensor,
215
use_fp32_reduce=True,
216
is_zp_float=False).view(-1, topk, K)
217
203
- output = hidden_states if inplace else torch.empty_like(hidden_states)
204
- return torch.sum(intermediate_cache3.view(*intermediate_cache3.shape),
205
- dim=1,
206
- out=output)
218
+ if output is None:
219
+ output = hidden_states if inplace else torch.empty_like(hidden_states)
220
+ return torch.sum(intermediate_cache3.view(-1, topk, K), dim=1, out=output)
221
222
223
def fused_marlin_moe_fake(hidden_states: torch.Tensor,
@@ -211,7 +225,7 @@ def fused_marlin_moe_fake(hidden_states: torch.Tensor,
225
w2: torch.Tensor,
226
w1_scale: torch.Tensor,
227
w2_scale: torch.Tensor,
214
- gating_output: torch.Tensor,
228
+ gating_output: Optional[torch.Tensor],
229
topk_weights: torch.Tensor,
230
topk_ids: torch.Tensor,
231
quant_type_id: int,
@@ -227,7 +241,10 @@ def fused_marlin_moe_fake(hidden_states: torch.Tensor,
241
w1_zeros: Optional[torch.Tensor] = None,
242
w2_zeros: Optional[torch.Tensor] = None,
243
workspace: Optional[torch.Tensor] = None,
244
+ intermediate_cache13: Optional[torch.Tensor] = None,
245
+ intermediate_cache2: Optional[torch.Tensor] = None,
246
is_k_full: bool = True,
247
+ output: Optional[torch.Tensor] = None,
248
inplace: bool = False) -> torch.Tensor:
249
return torch.empty_like(hidden_states)
250
@@ -237,3 +254,124 @@ direct_register_custom_op(
254
op_func=fused_marlin_moe,
255
fake_impl=fused_marlin_moe_fake,
256
)
257
+
258
+
259
+class MarlinExperts(mk.FusedMoEPermuteExpertsUnpermute):
260
+
261
+ def __init__(self, quant_config: FusedMoEQuantConfig):
262
+ # TODO (varun) : Enable activation quantization
263
+ assert quant_config.use_mxfp4_w4a16, "Supports only mxfp4_w4a16"
264
+ super().__init__(quant_config)
265
+
266
+ @override
267
+ def moe_problem_size(
268
+ self,
269
+ a1: torch.Tensor,
270
+ w1: torch.Tensor,
271
+ w2: torch.Tensor,
272
+ topk_ids: torch.Tensor,
273
+ ) -> tuple[int, int, int, int, int]:
274
+ assert w1.dim() == 3 and w2.dim() == 3
275
+
276
+ E = w1.size(0)
277
+ K = a1.size(-1)
278
+ N = marlin_moe_intermediate_size(w1, w2)
279
+
280
+ if a1.dim() == 2:
281
+ # Make sure we are using the correct a1 (pre-permute).
282
+ assert topk_ids.size(0) == a1.size(0), \
283
+ f"{topk_ids.size(0)} != {a1.size(0)}"
284
+ M = a1.size(0)
285
+ else:
286
+ assert a1.dim() == 3
287
+ assert a1.size(0) == E, f"{a1.size(0)} == {E}"
288
+ M = a1.size(1) # This is max_num_tokens
289
+
290
+ assert topk_ids.dim() == 2
291
+ topk = topk_ids.size(1)
292
+
293
+ return E, M, N, K, topk
294
+
295
+ def supports_expert_map(self) -> bool:
296
+ return True
297
+
298
+ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
299
+ return TopKWeightAndReduceNoOP()
300
+
301
+ @property
302
+ def activation_formats(
303
+ self
304
+ ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
305
+ return (mk.FusedMoEActivationFormat.Standard,
306
+ mk.FusedMoEActivationFormat.Standard)
307
+
308
+ def supports_chunking(self) -> bool:
309
+ return True
310
+
311
+ def workspace_shapes(
312
+ self, a: torch.Tensor, aq: torch.Tensor, M: int, N: int, K: int,
313
+ topk: int, global_num_experts: int, local_num_experts: int,
314
+ expert_tokens_meta: Optional[mk.ExpertTokensMetadata]
315
+ ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
316
+ # Modular Kernel provisions output buffer from workspace1. However in
317
+ # the fused_marlin_moe() function, the final torch.sum(), is defined
318
+ # essentially as,
319
+ # `torch.sum(workspace1, dim=1, out=output)`
320
+ # Having overlapping input and output tensors for torch.sum seems
321
+ # error prone and depends on how the torch.sum is implemented.
322
+ # For this reason we swap let the output buffer provision from
323
+ # workspace2.
324
+
325
+ # Workspace/IntermediateCache allocation matching fused_marlin_moe()
326
+ #workspace1 = (M * topk * max(2 * N, K),)
327
+ #workspace2 = (M * topk, N)
328
+
329
+ # Workspace/IntermediateCache allocation accounting for output buffer
330
+ # provisioning
331
+ workspace1 = (M * topk, max(N, K))
332
+ workspace2 = (M * topk * max(2 * N, K), )
333
+ output = (M, K)
334
+
335
+ return (workspace1, workspace2, output, a.dtype)
336
+
337
+ def apply(
338
+ self,
339
+ output: torch.Tensor,
340
+ hidden_states: torch.Tensor,
341
+ w1: torch.Tensor,
342
+ w2: torch.Tensor,
343
+ topk_weights: torch.Tensor,
344
+ topk_ids: torch.Tensor,
345
+ activation: str,
346
+ global_num_experts: int,
347
+ expert_map: Optional[torch.Tensor],
348
+ a1q_scale: Optional[torch.Tensor],
349
+ a2_scale: Optional[torch.Tensor],
350
+ workspace13: torch.Tensor,
351
+ workspace2: torch.Tensor,
352
+ expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
353
+ apply_router_weight_on_input: bool,
354
+ ):
355
+ assert self.w1_scale is not None
356
+ assert self.w2_scale is not None
357
+ return fused_marlin_moe(
358
+ hidden_states=hidden_states,
359
+ w1=w1,
360
+ w2=w2,
361
+ bias1=self.w1_bias,
362
+ bias2=self.w2_bias,
363
+ w1_scale=self.w1_scale,
364
+ w2_scale=self.w2_scale,
365
+ gating_output=None,
366
+ topk_weights=topk_weights,
367
+ topk_ids=topk_ids,
368
+ quant_type_id=scalar_types.float4_e2m1f.id, # works only for w4a16
369
+ apply_router_weight_on_input=apply_router_weight_on_input,
370
+ global_num_experts=global_num_experts,
371
+ activation=activation,
372
+ expert_map=expert_map,
373
+ output=output,
374
+ # Workspaces are swapped in workspace_shapes() to account for proper
375
+ # output buffer allocation. Please refer to workspace_shapes().
376
+ intermediate_cache13=workspace2,
377
+ intermediate_cache2=workspace13)
vllm/model_executor/layers/fused_moe/fused_moe.py
+1
-1
@@ -1780,7 +1780,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1780
torch.float32, torch.float16, torch.bfloat16, torch.float8_e4m3fn
1781
]
1782
1783
- E, num_tokens, N, K, top_k_num = mk._moe_problem_size(
1783
+ E, num_tokens, N, K, top_k_num = self.moe_problem_size(
1784
hidden_states, w1, w2, topk_ids)
1785
1786
if global_num_experts == -1:
vllm/model_executor/layers/fused_moe/modular_kernel.py
+48
-42
@@ -55,46 +55,6 @@ from vllm.v1.worker.ubatching import (dbo_current_ubatch_id, dbo_enabled,
55
#
56
57
58
-def _moe_problem_size(
59
- a1: torch.Tensor,
60
- w1: torch.Tensor,
61
- w2: torch.Tensor,
62
- topk_ids: torch.Tensor,
63
-) -> tuple[int, int, int, int, int]:
64
- """
65
- Extract the MoE problem size from the given tensor arguments:
66
- - a: The hidden states, input to the MoE layer.
67
- - w1: The first set of expert weights.
68
- - w2: The second set of expert weights.
69
- - topk_ids: The topk ids.
70
-
71
- Note: extracting the problem shape from the weight and activation tensors is
72
- not obvious. It needs to be done this way specifically due to subtle issues
73
- with particular kernels, e.g. the int4 kernels divide the trailing dimension
74
- by two, so it's not "correct" to extract N or K from the trailing dimension
75
- of w1 or w2. Similarly, some kernels transpose the weights, so this needs
76
- to be kept in mind.
77
- """
78
- assert w1.dim() == 3 and w2.dim() == 3
79
- E, N, _ = w1.size()
80
- K = a1.size(-1)
81
-
82
- if a1.dim() == 2:
83
- # Make sure we are using the correct a1 (pre-permute).
84
- assert topk_ids.size(0) == a1.size(0), \
85
- f"{topk_ids.size(0)} != {a1.size(0)}"
86
- M = a1.size(0)
87
- else:
88
- assert a1.dim() == 3
89
- assert a1.size(0) == E, f"{a1.size(0)} == {E}"
90
- M = a1.size(1) # This is max_num_tokens
91
-
92
- assert topk_ids.dim() == 2
93
- topk = topk_ids.size(1)
94
-
95
- return E, M, N, K, topk
96
-
97
-
58
class FusedMoEActivationFormat(Enum):
59
"""
60
The standard activation format (num_tokens, hidden dim).
@@ -391,6 +351,50 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
351
"""
352
raise NotImplementedError
353
354
+ def moe_problem_size(
355
+ self,
356
+ a1: torch.Tensor,
357
+ w1: torch.Tensor,
358
+ w2: torch.Tensor,
359
+ topk_ids: torch.Tensor,
360
+ ) -> tuple[int, int, int, int, int]:
361
+ """
362
+ Extract the MoE problem size from the given tensor arguments:
363
+ - a: The hidden states, input to the MoE layer.
364
+ - w1: The first set of expert weights.
365
+ - w2: The second set of expert weights.
366
+ - topk_ids: The topk ids.
367
+
368
+ Note: extracting the problem shape from the weight and activation
369
+ tensors is not obvious. It needs to be done this way specifically
370
+ due to subtle issues with particular kernels, e.g. the int4 kernels
371
+ divide the trailing dimension by two, so it's not "correct" to
372
+ extract N or K from the trailing dimension of w1 or w2. Similarly,
373
+ some kernels transpose the weights, so this needs to be kept in mind.
374
+
375
+ Note: This implementation covers most cases. However, if experts
376
+ require a specialized implementation, like MarlinExperts, they are free
377
+ to override this function.
378
+ """
379
+ assert w1.dim() == 3 and w2.dim() == 3
380
+ E, N, _ = w1.size()
381
+ K = a1.size(-1)
382
+
383
+ if a1.dim() == 2:
384
+ # Make sure we are using the correct a1 (pre-permute).
385
+ assert topk_ids.size(0) == a1.size(0), \
386
+ f"{topk_ids.size(0)} != {a1.size(0)}"
387
+ M = a1.size(0)
388
+ else:
389
+ assert a1.dim() == 3
390
+ assert a1.size(0) == E, f"{a1.size(0)} == {E}"
391
+ M = a1.size(1) # This is max_num_tokens
392
+
393
+ assert topk_ids.dim() == 2
394
+ topk = topk_ids.size(1)
395
+
396
+ return E, M, N, K, topk
397
+
398
#
399
# Various helpers for accessing quantization parameters from the
400
# quant_config.
@@ -674,7 +678,8 @@ class FusedMoEModularKernel(torch.nn.Module):
678
apply_router_weight_on_input: bool,
679
) -> torch.Tensor:
680
677
- _, M, N, K, top_k = _moe_problem_size(a1q, w1, w2, topk_ids)
681
+ _, M, N, K, top_k = self.fused_experts.moe_problem_size(
682
+ a1q, w1, w2, topk_ids)
683
684
(workspace13_shape, workspace2_shape, fused_out_shape,
685
workspace_dtype) = self.fused_experts.workspace_shapes(
@@ -737,7 +742,8 @@ class FusedMoEModularKernel(torch.nn.Module):
742
apply_router_weight_on_input: bool,
743
) -> torch.Tensor:
744
740
- _, M, N, K, top_k = _moe_problem_size(a1q, w1, w2, topk_ids)
745
+ _, M, N, K, top_k = self.fused_experts.moe_problem_size(
746
+ a1q, w1, w2, topk_ids)
747
748
CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
749
num_chunks = cdiv(M, CHUNK_SIZE)
vllm/model_executor/layers/quantization/mxfp4.py
+34
-27
@@ -15,6 +15,7 @@ from vllm.model_executor.layers.fused_moe import modular_kernel as mk
15
from vllm.model_executor.layers.fused_moe.config import (
16
FusedMoEQuantConfig, mxfp4_w4a4_moe_quant_config,
17
mxfp4_w4a16_moe_quant_config)
18
+from vllm.model_executor.layers.fused_moe.fused_marlin_moe import MarlinExperts
19
from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
20
OAITritonExperts)
21
from vllm.model_executor.layers.fused_moe.trtllm_moe import TrtLlmGenExperts
@@ -92,7 +93,7 @@ def get_mxfp4_backend():
93
"Please `pip install vllm[flashinfer]` for best results.")
94
95
# If FlashInfer is not available, try either Marlin or Triton
95
- if current_platform.get_device_capability(
96
+ if envs.VLLM_MXFP4_USE_MARLIN or current_platform.get_device_capability(
97
)[0] < 9 or not has_triton_kernels() or not is_torch_equal_or_newer(
98
"2.8.0"):
99
logger.info_once("Using Marlin backend")
@@ -646,9 +647,13 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
647
self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:
648
649
if self.mxfp4_backend == Mxfp4Backend.MARLIN:
649
- return None
650
-
651
- if self.mxfp4_backend == Mxfp4Backend.TRITON:
650
+ return mxfp4_w4a16_moe_quant_config(
651
+ w1_bias=layer.w13_bias,
652
+ w2_bias=layer.w2_bias,
653
+ w1_scale=layer.w13_weight_scale,
654
+ w2_scale=layer.w2_weight_scale,
655
+ )
656
+ elif self.mxfp4_backend == Mxfp4Backend.TRITON:
657
w1_scale = self.w13_precision_config
658
w2_scale = self.w2_precision_config
659
return mxfp4_w4a16_moe_quant_config(
@@ -690,6 +695,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
695
}
696
return TrtLlmGenExperts(self.moe, self.moe_quant_config,
697
**kwargs)
698
+ elif (self.mxfp4_backend == Mxfp4Backend.MARLIN):
699
+ return MarlinExperts(self.moe_quant_config)
700
else:
701
return OAITritonExperts(self.moe_quant_config)
702
@@ -782,6 +789,29 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
789
if enable_eplb:
790
raise NotImplementedError("EPLB is not supported for mxfp4")
791
792
+ if self.fused_experts is not None:
793
+ return self._route_and_experts(
794
+ layer,
795
+ x,
796
+ router_logits,
797
+ top_k,
798
+ renormalize,
799
+ use_grouped_topk,
800
+ topk_group,
801
+ num_expert_group,
802
+ global_num_experts,
803
+ expert_map,
804
+ custom_routing_function,
805
+ scoring_func,
806
+ e_score_correction_bias,
807
+ apply_router_weight_on_input,
808
+ activation,
809
+ enable_eplb,
810
+ expert_load_view,
811
+ logical_to_physical_map,
812
+ logical_replica_count,
813
+ )
814
+
815
if self.mxfp4_backend == Mxfp4Backend.MARLIN:
816
topk_weights, topk_ids, _ = FusedMoE.select_experts(
817
hidden_states=x,
@@ -815,29 +845,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
845
activation=activation,
846
expert_map=expert_map)
847
818
- if self.fused_experts is not None:
819
- return self._route_and_experts(
820
- layer,
821
- x,
822
- router_logits,
823
- top_k,
824
- renormalize,
825
- use_grouped_topk,
826
- topk_group,
827
- num_expert_group,
828
- global_num_experts,
829
- expert_map,
830
- custom_routing_function,
831
- scoring_func,
832
- e_score_correction_bias,
833
- apply_router_weight_on_input,
834
- activation,
835
- enable_eplb,
836
- expert_load_view,
837
- logical_to_physical_map,
838
- logical_replica_count,
839
- )
840
-
848
assert _can_support_mxfp4(
849
use_grouped_topk, topk_group, num_expert_group, expert_map,
850
custom_routing_function, e_score_correction_bias,
vllm/model_executor/layers/quantization/utils/marlin_utils.py
+10
@@ -187,6 +187,16 @@ def check_moe_marlin_supports_layer(layer: LinearBase, group_size: int) \
187
supports_router_weight and supports_activation
188
189
190
+def marlin_moe_intermediate_size(w1_packed: torch.Tensor,
191
+ w2_packed: torch.Tensor):
192
+ """
193
+ Given Marlin packed weight matrices w1_packed, and w2_packed,
194
+ return the MoE intermediate size N
195
+ """
196
+ marlin_tile_size = 16
197
+ return w2_packed.size(1) * marlin_tile_size
198
+
199
+
200
def marlin_make_workspace(output_size_per_partition: int,
201
device: torch.device) -> torch.Tensor:
202
max_workspace_size = (output_size_per_partition //