Add an opt-in fused weighted restore for AutoEP

deepspeed/module_inject/auto_ep_config.py

@@ -122,6 +122,22 @@ 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.")
+        if config.comm_backend == "deepep" and config.autoep_size > 1:
+            raise ValueError('combine_impl="fused_weighted_sum" cannot be used with comm_backend="deepep" because '
+                             "DeepEP already restores and reduces the routed rows. Set comm_backend to "
+                             '"comm", or leave combine_impl unset.')
+
     folding_spec = build_folding_spec(
         world_size=world_size,
         pp_size=pp_size,
@@ -156,7 +172,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}'")
@@ -297,6 +313,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.module_inject.auto_ep_comm import (DEEPEP_BACKEND, DeepEPExchange, assert_dtype_supported,
                                                   deepep_combine, deepep_dispatch)
@@ -63,7 +64,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
@@ -379,6 +381,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
@@ -559,6 +562,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.')
             if self.comm_backend == DEEPEP_BACKEND and folding_group_handles.spec.tp_size > 1:
                 # DeepEP's combine returns token-major rows, which folded TP's
                 # assignment-metadata restore can't consume. Refuse rather than
@@ -672,6 +680,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))
 
@@ -790,6 +803,14 @@ def forward(
             # backend being selected: with ep_size == 1 the local path runs
             # instead and still has one row per assignment to reduce.
             output = expert_output.reshape(bsz, seqlen, hdim)
+        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"
     comm_backend: Literal["comm", "deepep"] = "comm"
     comm_num_sm: int = 12
     comm_qp_margin: int = 4

deepspeed/ops/triton_ops/autoep_fused_token_ops.py

@@ -0,0 +1,274 @@
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+"""Fused AutoEP token restore without the eager scatter and FP32 intermediate.
+
+The kernel reduces each token's top-k rows in FP32. Communication, routing,
+expert reorder, and grouped GEMM remain unchanged.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from deepspeed.ops.triton_ops._triton import _TRITON_AVAILABLE, triton, tl
+
+_IS_ROCM_PYTORCH = getattr(torch.version, "hip", None) is not None
+
+SUPPORTED_ROW_DTYPES = (torch.bfloat16, torch.float16, torch.float32)
+
+_MAX_BLOCK_HIDDEN = 512
+_INVERT_INDEX_BLOCK = 256
+# The kernels hold a [slots, BLOCK_H] FP32 block live, so the hidden tile shrinks
+# as top-k grows to keep that block in registers rather than spilling.
+_MAX_BLOCK_ELEMENTS = 2048
+
+if _TRITON_AVAILABLE:
+
+    @triton.jit
+    def _invert_index_kernel(
+        index_ptr,
+        inverse_ptr,
+        num_indices,
+        num_inverse_rows,
+        BLOCK: tl.constexpr,
+    ):
+        offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+        in_range = offsets < num_indices
+
+        targets = tl.load(index_ptr + offsets, mask=in_range, other=-1).to(tl.int64)
+        writable = in_range & (targets >= 0) & (targets < num_inverse_rows)
+        tl.store(inverse_ptr + tl.where(writable, targets, 0), offsets.to(tl.int32), mask=writable)
+
+    @triton.jit
+    def _weighted_restore_forward_kernel(
+        rows_ptr,
+        inverse_ptr,
+        scores_ptr,
+        out_ptr,
+        hidden,
+        rows_stride,
+        scores_stride,
+        out_stride,
+        TOP_K: tl.constexpr,
+        K_PADDED: tl.constexpr,
+        BLOCK_H: tl.constexpr,
+    ):
+        token = tl.program_id(0)
+        hidden_offsets = tl.program_id(1) * BLOCK_H + tl.arange(0, BLOCK_H)
+        hidden_mask = hidden_offsets < hidden
+
+        slots = tl.arange(0, K_PADDED)
+        slot_mask = slots < TOP_K
+
+        source_rows = tl.load(inverse_ptr + token * TOP_K + slots, mask=slot_mask, other=-1).to(tl.int64)
+        row_valid = slot_mask & (source_rows >= 0)
+        safe_rows = tl.where(row_valid, source_rows, 0)
+
+        scores = tl.load(scores_ptr + token * scores_stride + slots, mask=slot_mask, other=0.0).to(tl.float32)
+
+        block_mask = row_valid[:, None] & hidden_mask[None, :]
+        values = tl.load(
+            rows_ptr + safe_rows[:, None] * rows_stride + hidden_offsets[None, :],
+            mask=block_mask,
+            other=0.0,
+        ).to(tl.float32)
+
+        # Match the eager path's FP32 product and accumulation.
+        weighted = tl.sum(values * scores[:, None], axis=0)
+        tl.store(
+            out_ptr + token * out_stride + hidden_offsets,
+            weighted.to(out_ptr.dtype.element_ty),
+            mask=hidden_mask,
+        )
+
+    @triton.jit
+    def _weighted_restore_backward_kernel(
+        grad_out_ptr,
+        rows_ptr,
+        inverse_ptr,
+        scores_ptr,
+        grad_rows_ptr,
+        grad_scores_ptr,
+        hidden,
+        grad_out_stride,
+        rows_stride,
+        scores_stride,
+        grad_rows_stride,
+        grad_scores_stride,
+        TOP_K: tl.constexpr,
+        K_PADDED: tl.constexpr,
+        BLOCK_H: tl.constexpr,
+    ):
+        token = tl.program_id(0)
+
+        slots = tl.arange(0, K_PADDED)
+        slot_mask = slots < TOP_K
+
+        source_rows = tl.load(inverse_ptr + token * TOP_K + slots, mask=slot_mask, other=-1).to(tl.int64)
+        row_valid = slot_mask & (source_rows >= 0)
+        safe_rows = tl.where(row_valid, source_rows, 0)
+        scores = tl.load(scores_ptr + token * scores_stride + slots, mask=slot_mask, other=0.0).to(tl.float32)
+
+        grad_rows_dtype = grad_rows_ptr.dtype.element_ty
+        # Keeping one token per program avoids a second reduction pass for scores.
+        score_partials = tl.zeros([K_PADDED, BLOCK_H], dtype=tl.float32)
+
+        for hidden_start in range(0, hidden, BLOCK_H):
+            hidden_offsets = hidden_start + tl.arange(0, BLOCK_H)
+            hidden_mask = hidden_offsets < hidden
+            block_mask = row_valid[:, None] & hidden_mask[None, :]
+
+            upstream = tl.load(
+                grad_out_ptr + token * grad_out_stride + hidden_offsets,
+                mask=hidden_mask,
+                other=0.0,
+            ).to(tl.float32)
+
+            values = tl.load(
+                rows_ptr + safe_rows[:, None] * rows_stride + hidden_offsets[None, :],
+                mask=block_mask,
+                other=0.0,
+            ).to(tl.float32)
+            score_partials += values * upstream[None, :]
+
+            tl.store(
+                grad_rows_ptr + safe_rows[:, None] * grad_rows_stride + hidden_offsets[None, :],
+                (upstream[None, :] * scores[:, None]).to(grad_rows_dtype),
+                mask=block_mask,
+            )
+
+        grad_scores = tl.sum(score_partials, axis=1)
+        tl.store(
+            grad_scores_ptr + token * grad_scores_stride + slots,
+            grad_scores.to(grad_scores_ptr.dtype.element_ty),
+            mask=slot_mask,
+        )
+
+
+def is_available() -> bool:
+    """Whether this build can run the fused weighted restore at all."""
+    return _TRITON_AVAILABLE and not _IS_ROCM_PYTORCH
+
+
+def assert_supported(rows: torch.Tensor, *, score_apply: str) -> None:
+    """Reject unsupported configurations before collectives begin."""
+    if not _TRITON_AVAILABLE:
+        raise RuntimeError('combine_impl="fused_weighted_sum" needs Triton, which is not installed in this '
+                           "environment. Install Triton, or leave combine_impl unset.")
+    if _IS_ROCM_PYTORCH:
+        raise RuntimeError('combine_impl="fused_weighted_sum" is not yet supported on ROCm. Leave combine_impl '
+                           "unset to run here.")
+    if rows.device.type != "cuda":
+        raise RuntimeError('combine_impl="fused_weighted_sum" runs CUDA kernels but this layer is on device '
+                           f'"{rows.device.type}". Leave combine_impl unset to run here.')
+    if rows.dtype not in SUPPORTED_ROW_DTYPES:
+        raise RuntimeError('combine_impl="fused_weighted_sum" supports bfloat16, float16, and float32 rows, got '
+                           f"{rows.dtype}. Leave combine_impl unset, or use a supported floating-point dtype.")
+    if score_apply != "post":
+        raise RuntimeError('combine_impl="fused_weighted_sum" folds the routing weight into the top-k reduction, '
+                           f'which only exists for score_apply="post", but this layer resolved '
+                           f'score_apply="{score_apply}". Leave combine_impl unset.')
+
+
+def _block_hidden(hidden: int, slots: int) -> int:
+    """Choose a power-of-two tile within the FP32 register budget."""
+    budget = max(16, _MAX_BLOCK_ELEMENTS // slots)
+    return min(_MAX_BLOCK_HIDDEN, budget, max(16, triton.next_power_of_2(hidden)))
+
+
+def _padded_top_k(top_k: int) -> int:
+    """Round top-k up to a power of two, which ``tl.arange`` requires."""
+    return max(2, triton.next_power_of_2(top_k))
+
+
+def _invert_index(index: torch.Tensor, num_inverse_rows: int) -> torch.Tensor:
+    """Invert the row permutation produced by sorting routed assignments."""
+    inverse = torch.empty((num_inverse_rows, ), dtype=torch.int32, device=index.device)
+    num_indices = index.numel()
+
+    grid = (triton.cdiv(num_indices, _INVERT_INDEX_BLOCK), )
+    _invert_index_kernel[grid](
+        index.contiguous(),
+        inverse,
+        num_indices,
+        num_inverse_rows,
+        BLOCK=_INVERT_INDEX_BLOCK,
+    )
+    return inverse
+
+
+class _FusedWeightedRestore(torch.autograd.Function):
+    """Weight rows by their routing score and reduce over top-k in one pass."""
+
+    @staticmethod
+    def forward(ctx, combined_rows, top_scores, inverse, top_k):
+        combined_rows = combined_rows.contiguous()
+        n_tokens, hidden = top_scores.shape[0], combined_rows.shape[-1]
+        output = torch.empty((n_tokens, hidden), dtype=combined_rows.dtype, device=combined_rows.device)
+
+        ctx.save_for_backward(combined_rows, top_scores, inverse)
+        ctx.top_k = top_k
+
+        k_padded = _padded_top_k(top_k)
+        block_hidden = _block_hidden(hidden, slots=k_padded)
+        grid = (n_tokens, triton.cdiv(hidden, block_hidden))
+        _weighted_restore_forward_kernel[grid](
+            combined_rows,
+            inverse,
+            top_scores,
+            output,
+            hidden,
+            combined_rows.stride(0),
+            top_scores.stride(0),
+            output.stride(0),
+            TOP_K=top_k,
+            K_PADDED=k_padded,
+            BLOCK_H=block_hidden,
+        )
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        combined_rows, top_scores, inverse = ctx.saved_tensors
+        grad_output = grad_output.contiguous()
+
+        grad_rows = torch.empty_like(combined_rows)
+        grad_scores = torch.empty_like(top_scores)
+
+        n_tokens, hidden = top_scores.shape[0], combined_rows.shape[-1]
+
+        k_padded = _padded_top_k(ctx.top_k)
+        _weighted_restore_backward_kernel[(n_tokens, )](
+            grad_output,
+            combined_rows,
+            inverse,
+            top_scores,
+            grad_rows,
+            grad_scores,
+            hidden,
+            grad_output.stride(0),
+            combined_rows.stride(0),
+            top_scores.stride(0),
+            grad_rows.stride(0),
+            grad_scores.stride(0),
+            TOP_K=ctx.top_k,
+            K_PADDED=k_padded,
+            BLOCK_H=_block_hidden(hidden, slots=k_padded),
+        )
+        return grad_rows, grad_scores, None, None
+
+
+def fused_weighted_restore(
+    combined_rows: torch.Tensor,
+    top_scores: torch.Tensor,
+    token_indices_sorted: torch.Tensor,
+    top_k: int,
+    shape: tuple[int, int, int],
+) -> torch.Tensor:
+    """Restore ``[T * K, H]`` rows directly to weighted ``[B, S, H]`` output."""
+    bsz, seqlen, hidden = shape
+    n_tokens = bsz * seqlen
+    expected_rows = n_tokens * top_k
+    inverse = _invert_index(token_indices_sorted, expected_rows)
+    output = _FusedWeightedRestore.apply(combined_rows, top_scores.contiguous(), inverse, top_k)
+    return output.reshape(bsz, seqlen, hidden)