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 13, 23:15 UTC
Flashinfer_CUTLASS_MOE fuses quantization for TP (#27223)
Signed-off-by: Shu Wang. <shuw@nvidia.com>
Shu Wang committed
Oct 31, 2025 at 10:54 UTC
fc16f1c4779023e2e9ca3efe7a8b78a27cd318ac
4 files changed
+15
-32
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
+5
-1
@@ -56,6 +56,7 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
56
ep_size: int = 1,
57
tp_rank: int = 0,
58
tp_size: int = 1,
59
+ use_dp: bool = False,
60
):
61
super().__init__(quant_config)
62
assert quant_config.quant_dtype in ("nvfp4", torch.float8_e4m3fn, None), (
@@ -67,6 +68,7 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
68
self.tp_rank = tp_rank
69
self.tp_size = tp_size
70
self.out_dtype = out_dtype
71
+ self.use_dp = use_dp
72
73
@property
74
def activation_formats(
@@ -117,7 +119,8 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
119
"""
120
workspace1 = (M, K)
121
workspace2 = (0,)
120
- output_shape = (M, K * 2 if self.quant_dtype == "nvfp4" else K)
122
+ # For TP, the quantization is fused with fused_moe call.
123
+ output_shape = (M, K * 2 if self.quant_dtype == "nvfp4" and self.use_dp else K)
124
# The workspace is determined by `aq`, since it comes after any
125
# potential communication op and is involved in the expert computation.
126
return (workspace1, workspace2, output_shape)
@@ -214,6 +217,7 @@ def flashinfer_cutlass_moe_fp4(
217
FlashInferExperts(
218
out_dtype=hidden_states.dtype,
219
quant_config=quant_config,
220
+ use_dp=False,
221
),
222
)
223
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py
+9
-8
@@ -170,6 +170,8 @@ class FlashInferAllGatherMoEPrepareAndFinalize(FlashInferCutlassMoEPrepareAndFin
170
self._apply_router_weight_on_input(
171
a1, topk_weights, topk_ids, apply_router_weight_on_input
172
)
173
+ if not self.use_dp:
174
+ return a1, None, None, topk_ids, topk_weights
175
176
a1q, a1q_scale = moe_kernel_quantize_input(
177
a1,
@@ -179,14 +181,13 @@ class FlashInferAllGatherMoEPrepareAndFinalize(FlashInferCutlassMoEPrepareAndFin
181
quant_config.block_shape,
182
is_fp4_scale_swizzled=not self.use_dp,
183
)
182
- if self.use_dp:
183
- topk_weights, topk_ids, a1q, a1q_scale = get_dp_group().all_gatherv(
184
- [topk_weights, topk_ids, a1q, a1q_scale],
185
- dim=0,
186
- sizes=get_local_sizes(),
187
- )
188
- if quant_config.quant_dtype == "nvfp4":
189
- a1q_scale = nvfp4_block_scale_interleave(a1q_scale)
184
+ topk_weights, topk_ids, a1q, a1q_scale = get_dp_group().all_gatherv(
185
+ [topk_weights, topk_ids, a1q, a1q_scale],
186
+ dim=0,
187
+ sizes=get_local_sizes(),
188
+ )
189
+ if quant_config.quant_dtype == "nvfp4":
190
+ a1q_scale = nvfp4_block_scale_interleave(a1q_scale)
191
192
return a1q, a1q_scale, None, topk_ids, topk_weights
193
vllm/model_executor/layers/quantization/modelopt.py
-23
@@ -1769,29 +1769,6 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1769
expert_map=expert_map,
1770
apply_router_weight_on_input=apply_router_weight_on_input,
1771
)
1772
- elif (
1773
- self.allow_flashinfer
1774
- and self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS
1775
- ):
1776
- from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( # noqa: E501
1777
- flashinfer_cutlass_moe_fp4,
1778
- )
1779
-
1780
- assert self.moe_quant_config is not None
1781
-
1782
- return flashinfer_cutlass_moe_fp4(
1783
- hidden_states=x,
1784
- w1=layer.w13_weight,
1785
- w2=layer.w2_weight,
1786
- topk_weights=topk_weights,
1787
- topk_ids=topk_ids,
1788
- quant_config=self.moe_quant_config,
1789
- inplace=False,
1790
- activation=activation,
1791
- global_num_experts=global_num_experts,
1792
- expert_map=expert_map,
1793
- apply_router_weight_on_input=apply_router_weight_on_input,
1794
- )
1772
else:
1773
# If no modular kernel is provided, use cutlass_moe_fp4 for TP case
1774
# only (no EP).
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+1
@@ -79,6 +79,7 @@ def select_nvfp4_gemm_impl(
79
ep_size=moe.moe_parallel_config.ep_size,
80
tp_rank=moe.moe_parallel_config.tp_rank,
81
tp_size=moe.moe_parallel_config.tp_size,
82
+ use_dp=moe.moe_parallel_config.dp_size > 1,
83
)
84
85
# native cutlass experts currently don't support DP; TP case won't call this