auto_ep_layer.py

# Copyright (c) DeepSpeed Team.
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# SPDX-License-Identifier: Apache-2.0 AND BSD-3-Clause
#
# Portions of this file are derived from TorchTitan.
# See THIRD_PARTY_NOTICES.md for the BSD-3-Clause notice.

# DeepSpeed Team
"""AutoEP MoE Layer: drop-in replacement for HF MoE blocks with EP support.

Contains AutoEPMoELayer, compute_split_plan, _AllToAllV, and helper functions.
"""

from __future__ import annotations

from typing import Literal, NamedTuple

import torch
import torch.nn as nn
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
from deepspeed.moe.ep_experts import GroupedExperts
from deepspeed.moe.ep_repack import _gather_source_zero_params, repack_expert_requires_grad_flags, repack_expert_weights

# ---------------------------------------------------------------------------
# Named tuples
# ---------------------------------------------------------------------------


class RouterOutput(NamedTuple):
    top_scores: torch.Tensor  # [T, K]
    selected_experts: torch.Tensor  # [T, K]
    num_tokens_per_expert: torch.Tensor  # [E_global]


class SplitPlan(NamedTuple):
    input_splits: list[int]  # len=ep_size
    output_splits: list[int]  # len=ep_size
    local_counts: torch.Tensor  # [E_local]
    local_counts_by_source: torch.Tensor  # [ep_size, E_local]


# ---------------------------------------------------------------------------
# Helper functions
# ---------------------------------------------------------------------------


def resolve_score_apply_mode(
    spec: MoELayerSpec,
    config_override: Literal["auto", "pre", "post"],
) -> Literal["pre", "post"]:
    """Resolve score-application mode from config override or preset default."""
    if config_override != "auto":
        return config_override
    return spec.score_apply


def resolve_combine_impl(
    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
    return "weighted_sum"


def _copy_parameter_data(target: nn.Parameter, source: torch.Tensor) -> None:
    full_shape = torch.Size(getattr(source, "ds_shape", source.shape))
    with torch.no_grad():
        source_data = source.data
        if torch.Size(source_data.shape) != full_shape:
            raise RuntimeError("AutoEP source parameter must be gathered before copying: "
                               f"expected full shape {tuple(full_shape)}, got {tuple(source_data.shape)}")
        if (torch.Size(target.data.shape) != full_shape or target.data.dtype != source_data.dtype
                or target.data.device != source_data.device):
            target.data = torch.empty(full_shape, dtype=source_data.dtype, device=source_data.device)
        target.data.copy_(source_data)


def apply_scores_before_experts_if_enabled(
    routed_input: torch.Tensor,
    top_scores: torch.Tensor,
    score_apply: Literal["pre", "post"],
) -> torch.Tensor:
    """Pre-multiply token representations by router scores before expert compute."""
    if score_apply == "pre":
        return (routed_input.to(torch.float32) * top_scores.reshape(-1, 1)).to(routed_input.dtype)
    return routed_input


def _split_plan_from_expert_counts(
    num_tokens_per_expert: torch.Tensor,  # [E_global], int32
    ep_size: int,
    num_local_experts: int,
    ep_group: dist.ProcessGroup | None,
) -> SplitPlan:
    """Build a SplitPlan from a global per-expert token histogram (ep_size > 1).

    A single all-to-all exchanges the whole ``[ep_size, E_local]`` count matrix.
    Row ``s`` of the received matrix holds the per-local-expert counts sent by
    source rank ``s``, so summing it over experts recovers exactly the per-rank
    totals that a separate rank-count exchange would have produced.
    """
    count_matrix = num_tokens_per_expert.view(ep_size, num_local_experts)

    send_counts = count_matrix.reshape(-1).contiguous()  # [ep_size * E_local]
    received_counts_flat = torch.empty_like(send_counts)
    dist.all_to_all_single(
        received_counts_flat,
        send_counts,
        group=ep_group,
    )
    received_counts = received_counts_flat.view(ep_size, num_local_experts)

    # input_splits: tokens THIS rank sends to each destination rank.
    # output_splits: tokens THIS rank receives from each source rank.
    # Stacked so both split lists share a single device-to-host sync.
    input_splits, output_splits = torch.stack((
        count_matrix.sum(dim=1),
        received_counts.sum(dim=1),
    )).cpu().tolist()

    return SplitPlan(
        input_splits=input_splits,
        output_splits=output_splits,
        local_counts=received_counts.sum(dim=0),  # [E_local]
        local_counts_by_source=received_counts,
    )


def compute_split_plan(
        selected_experts: torch.Tensor,  # [T, K]
        num_experts: int,
        ep_size: int,
        num_local_experts: int,
        ep_group: dist.ProcessGroup | None,
        num_tokens_per_expert: torch.Tensor | None = None,  # [E_global], int32
) -> SplitPlan:
    """Compute AllToAllV split sizes for token dispatch/combine.

    ``num_tokens_per_expert`` may be supplied by the caller to reuse the
    histogram already computed by the router; when omitted it is derived from
    ``selected_experts``.

    Returns SplitPlan with input_splits, output_splits, local_counts, and
    local_counts_by_source.
    """
    if num_tokens_per_expert is None:
        num_tokens_per_expert = count_tokens_per_expert(selected_experts, num_experts)

    if ep_size == 1:
        # No dispatch needed - all tokens stay local
        T_K = selected_experts.numel()
        return SplitPlan(
            input_splits=[T_K],
            output_splits=[T_K],
            local_counts=num_tokens_per_expert,
            local_counts_by_source=num_tokens_per_expert.view(1, num_local_experts),
        )

    return _split_plan_from_expert_counts(num_tokens_per_expert, ep_size, num_local_experts, ep_group)


def compute_split_plan_from_expert_indices(
    expert_indices: torch.Tensor,
    num_experts: int,
    ep_size: int,
    num_local_experts: int,
    ep_group: dist.ProcessGroup | None,
) -> SplitPlan:
    """Compute EP AllToAllV splits for an already partitioned assignment list."""
    counts = count_tokens_per_expert(expert_indices, num_experts)
    if ep_size == 1:
        return SplitPlan([int(expert_indices.numel())], [int(expert_indices.numel())], counts,
                         counts.view(1, num_local_experts))

    return _split_plan_from_expert_counts(counts, ep_size, num_local_experts, ep_group)


class _AllToAllV(torch.autograd.Function):
    """Autograd-compatible all-to-all with variable split sizes."""

    @staticmethod
    def forward(ctx, group, x, input_splits, output_splits):
        ctx.group = group
        ctx.input_splits = input_splits
        ctx.output_splits = output_splits

        output_size = sum(output_splits)
        output = torch.empty(
            (output_size, x.shape[1]),
            dtype=x.dtype,
            device=x.device,
        )

        dist.all_to_all_single(
            output,
            x.contiguous(),
            output_split_sizes=output_splits,
            input_split_sizes=input_splits,
            group=group,
        )
        return output

    @staticmethod
    def backward(ctx, grad_out):
        # Reverse the splits for backward
        grad_out = grad_out.contiguous()
        input_size = sum(ctx.input_splits)
        grad_input = torch.empty(
            (input_size, grad_out.shape[1]),
            dtype=grad_out.dtype,
            device=grad_out.device,
        )

        dist.all_to_all_single(
            grad_input,
            grad_out,
            output_split_sizes=ctx.input_splits,
            input_split_sizes=ctx.output_splits,
            group=ctx.group,
        )
        return None, grad_input, None, None


def permute_by_local_expert(
    tokens: torch.Tensor,
    local_counts: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
    """Reorder tokens so they are grouped contiguously by local expert ID.

    Uses TorchTitan's Triton kernel for permutation index generation.

    Returns:
        tokens_permuted: [N_padded, H] (alignment-padded)
        permuted_indices: [N_padded] (maps padded positions -> original positions)
        aligned_counts: [E_local] aligned token counts per expert (for expert computation)
        n_tokens: original token count before padding (for unpermute)
    """
    from deepspeed.moe.ep_kernels import generate_permute_indices, TOKEN_GROUP_ALIGN_SIZE_M

    if local_counts.ndim == 1:
        # [E_local]: already aggregated over sources (ep_degree=1)
        ep_degree = 1
        num_local_experts = local_counts.shape[0]
        local_counts_flat = local_counts
    elif local_counts.ndim == 2:
        # [ep_size, E_local]: preserve per-source layout for correct regrouping
        ep_degree, num_local_experts = local_counts.shape
        local_counts_flat = local_counts.reshape(-1)
    else:
        raise ValueError(
            f"local_counts must have shape [E_local] or [ep_degree, E_local], got {tuple(local_counts.shape)}")

    n_tokens = tokens.shape[0]
    alignment = TOKEN_GROUP_ALIGN_SIZE_M

    # Compute padded max length
    x_padded_per_expert = n_tokens + num_local_experts * alignment
    padded_max_len = ((x_padded_per_expert + alignment - 1) // alignment) * alignment

    # Use the pure-PyTorch path for host tensors. The CPU accelerator reports
    # CPU tensors as "on accelerator", but Triton still requires a GPU driver.
    use_cpu = tokens.device.type == "cpu"
    counts_for_permute = local_counts_flat.cpu() if use_cpu else local_counts_flat
    with torch.no_grad():
        permuted_indices, m_sizes, _offsets = generate_permute_indices(
            counts_for_permute,
            num_local_experts,
            ep_degree,
            padded_max_len,
            alignment,
            use_cpu=use_cpu,
        )
    if not use_cpu:
        permuted_indices = permuted_indices.to(tokens.device)
        m_sizes = m_sizes.to(tokens.device)

    # Add padding row for out-of-bounds indices (index n_tokens -> zero row)
    tokens_padded = torch.vstack((tokens, tokens.new_zeros((tokens.shape[-1], ))))
    tokens_permuted = tokens_padded[permuted_indices, :]

    return tokens_permuted, permuted_indices, m_sizes, n_tokens


def unpermute_by_local_expert(
    expert_output: torch.Tensor,
    permuted_indices: torch.Tensor,
    n_tokens: int,
) -> torch.Tensor:
    """Reverse permute_by_local_expert: restore original token order and strip padding.

    Args:
        expert_output: [N_padded, H] from expert computation
        permuted_indices: [N_padded] index mapping from permute_by_local_expert
        n_tokens: original token count before alignment padding
    """
    # Scatter expert outputs back to original positions.
    # permuted_indices values range 0..n_tokens, where n_tokens is the zero-padding row.
    out_unpermuted = expert_output.new_zeros((n_tokens + 1, expert_output.shape[-1]))
    out_unpermuted[permuted_indices, :] = expert_output
    # Strip the zero-padding row to get [n_tokens, H]
    return out_unpermuted[:-1]


def combine_from_routed(
        expert_output: torch.Tensor,  # [N, H]
        top_scores: torch.Tensor,  # [T, K]
        token_indices_sorted: torch.Tensor,  # [N]
        top_k: int,
        score_apply: Literal["pre", "post"],
        combine_impl: Literal["weighted_sum", "legacy_bmm"],
        shape: tuple[int, int, int],  # (B, S, H)
) -> torch.Tensor:
    """Scatter-add expert outputs back to original token positions."""
    bsz, seqlen, hdim = shape
    T = bsz * seqlen

    # Create output tensor
    output = torch.zeros(T * top_k, hdim, dtype=expert_output.dtype, device=expert_output.device)

    # Place expert outputs back in unsorted order
    output[token_indices_sorted] = expert_output

    # Reshape to [T, K, H]
    output = output.reshape(T, top_k, hdim)

    if score_apply == "post":
        if combine_impl == "legacy_bmm":
            # Legacy reduction path retained as a debug option for model-family
            # verification. The weighted-sum path is the default.
            output = torch.bmm(
                top_scores.reshape(-1, 1, top_k).float(),
                output.float(),
            ).to(expert_output.dtype).squeeze(1)
        else:
            # Match the runtime HF grouped-mm path: apply routing weights per
            # token-slot sample, then reduce over top-k.
            output = (output.float() * top_scores.reshape(T, top_k, 1).float()).sum(dim=1).to(expert_output.dtype)
    else:
        # Scores already applied pre-experts, just sum over top_k
        output = output.sum(dim=1)

    return output.reshape(bsz, seqlen, hdim)


# ---------------------------------------------------------------------------
# AutoEPMoELayer
# ---------------------------------------------------------------------------


class AutoEPMoELayer(nn.Module):
    """Drop-in replacement for HF MoE blocks with Expert Parallelism support."""

    _is_autoep_layer = True  # Marker for AutoTP skip handshake

    def __init__(
        self,
        spec: MoELayerSpec,
        source_module: nn.Module,
        ep_size: int,
        ep_rank: int,
        config: AutoEPConfig,
    ) -> None:
        super().__init__()

        self.model_family = spec.model_family
        self.return_router_logits = spec.return_router_logits
        self.router_logits_capture_target = spec.router_logits_capture_target
        self.router_logits_capture_index = spec.router_logits_capture_index
        self.router_logits_capture_mode = spec.router_logits_capture_mode
        self.moe_output_shape = spec.moe_output_shape
        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
        self.num_experts = spec.num_experts
        self.num_local_experts = spec.num_experts // ep_size
        self.hidden_size = spec.hidden_size
        self.ep_group_name = f"ep_size_{ep_size}"
        self.ep_group = None  # Set by set_deepspeed_parallelism()
        self.folding_group_handles = None
        self.tp_group = None
        resolved_config = resolve_autoep_config_defaults(config, spec.model_family)
        self.validate_folding_routing = bool(resolved_config.validate_folding_routing)

        # Router: copy gate weights from source
        source_gate = getattr(source_module, spec.router_name)
        source_gate_bias = getattr(source_gate, 'bias', None)
        source_ecb = getattr(source_gate, 'e_score_correction_bias', None)
        unsupported_router_biases = [
            getattr(source_gate, bias_name, None) for bias_name in spec.unsupported_router_bias_names
        ]
        if not spec.supports_expert_bias and resolved_config.load_balance_coeff is not None:
            raise ValueError(f"AutoEP preset '{spec.model_family}' does not support load_balance_coeff/expert_bias "
                             "yet. Set load_balance_coeff=None.")
        with _gather_source_zero_params([source_gate.weight, source_gate_bias, source_ecb,
                                         *unsupported_router_biases]):
            for bias_name, router_bias in zip(spec.unsupported_router_bias_names, unsupported_router_biases):
                if router_bias is None:
                    continue
                if torch.is_tensor(router_bias) and torch.count_nonzero(router_bias.detach()).item() == 0:
                    continue
                raise ValueError(f"AutoEP preset '{spec.model_family}' does not support nonzero router bias "
                                 f"'{bias_name}' yet.")
            self.router = TokenChoiceTopKRouter(
                dim=spec.hidden_size,
                num_experts=spec.num_experts,
                num_expert_groups=spec.num_expert_groups,
                num_limited_groups=spec.num_limited_groups,
                top_k=spec.top_k,
                score_func=spec.score_func,
                route_norm=route_norm,
                route_scale=spec.route_scale,
                gate_bias=spec.gate_bias,
                group_score_func=spec.group_score_func,
            )
            # Copy gate weights
            _copy_parameter_data(self.router.gate.weight, source_gate.weight)
            self.router.gate.weight.requires_grad_(source_gate.weight.requires_grad)
            if spec.gate_bias and source_gate_bias is not None:
                _copy_parameter_data(self.router.gate.bias, source_gate_bias)
                self.router.gate.bias.requires_grad_(source_gate_bias.requires_grad)

            # Copy pre-trained score correction bias (DeepSeek-V3/Moonlight noaux_tc routing)
            if source_ecb is not None and isinstance(source_ecb, nn.Parameter):
                self.router.e_score_correction_bias = nn.Parameter(source_ecb.data.clone(),
                                                                   requires_grad=source_ecb.requires_grad)
                logger.info('AutoEP: copied e_score_correction_bias from source gate '
                            '(shape=%s)', source_ecb.shape)

        # Alias router under the name OutputRecorder expects (layer_name if provided),
        # but only when OutputRecorder captures from the router child and the alias is safe.
        alias_target = spec.router_logits_capture_layer_name or spec.router_name
        if spec.router_logits_capture_target == "router" and alias_target != "router":
            if "." in alias_target or alias_target in ("experts", "shared_experts") or hasattr(self, alias_target):
                logger.warning(f"Skipping router alias '{alias_target}' to avoid name collision.")
            else:
                setattr(self, alias_target, self.router)

        # Experts: extract local expert weights
        w1, w2, w3 = repack_expert_weights(
            experts_source=getattr(source_module, spec.experts_name),
            spec=spec,
            ep_rank=ep_rank,
            ep_size=ep_size,
        )
        w1_requires_grad, w2_requires_grad, w3_requires_grad = repack_expert_requires_grad_flags(
            experts_source=getattr(source_module, spec.experts_name),
            spec=spec,
            ep_rank=ep_rank,
            ep_size=ep_size,
        )
        self.experts = GroupedExperts(
            dim=spec.hidden_size,
            hidden_dim=spec.ffn_hidden_size,
            num_experts=self.num_local_experts,
            use_grouped_mm=config.use_grouped_mm,
            disable_triton_grouped_mm=config.disable_triton_grouped_mm,
        )
        _copy_parameter_data(self.experts.w1, w1)
        _copy_parameter_data(self.experts.w2, w2)
        _copy_parameter_data(self.experts.w3, w3)
        self.experts.w1.requires_grad_(w1_requires_grad)
        self.experts.w2.requires_grad_(w2_requires_grad)
        self.experts.w3.requires_grad_(w3_requires_grad)

        self.shared_experts = getattr(source_module, spec.shared_experts_name,
                                      None) if spec.has_shared_experts else None
        self.shared_experts_gate = getattr(source_module, spec.shared_experts_gate_name,
                                           None) if spec.shared_experts_gate_name else None

        # Mark expert params for EDP gradient reduction
        for param in self.experts.parameters():
            param.allreduce = False
            param.group_name = self.ep_group_name
            param.ds_zero_placement_family = "autoep_expert"
            param.ds_zero_partition_group_name = self.ep_group_name

        # Mark shared expert and router params for global DP reduction.
        # The router runs redundantly on every TP peer and its gradient is
        # rebuilt into a replicated full view by the restore all-gather, so it
        # is tagged as the replicated family (AVERAGE TP reduction); a SUM would
        # double it under tp_size=2. See mark_autoep_folding_router_parameter.
        for param in self.router.parameters():
            param.allreduce = True
            mark_autoep_folding_router_parameter(param)
            param.ds_zero_placement_family = "replicated"
        if self.shared_experts is not None:
            for param in self.shared_experts.parameters():
                param.allreduce = True
                param.ds_zero_placement_family = "replicated"
        if self.shared_experts_gate is not None:
            for param in self.shared_experts_gate.parameters():
                param.allreduce = True
                param.ds_zero_placement_family = "replicated"

        # Load balancing buffers
        self.load_balance_coeff = resolved_config.load_balance_coeff
        buf_device = source_gate.weight.device
        if self.load_balance_coeff is not None:
            self.register_buffer(
                "expert_bias",
                torch.zeros(spec.num_experts, dtype=torch.float32, device=buf_device),
                persistent=True,
            )
        else:
            self.expert_bias = None
        self.register_buffer(
            "tokens_per_expert",
            torch.zeros(spec.num_experts, dtype=torch.float32, device=buf_device),
            persistent=False,
        )

        # Router-logit cache
        self._cached_router_logits = None
        self._register_logit_hook()

    def _register_logit_hook(self):
        """Register a forward hook that caches gate logits for OutputRecorder capture."""
        if self.router_logits_capture_target != "router":
            return

        def hook_fn(module, input, output):
            x = input[0]  # [T, H]
            logits = module.gate(x)  # [T, E_global]
            if self.router_logits_capture_mode == "post_score":
                if self.router.score_func == "softmax":
                    logits = torch.softmax(logits.float(), dim=-1).to(logits.dtype)
                elif self.router.score_func == "sigmoid":
                    logits = torch.sigmoid(logits.float()).to(logits.dtype)
            self._cached_router_logits = logits

        self.router.register_forward_hook(hook_fn)

    def set_deepspeed_parallelism(
        self,
        use_data_before_expert_parallel_: bool = False,
        folding_group_handles=None,
    ) -> None:
        """Bind EP group handle to this module."""
        from deepspeed.utils import groups
        from deepspeed.utils.bwc import bwc_pipeline_parallel_world_size

        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
            self.ep_rank = dist.get_rank(group=self.ep_group)
            return

        if self.ep_group_name not in groups._get_expert_parallel_group_dict():
            mp_size = max(
                getattr(groups, '_get_model_parallel_world_size', lambda: 1)(),
                getattr(groups, '_get_sequence_parallel_world_size', lambda: 1)(),
            )
            mp_mode = "tp" if getattr(groups, '_get_model_parallel_world_size', lambda: 1)() > 1 else "sp"
            pp_size = 1 if groups.mpu is None else bwc_pipeline_parallel_world_size(groups.mpu)
            groups._create_expert_and_data_parallel(
                expert_parallel_size_=self.ep_size,
                mp_size=mp_size,
                pp_size=pp_size,
                mp_mode=mp_mode,
                use_data_before_expert_parallel_=use_data_before_expert_parallel_,
            )
        self.ep_group = groups._get_expert_parallel_group(self.ep_group_name)

    def forward(
        self,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
        """Forward pass.

        Args:
            hidden_states: [B, S, H]

        Returns:
            [B, S, H] or ([B, S, H], [T, E]) if return_router_logits.
            Some HF MoE contracts return ([T, H], [T, E]) instead.
        """
        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))

        # Accumulate expert utilization
        with torch.no_grad():
            self.tokens_per_expert.add_(ro.num_tokens_per_expert)

        # Reorder tokens into expert-contiguous order.
        token_indices_sorted = torch.argsort(ro.selected_experts.view(-1), stable=True)
        top_scores_sorted = ro.top_scores.view(-1)[token_indices_sorted]
        expert_indices_sorted = ro.selected_experts.reshape(-1).index_select(0, token_indices_sorted)

        folded_tp = self.folding_group_handles is not None and self.folding_group_handles.spec.tp_size > 1
        restore_ctx = None
        if folded_tp:
            from deepspeed.moe.ep_tp_dispatch import (
                RoutedAssignmentPayload,
                assignment_ordinals_by_expert,
                assert_tp_payload_consistent,
                dispatch_counters,
                partition_assignments,
                restore_combined,
            )
            payload = RoutedAssignmentPayload(
                token_indices=(token_indices_sorted // self.top_k).to(torch.long),
                expert_indices=expert_indices_sorted.to(torch.long),
                assignment_indices=assignment_ordinals_by_expert(expert_indices_sorted.to(torch.long)),
                capacity_slots=(token_indices_sorted % self.top_k).to(torch.long),
                combine_weights=top_scores_sorted
                if self.score_apply == "post" else torch.ones_like(top_scores_sorted),
                drop_mask=torch.zeros_like(top_scores_sorted, dtype=torch.bool),
                pad_mask=torch.zeros_like(top_scores_sorted, dtype=torch.bool),
                input_splits=[0 for _ in range(self.ep_size)],
                output_splits=[0 for _ in range(self.ep_size)],
                extra={
                    "destination_ranks": (expert_indices_sorted // self.num_local_experts).to(torch.long),
                    "top_scores": top_scores_sorted,
                    "num_tokens": torch.tensor(bsz * seqlen, device=hidden_states.device, dtype=torch.long),
                },
            )
            if self.validate_folding_routing:
                assert_tp_payload_consistent(payload,
                                             tp_group=self.tp_group,
                                             tp_size=self.folding_group_handles.spec.tp_size)
            tp_rank = dist.get_rank(group=self.tp_group)
            local_payload, restore_ctx = partition_assignments(payload,
                                                               tp_group=self.tp_group,
                                                               tp_rank=tp_rank,
                                                               tp_size=self.folding_group_handles.spec.tp_size)
            token_indices_for_compute = token_indices_sorted.index_select(0, restore_ctx.local_indices)
            top_scores_for_compute = top_scores_sorted.index_select(0, restore_ctx.local_indices)
            expert_indices_for_plan = local_payload.expert_indices
        else:
            token_indices_for_compute = token_indices_sorted
            top_scores_for_compute = top_scores_sorted
            expert_indices_for_plan = expert_indices_sorted

        routed_input = x[token_indices_for_compute // self.top_k]  # [N, H]
        routed_input = apply_scores_before_experts_if_enabled(routed_input,
                                                              top_scores_for_compute,
                                                              score_apply=self.score_apply)

        if self.ep_size == 1:
            # No AllToAll needed - local computation only
            local_counts = ro.num_tokens_per_expert

            routed_input_permuted, perm_indices, aligned_counts, n_tokens = permute_by_local_expert(
                routed_input, local_counts)
            expert_output = self.experts(routed_input_permuted, aligned_counts)
            expert_output = unpermute_by_local_expert(expert_output, perm_indices, n_tokens)
        else:
            # EP dispatch/compute/combine
            if folded_tp:
                plan = compute_split_plan_from_expert_indices(
                    expert_indices=expert_indices_for_plan,
                    num_experts=self.num_experts,
                    ep_size=self.ep_size,
                    num_local_experts=self.num_local_experts,
                    ep_group=self.ep_group,
                )
            else:
                plan = compute_split_plan(
                    selected_experts=ro.selected_experts,
                    num_experts=self.num_experts,
                    ep_size=self.ep_size,
                    num_local_experts=self.num_local_experts,
                    ep_group=self.ep_group,
                    num_tokens_per_expert=ro.num_tokens_per_expert,
                )

            routed_input = _AllToAllV.apply(self.ep_group, routed_input, plan.input_splits, plan.output_splits)

            routed_input, perm_indices, aligned_counts, n_tokens = permute_by_local_expert(
                routed_input, plan.local_counts_by_source)
            expert_output = self.experts(routed_input, aligned_counts)
            expert_output = unpermute_by_local_expert(expert_output, perm_indices, n_tokens)

            expert_output = _AllToAllV.apply(self.ep_group, expert_output, plan.output_splits, plan.input_splits)

        if folded_tp:
            output = restore_combined(expert_output,
                                      restore_ctx,
                                      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,
                top_scores=ro.top_scores,
                token_indices_sorted=token_indices_sorted,
                top_k=self.top_k,
                score_apply=self.score_apply,
                combine_impl=self.combine_impl,
                shape=(bsz, seqlen, hdim),
            )

        if self.moe_output_shape == "flat":
            output = output.reshape(-1, hdim)
            shared_expert_input = x
        elif self.shared_experts_gate is not None:
            shared_expert_input = x
        else:
            shared_expert_input = hidden_states

        if self.shared_experts is not None:
            shared_expert_output = self.shared_experts(shared_expert_input)
            if self.shared_experts_gate is not None:
                shared_expert_gate = torch.sigmoid(self.shared_experts_gate(shared_expert_input))
                shared_expert_output = shared_expert_gate * shared_expert_output
            if shared_expert_output.shape != output.shape:
                shared_expert_output = shared_expert_output.reshape_as(output)
            output = output + shared_expert_output

        if self.return_router_logits:
            logits = self._cached_router_logits
            self._cached_router_logits = None
            return output, logits

        self._cached_router_logits = None
        return output