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
DP/EP Support for gpt-oss with deepep-ht comm kernel on SM100 (#23608)
Yongye Zhu committed
Aug 27, 2025 at 17:33 UTC
082cc07ef8f810bea61eaed77a60137684ca78f8
11 files changed
+365
-12
vllm/distributed/device_communicators/base_device_communicator.py
+1
-1
@@ -255,7 +255,7 @@ class DeviceCommunicatorBase:
255
if module.__class__.__name__ == "FusedMoE"
256
]
257
for module in moe_modules:
258
- module.quant_method.init_prepare_finalize()
258
+ module.quant_method.init_prepare_finalize(module)
259
260
def dispatch(
261
self, hidden_states: torch.Tensor,
vllm/model_executor/layers/fused_moe/config.py
+6
@@ -450,6 +450,12 @@ class FusedMoEConfig:
450
if quant_dtype is None and isinstance(quant_config, Fp8Config):
451
quant_dtype = torch.float8_e4m3fn
452
453
+ from vllm.model_executor.layers.quantization.mxfp4 import (
454
+ Mxfp4Config)
455
+ if (quant_dtype is None and isinstance(quant_config, Mxfp4Config)
456
+ and envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8):
457
+ quant_dtype = "mxfp8"
458
+
459
from vllm.model_executor.layers.quantization.modelopt import (
460
ModelOptNvFp4Config)
461
if quant_dtype is None and isinstance(quant_config,
vllm/model_executor/layers/fused_moe/layer.py
+4
-2
@@ -200,7 +200,7 @@ class FusedMoEMethodBase(QuantizeMethodBase):
200
201
# Note: init_prepare_finalize should only be called by
202
# prepare_communication_buffer_for_model.
203
- def init_prepare_finalize(self):
203
+ def init_prepare_finalize(self, layer: torch.nn.Module):
204
assert self.moe is not None
205
prepare_finalize = self.maybe_make_prepare_finalize(self.moe)
206
@@ -211,7 +211,7 @@ class FusedMoEMethodBase(QuantizeMethodBase):
211
assert self.fused_experts is None, \
212
f"Attempt to override experts for {id(self)}!"
213
self.topk_indices_dtype = prepare_finalize.topk_indices_dtype()
214
- experts = self.select_gemm_impl(prepare_finalize, self.moe)
214
+ experts = self.select_gemm_impl(prepare_finalize, self.moe, layer)
215
self.fused_experts = FusedMoEModularKernel(
216
prepare_finalize,
217
experts,
@@ -221,6 +221,7 @@ class FusedMoEMethodBase(QuantizeMethodBase):
221
self,
222
prepare_finalize: FusedMoEPrepareAndFinalize,
223
moe: FusedMoEConfig,
224
+ layer: torch.nn.Module,
225
) -> FusedMoEPermuteExpertsUnpermute:
226
# based on the all2all implementation, select the appropriate
227
# gemm implementation
@@ -273,6 +274,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
274
prepare_finalize: FusedMoEPrepareAndFinalize,
275
# TODO(bnell): Remove. Every layer should have an moe config object.
276
moe: FusedMoEConfig,
277
+ layer: torch.nn.Module,
278
) -> FusedMoEPermuteExpertsUnpermute:
279
if (prepare_finalize.activation_format ==
280
FusedMoEActivationFormat.BatchedExperts):
vllm/model_executor/layers/fused_moe/trtllm_moe.py
new
+197
@@ -0,0 +1,197 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import Optional
4
+
5
+import torch
6
+
7
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8
+from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
9
+from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
10
+ TopKWeightAndReduceNoOP)
11
+from vllm.utils import next_power_of_2
12
+
13
+
14
+class TrtLlmGenExperts(mk.FusedMoEPermuteExpertsUnpermute):
15
+
16
+ def __init__(
17
+ self,
18
+ moe: FusedMoEConfig,
19
+ gemm1_alpha,
20
+ gemm1_beta,
21
+ gemm1_clamp_limit,
22
+ w13_bias,
23
+ w2_bias,
24
+ max_capture_size,
25
+ ):
26
+ super().__init__(moe.quant_config)
27
+ self.moe = moe
28
+ self.gemm1_alpha = gemm1_alpha
29
+ self.gemm1_beta = gemm1_beta
30
+ self.gemm1_clamp_limit = gemm1_clamp_limit
31
+ self.w13_bias = w13_bias
32
+ self.w2_bias = w2_bias
33
+ self.max_capture_size = max_capture_size
34
+
35
+ @property
36
+ def activation_formats(
37
+ self
38
+ ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
39
+ return (mk.FusedMoEActivationFormat.Standard,
40
+ mk.FusedMoEActivationFormat.Standard)
41
+
42
+ def supports_chunking(self) -> bool:
43
+ return True
44
+
45
+ def supports_expert_map(self) -> bool:
46
+ return True
47
+
48
+ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
49
+ return TopKWeightAndReduceNoOP()
50
+
51
+ def workspace_shapes(
52
+ self,
53
+ a: torch.Tensor,
54
+ aq: torch.Tensor,
55
+ M: int,
56
+ N: int,
57
+ K: int,
58
+ topk: int,
59
+ global_num_experts: int,
60
+ local_num_experts: int,
61
+ expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
62
+ ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
63
+ # The workspaces for this implementation are managed by flashinfer.
64
+ # TODO(varun) : workspace1 is could be used as the output tensor. This
65
+ # is error-prone. Allow the `workspace_shapes` to return None workspaces
66
+ workspace1 = (M, K)
67
+ workspace2 = (0, 0)
68
+ output = (M, K)
69
+ return (workspace1, workspace2, output, a.dtype)
70
+
71
+ def _get_tile_tokens_dim(self, x: torch.Tensor, top_k: int,
72
+ local_num_experts: int):
73
+ # Number of tokens in the input tensor.
74
+ num_tokens = x.shape[0]
75
+ # Factor to account for the imbalance of the experts.
76
+ # factor equals to the
77
+ # max_real_num_tokens_per_expert / perfect_num_tokens_per_expert
78
+ # 1.0 means perfect expert distribution.
79
+ # > 1.0 means some experts have more tokens than the perfect
80
+ # distribution.
81
+ # < 1.0 does not make sense.
82
+ imbalance_factor = 1.3
83
+ # Calculate the number of tokens per expert assuming perfect
84
+ # distribution.
85
+ num_tokens_per_expert = (num_tokens * top_k) // local_num_experts
86
+ # Apply the imbalance factor.
87
+ num_tokens_per_expert = int(num_tokens_per_expert * imbalance_factor)
88
+ # And pad the number to the next power of 2.
89
+ tile_tokens_dim = next_power_of_2(num_tokens_per_expert)
90
+ # Cap to 8-64 tokens per CTA tile as it's the range supported by the
91
+ # kernel.
92
+ tile_tokens_dim = min(max(tile_tokens_dim, 8), 64)
93
+
94
+ return tile_tokens_dim
95
+
96
+ def apply(
97
+ self,
98
+ output: torch.Tensor,
99
+ hidden_states: torch.Tensor,
100
+ w1: torch.Tensor,
101
+ w2: torch.Tensor,
102
+ topk_weights: torch.Tensor,
103
+ topk_ids: torch.Tensor,
104
+ activation: str,
105
+ global_num_experts: int,
106
+ expert_map: Optional[torch.Tensor],
107
+ w1_scale: Optional[torch.Tensor],
108
+ w2_scale: Optional[torch.Tensor],
109
+ w1_zp: Optional[torch.Tensor],
110
+ w2_zp: Optional[torch.Tensor],
111
+ a1q_scale: Optional[torch.Tensor],
112
+ a2_scale: Optional[torch.Tensor],
113
+ workspace13: torch.Tensor,
114
+ workspace2: torch.Tensor,
115
+ expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
116
+ apply_router_weight_on_input: bool,
117
+ ):
118
+ topk = topk_ids.size(-1)
119
+ local_num_experts = w1.size(0)
120
+ intermediate_size = w2.size(1)
121
+ local_expert_offset = self.moe.ep_rank * local_num_experts
122
+
123
+ x_quant = hidden_states
124
+ x_scale = a1q_scale
125
+ if x_scale is not None:
126
+ x_scale = x_scale.view(torch.float8_e4m3fn).reshape(
127
+ *x_quant.shape[:-1], -1)
128
+
129
+ packed_tensor = (topk_ids.to(torch.int32) << 16) | topk_weights.to(
130
+ torch.bfloat16).view(torch.int16)
131
+
132
+ assert w1_scale is not None
133
+ assert w2_scale is not None
134
+ kwargs = {
135
+ "topk_ids":
136
+ packed_tensor,
137
+ "routing_bias":
138
+ None,
139
+ "hidden_states":
140
+ x_quant,
141
+ "hidden_states_scale":
142
+ x_scale,
143
+ "gemm1_weights":
144
+ w1,
145
+ "gemm1_weights_scale":
146
+ w1_scale,
147
+ "gemm1_bias":
148
+ self.w13_bias,
149
+ "gemm1_alpha":
150
+ self.gemm1_alpha,
151
+ "gemm1_beta":
152
+ self.gemm1_beta,
153
+ "gemm1_clamp_limit":
154
+ self.gemm1_clamp_limit,
155
+ "gemm2_weights":
156
+ w2,
157
+ "gemm2_weights_scale":
158
+ w2_scale,
159
+ "gemm2_bias":
160
+ self.w2_bias,
161
+ "output1_scale_scalar":
162
+ None,
163
+ "output1_scale_gate_scalar":
164
+ None,
165
+ "output2_scale_scalar":
166
+ None,
167
+ "num_experts":
168
+ global_num_experts,
169
+ "top_k":
170
+ topk,
171
+ "n_group":
172
+ None,
173
+ "topk_group":
174
+ None,
175
+ "intermediate_size":
176
+ intermediate_size,
177
+ "local_expert_offset":
178
+ local_expert_offset,
179
+ "local_num_experts":
180
+ local_num_experts,
181
+ "routed_scaling_factor":
182
+ None,
183
+ "tile_tokens_dim":
184
+ self._get_tile_tokens_dim(x_quant, topk, local_num_experts),
185
+ "routing_method_type":
186
+ 1,
187
+ "do_finalize":
188
+ True,
189
+ "output":
190
+ output,
191
+ "tune_max_num_tokens":
192
+ self.max_capture_size,
193
+ }
194
+
195
+ from flashinfer import trtllm_fp4_block_scale_routed_moe
196
+ trtllm_fp4_block_scale_routed_moe(**kwargs)
197
+ return output
vllm/model_executor/layers/fused_moe/utils.py
+16
@@ -12,6 +12,8 @@ from vllm.model_executor.layers.quantization.utils.int8_utils import (
12
per_token_group_quant_int8, per_token_quant_int8)
13
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
14
quant_dequant_mxfp4)
15
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
16
+ mxfp8_quantize)
17
from vllm.platforms import current_platform
18
from vllm.triton_utils import tl, triton
19
from vllm.utils import cdiv
@@ -177,6 +179,18 @@ def _mxfp4_quantize(
179
return A, None
180
181
182
+def _mxfp8_quantize(
183
+ A: torch.Tensor,
184
+ A_scale: Optional[torch.Tensor],
185
+ per_act_token_quant: bool,
186
+ block_shape: Optional[list[int]] = None,
187
+) -> tuple[torch.Tensor, torch.Tensor]:
188
+ assert A_scale is None
189
+ assert not per_act_token_quant
190
+ assert block_shape is None
191
+ return mxfp8_quantize(A)
192
+
193
+
194
def moe_kernel_quantize_input(
195
A: torch.Tensor,
196
A_scale: Optional[torch.Tensor],
@@ -195,6 +209,8 @@ def moe_kernel_quantize_input(
209
is_sf_swizzled_layout=is_fp4_scale_swizzled)
210
elif quant_dtype == "mxfp4":
211
return _mxfp4_quantize(A, A_scale, per_act_token_quant, block_shape)
212
+ elif quant_dtype == "mxfp8":
213
+ return _mxfp8_quantize(A, A_scale, per_act_token_quant, block_shape)
214
else:
215
return A, A_scale
216
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+4
-4
@@ -322,6 +322,7 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
322
self,
323
prepare_finalize: mk.FusedMoEPrepareAndFinalize,
324
moe: FusedMoEConfig,
325
+ layer: torch.nn.Module,
326
) -> mk.FusedMoEPermuteExpertsUnpermute:
327
"""Return the appropriate GEMM experts implementation."""
328
experts = select_nvfp4_gemm_impl(
@@ -719,10 +720,9 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
720
dtype=torch.int64)
721
722
def select_gemm_impl(
722
- self,
723
- prepare_finalize: FusedMoEPrepareAndFinalize,
724
- moe: FusedMoEConfig,
725
- ) -> FusedMoEPermuteExpertsUnpermute:
723
+ self, prepare_finalize: FusedMoEPrepareAndFinalize,
724
+ moe: FusedMoEConfig,
725
+ layer: torch.nn.Module) -> FusedMoEPermuteExpertsUnpermute:
726
# cutlass path
727
if self.use_cutlass:
728
from vllm.model_executor.layers.fused_moe import (
vllm/model_executor/layers/quantization/fp8.py
+1
@@ -897,6 +897,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
897
self,
898
prepare_finalize: FusedMoEPrepareAndFinalize,
899
moe: FusedMoEConfig,
900
+ layer: torch.nn.Module,
901
) -> FusedMoEPermuteExpertsUnpermute:
902
from vllm.model_executor.layers.fused_moe import (
903
BatchedTritonOrDeepGemmExperts, TritonOrDeepGemmExperts)
vllm/model_executor/layers/quantization/modelopt.py
+2
@@ -311,6 +311,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
311
self,
312
prepare_finalize: mk.FusedMoEPrepareAndFinalize,
313
moe: FusedMoEConfig,
314
+ layer: torch.nn.Module,
315
) -> mk.FusedMoEPermuteExpertsUnpermute:
316
experts = select_cutlass_fp8_gemm_impl(
317
moe,
@@ -1032,6 +1033,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1033
self,
1034
prepare_finalize: mk.FusedMoEPrepareAndFinalize,
1035
moe: FusedMoEConfig,
1036
+ layer: torch.nn.Module,
1037
) -> mk.FusedMoEPermuteExpertsUnpermute:
1038
experts = select_nvfp4_gemm_impl(
1039
moe,
vllm/model_executor/layers/quantization/mxfp4.py
+110
@@ -10,6 +10,8 @@ from vllm.config import get_current_vllm_config
10
from vllm.logger import init_logger
11
from vllm.model_executor.layers.fused_moe import (FusedMoE, FusedMoEConfig,
12
FusedMoEMethodBase)
13
+from vllm.model_executor.layers.fused_moe import modular_kernel as mk
14
+from vllm.model_executor.layers.fused_moe.trtllm_moe import TrtLlmGenExperts
15
from vllm.model_executor.layers.linear import (LinearBase,
16
UnquantizedLinearMethod)
17
from vllm.model_executor.layers.quantization import QuantizationMethods
@@ -445,6 +447,91 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
447
448
return tile_tokens_dim
449
450
+ def select_gemm_impl(
451
+ self,
452
+ prepare_finalize: mk.FusedMoEPrepareAndFinalize,
453
+ moe: FusedMoEConfig,
454
+ layer: torch.nn.Module,
455
+ ) -> mk.FusedMoEPermuteExpertsUnpermute:
456
+ if (prepare_finalize.activation_format ==
457
+ mk.FusedMoEActivationFormat.BatchedExperts):
458
+ raise NotImplementedError(
459
+ "Mxfp4 does not support batched experts format for EP")
460
+ else:
461
+ if should_use_flashinfer_mxfp4():
462
+ # B200 code-path
463
+ kwargs = {
464
+ "gemm1_alpha": layer.gemm1_alpha,
465
+ "gemm1_beta": layer.gemm1_beta,
466
+ "gemm1_clamp_limit": layer.gemm1_clamp_limit,
467
+ "w13_bias": layer.w13_bias,
468
+ "w2_bias": layer.w2_bias,
469
+ "max_capture_size": self.max_capture_size,
470
+ }
471
+ return TrtLlmGenExperts(moe, **kwargs)
472
+ else:
473
+ # Use matmul_ogs from triton_kernels here!
474
+ raise NotImplementedError(
475
+ "Mxfp4 does not support non-batched experts format for EP")
476
+
477
+ def _route_and_experts(
478
+ self,
479
+ layer: torch.nn.Module,
480
+ x: torch.Tensor,
481
+ router_logits: torch.Tensor,
482
+ top_k: int,
483
+ renormalize: bool,
484
+ use_grouped_topk: bool = False,
485
+ topk_group: Optional[int] = None,
486
+ num_expert_group: Optional[int] = None,
487
+ global_num_experts: int = -1,
488
+ expert_map: Optional[torch.Tensor] = None,
489
+ custom_routing_function: Optional[Callable] = None,
490
+ scoring_func: str = "softmax",
491
+ e_score_correction_bias: Optional[torch.Tensor] = None,
492
+ apply_router_weight_on_input: bool = False,
493
+ activation: str = "silu",
494
+ enable_eplb: bool = False,
495
+ expert_load_view: Optional[torch.Tensor] = None,
496
+ logical_to_physical_map: Optional[torch.Tensor] = None,
497
+ logical_replica_count: Optional[torch.Tensor] = None
498
+ ) -> torch.Tensor:
499
+
500
+ assert isinstance(self.fused_experts, mk.FusedMoEModularKernel)
501
+
502
+ topk_weights, topk_ids = FusedMoE.select_experts(
503
+ hidden_states=x,
504
+ router_logits=router_logits,
505
+ use_grouped_topk=use_grouped_topk,
506
+ top_k=top_k,
507
+ renormalize=renormalize,
508
+ topk_group=topk_group,
509
+ num_expert_group=num_expert_group,
510
+ custom_routing_function=custom_routing_function,
511
+ scoring_func=scoring_func,
512
+ e_score_correction_bias=e_score_correction_bias,
513
+ indices_type=self.topk_indices_dtype,
514
+ enable_eplb=enable_eplb,
515
+ expert_map=expert_map,
516
+ expert_load_view=expert_load_view,
517
+ logical_to_physical_map=logical_to_physical_map,
518
+ logical_replica_count=logical_replica_count)
519
+
520
+ return self.fused_experts(
521
+ hidden_states=x,
522
+ w1=layer.w13_weight,
523
+ w2=layer.w2_weight,
524
+ topk_weights=topk_weights,
525
+ topk_ids=topk_ids,
526
+ inplace=True,
527
+ activation=activation,
528
+ global_num_experts=global_num_experts,
529
+ expert_map=expert_map,
530
+ w1_scale=layer.w13_weight_scale,
531
+ w2_scale=layer.w2_weight_scale,
532
+ apply_router_weight_on_input=apply_router_weight_on_input,
533
+ )
534
+
535
def apply(
536
self,
537
layer: torch.nn.Module,
@@ -503,6 +590,29 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
590
activation=activation,
591
expert_map=expert_map)
592
593
+ if self.fused_experts is not None:
594
+ return self._route_and_experts(
595
+ layer,
596
+ x,
597
+ router_logits,
598
+ top_k,
599
+ renormalize,
600
+ use_grouped_topk,
601
+ topk_group,
602
+ num_expert_group,
603
+ global_num_experts,
604
+ expert_map,
605
+ custom_routing_function,
606
+ scoring_func,
607
+ e_score_correction_bias,
608
+ apply_router_weight_on_input,
609
+ activation,
610
+ enable_eplb,
611
+ expert_load_view,
612
+ logical_to_physical_map,
613
+ logical_replica_count,
614
+ )
615
+
616
assert _can_support_mxfp4(
617
use_grouped_topk, topk_group, num_expert_group, expert_map,
618
custom_routing_function, e_score_correction_bias,
vllm/model_executor/layers/quantization/utils/mxfp4_utils.py
+4
-5
@@ -66,11 +66,10 @@ def _can_support_mxfp4(use_grouped_topk: bool = False,
66
logical_to_physical_map: Optional[torch.Tensor] = None,
67
logical_replica_count: Optional[torch.Tensor] = None):
68
return not (use_grouped_topk or topk_group or num_expert_group
69
- or expert_map or custom_routing_function
70
- or e_score_correction_bias or apply_router_weight_on_input
71
- or scoring_func != "softmax" or activation != "swigluoai"
72
- or expert_load_view or logical_to_physical_map
73
- or logical_replica_count)
69
+ or custom_routing_function or e_score_correction_bias
70
+ or apply_router_weight_on_input or scoring_func != "softmax"
71
+ or activation != "swigluoai" or expert_load_view
72
+ or logical_to_physical_map or logical_replica_count)
73
74
75
def _dequant_mxfp4(x: torch.Tensor, scale: torch.Tensor,
vllm/model_executor/layers/quantization/utils/mxfp8_utils.py
new
+20
@@ -0,0 +1,20 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+from vllm.logger import init_logger
7
+
8
+logger = init_logger(__name__)
9
+
10
+
11
+def mxfp8_quantize(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
12
+
13
+ try:
14
+ from flashinfer import mxfp8_quantize
15
+ except ImportError as err:
16
+ raise ImportError("The package `flashinfer` is required to do "
17
+ "MX-FP8 quantization. Please install it with" \
18
+ "`pip install flashinfer`") from err
19
+
20
+ return mxfp8_quantize(x, is_sf_swizzled_layout=False)