[Perf][Nemotron] Skip redundant latent-MoE all-reduce at TP>1 (~13% decode win) (#52301)
Signed-off-by: Junhao Shen <junshen@nvidia.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Misha Goin <mgoin64@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Junhao Shen committed
Sep 15, 2026 at 22:33 UTC
4bac767695d2dfecd4098a77f71dacd100061247
2 files changed
+21
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+9
@@ -457,9 +457,18 @@ class MoERunner(MoERunnerInterface):
457
Latent MoE output transforms may contain non-linear ops, e.g. RMSNorm.
458
TP partial routed outputs must be summed in latent space before such
459
transforms are applied.
460
+
461
+ A transform that commutes with the TP sum is exempt: if
462
+ ``sum_r T(x_r) == T(sum_r x_r)``, applying the transform to the local
463
+ partial output and letting the existing late all-reduce sum the
464
+ combined result is equivalent, and costs one collective instead of two.
465
+ Such a transform opts out by setting ``reduce_commutative = True``.
466
+ The default is False, so transforms that do not declare themselves
467
+ keep being reduced early.
468
"""
469
if (
470
self.routed_output_transform is not None
471
+ and not getattr(self.routed_output_transform, "reduce_commutative", False)
472
and not self.moe_config.is_sequence_parallel
473
and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1)
474
and not fused_output_is_reduced
vllm/model_executor/models/nemotron_h.py
+12
@@ -49,6 +49,7 @@ from vllm.model_executor.layers.linear import (
49
QKVParallelLinear,
50
ReplicatedLinear,
51
RowParallelLinear,
52
+ UnquantizedLinearMethod,
53
)
54
from vllm.model_executor.layers.logits_processor import LogitsProcessor
55
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
@@ -204,6 +205,17 @@ class NemotronHMoE(nn.Module):
205
disable_tp=self.is_sequence_parallel,
206
prefix=f"{prefix}.fc2_latent_proj",
207
)
208
+ # A bias-free, unquantized linear commutes with the TP sum
209
+ # (sum_r W x_r == W sum_r x_r), so one reduce after the transform
210
+ # suffices. Test the layer, not the model-wide `quant_config`: ModelOpt
211
+ # excludes the latent projections, so quantized checkpoints still get
212
+ # an UnquantizedLinearMethod here.
213
+ self.fc2_latent_proj.reduce_commutative = (
214
+ not config.mlp_bias
215
+ and isinstance(
216
+ self.fc2_latent_proj.quant_method, UnquantizedLinearMethod
217
+ )
218
+ )
219
else:
220
self.fc1_latent_proj = None
221
self.fc2_latent_proj = None