Add an opt-in fused weighted restore for AutoEP

deepspeed/module_inject/auto_ep_config.py

@@ -118,6 +118,18 @@ def validate_autoep_config(
     if not config.enabled:
         return
 
+    # Reject configurations that would bypass the requested fused reduction.
+    if config.combine_impl == "fused_weighted_sum":
+        if tp_size > 1:
+            raise ValueError('combine_impl="fused_weighted_sum" does not support folded tensor parallelism '
+                             f"(tensor_parallel.autotp_size={tp_size}), which restores combined tokens from "
+                             "assignment metadata instead of the weighted reduction it implements. Set "
+                             'tensor_parallel.autotp_size to 1, or leave combine_impl unset.')
+        if config.expert_tensor_parallel_size > 1:
+            raise ValueError('combine_impl="fused_weighted_sum" requires expert_tensor_parallel_size=1, but got '
+                             f"{config.expert_tensor_parallel_size}. Set expert_tensor_parallel_size to 1, or leave "
+                             "combine_impl unset.")
+
     folding_spec = build_folding_spec(
         world_size=world_size,
         pp_size=pp_size,
@@ -152,7 +164,7 @@ def validate_autoep_config(
                          f"got '{config.score_apply}'")
 
     # Validate combine_impl
-    valid_combine_impl = ("auto", "weighted_sum", "legacy_bmm")
+    valid_combine_impl = ("auto", "weighted_sum", "fused_weighted_sum", "legacy_bmm")
     if config.combine_impl not in valid_combine_impl:
         raise ValueError(f"combine_impl must be one of {valid_combine_impl}, "
                          f"got '{config.combine_impl}'")
@@ -272,6 +284,15 @@ def validate_autoep_post_detection(
         return
 
     for spec in specs:
+        # The fused reduction folds the routing weight into the top-k reduction,
+        # which only exists when scores are applied after the experts.
+        if config.combine_impl == "fused_weighted_sum":
+            resolved_score_apply = config.score_apply if config.score_apply != "auto" else spec.score_apply
+            if resolved_score_apply != "post":
+                raise ValueError(f'combine_impl="fused_weighted_sum" requires score_apply="post", but layer '
+                                 f"'{spec.moe_module_name}' resolved score_apply=\"{resolved_score_apply}\". "
+                                 "Leave combine_impl unset.")
+
         # ep_size must not exceed num_experts
         if config.autoep_size > spec.num_experts:
             valid_divisors = _divisors(spec.num_experts)

deepspeed/module_inject/auto_ep_layer.py

@@ -21,6 +21,7 @@
 import deepspeed.comm as dist
 from deepspeed.module_inject.auto_ep_config import AutoEPConfig, MoELayerSpec, resolve_autoep_config_defaults
 from deepspeed.module_inject.auto_ep_folding import mark_autoep_folding_router_parameter
+from deepspeed.ops.triton_ops import autoep_fused_token_ops as fused_token_ops
 from deepspeed.utils import logger
 from deepspeed.moe.ep_router import TokenChoiceTopKRouter
 from deepspeed.moe.ep_count import count_tokens_per_expert
@@ -61,7 +62,8 @@ def resolve_score_apply_mode(
 
 
 def resolve_combine_impl(
-    config_override: Literal["auto", "weighted_sum", "legacy_bmm"], ) -> Literal["weighted_sum", "legacy_bmm"]:
+    config_override: Literal["auto", "weighted_sum", "fused_weighted_sum", "legacy_bmm"],
+) -> Literal["weighted_sum", "fused_weighted_sum", "legacy_bmm"]:
     """Resolve combine implementation from config override or default."""
     if config_override != "auto":
         return config_override
@@ -377,6 +379,7 @@ def __init__(
         self.top_k = spec.top_k
         self.score_apply = resolve_score_apply_mode(spec, config.score_apply)
         self.combine_impl = resolve_combine_impl(config.combine_impl)
+        self._fused_combine_checked = False
         route_norm = spec.route_norm if config.route_norm is None else config.route_norm
         self.ep_size = ep_size
         self.ep_rank = ep_rank
@@ -550,6 +553,11 @@ def set_deepspeed_parallelism(
 
         if folding_group_handles is not None:
             self.folding_group_handles = folding_group_handles
+            if self.combine_impl == "fused_weighted_sum" and folding_group_handles.spec.tp_size > 1:
+                # Folded TP restores tokens through a different path.
+                raise ValueError('combine_impl="fused_weighted_sum" does not support folded tensor parallelism '
+                                 f"(tensor_parallel.autotp_size={folding_group_handles.spec.tp_size}). Set "
+                                 'tensor_parallel.autotp_size to 1, or leave combine_impl unset.')
             self.ep_group_name = folding_group_handles.ep_group_name
             self.ep_group = folding_group_handles.ep_group
             self.tp_group = folding_group_handles.tp_group
@@ -588,6 +596,11 @@ def forward(
         bsz, seqlen, hdim = hidden_states.shape
         x = hidden_states.reshape(-1, hdim)  # [T, H]
 
+        # Fail all ranks before any collective can stall.
+        if self.combine_impl == "fused_weighted_sum" and not self._fused_combine_checked:
+            fused_token_ops.assert_supported(x, score_apply=self.score_apply)
+            self._fused_combine_checked = True
+
         # Router
         ro: RouterOutput = RouterOutput(*self.router(x, self.expert_bias))
 
@@ -693,6 +706,14 @@ def forward(
                                       tp_group=self.tp_group,
                                       validate_coverage=self.validate_folding_routing).reshape(bsz, seqlen, hdim)
             self._last_folding_dispatch_counters = dispatch_counters(restore_ctx)
+        elif self.combine_impl == "fused_weighted_sum":
+            output = fused_token_ops.fused_weighted_restore(
+                expert_output,
+                top_scores=ro.top_scores,
+                token_indices_sorted=token_indices_sorted,
+                top_k=self.top_k,
+                shape=(bsz, seqlen, hdim),
+            )
         else:
             output = combine_from_routed(
                 expert_output,

deepspeed/module_inject/auto_ep_presets/base.py

@@ -109,7 +109,7 @@ class AutoEPConfig:
     route_norm: bool | None = None
     route_scale: float = 1.0
     score_apply: Literal["auto", "pre", "post"] = "auto"
-    combine_impl: Literal["auto", "weighted_sum", "legacy_bmm"] = "auto"
+    combine_impl: Literal["auto", "weighted_sum", "fused_weighted_sum", "legacy_bmm"] = "auto"
     num_expert_groups: int | None = None
     num_limited_groups: int | None = None
     score_func: Literal["auto", "softmax", "sigmoid"] = "auto"