Diff Coverage

def _differentiable_all_reduce(data: Tensor, op: str, group: Any) -> Tensor:
    """Autograd-aware all-reduce: backward is an identity, the loss is scaled by world size."""
    if not data.is_contiguous() or data.storage_offset() != 0:
        data = data.contiguous()
    reduce_op = _REDUCE_OP_MAP.get(op, torch.distributed.ReduceOp.SUM)
    return dist_func.all_reduce(data, op=reduce_op, group=group)

__all__ = [
    "vocab_parallel_cross_entropy_local",
    "distributed_log_softmax",

    """
    max_local = logits_local.max(dim=dim, keepdim=True).values

    group = mesh.get_group(mesh_dim)
    max_global = _differentiable_all_reduce(max_local, op="max", group=group)

    exp_local = (logits_local - max_global).exp()

    sum_local = exp_local.sum(dim=dim, keepdim=True)

    exp_local = (logits_local - max_global).exp()

    sum_local = exp_local.sum(dim=dim, keepdim=True)

    sum_global = _differentiable_all_reduce(sum_local, op="sum", group=group)

    log_softmax = logits_local - max_global - sum_global.log()

    return log_softmax

        )

        if reduction == "mean":
            group = mesh.get_group(mesh_dim)
            total_loss = _differentiable_all_reduce(loss, op="sum", group=group)
            total_weight_sum = _differentiable_all_reduce(
                total_weight, op="sum", group=group
            )

            ctx.save_for_backward(

                return torch.tensor(float('nan'), dtype=total_loss.dtype, device=total_loss.device)
            return total_loss / total_weight_sum
        if reduction == "sum":
            group = mesh.get_group(mesh_dim)
            total_loss = _differentiable_all_reduce(loss, op="sum", group=group)

            ctx.save_for_backward(
                log_probs_local,
                target,

def _is_npu_tensor(tensor: torch.Tensor) -> bool:
    """Return whether a tensor is stored on an Ascend NPU."""
    return tensor.device.type == "npu"


# ---------------------------------------------------------------------------
# Grouped expert computation kernels

    parts = []
    offset = 0
    for e, n in enumerate(counts_list):
        if n == 0:
            continue
        x_e = x[offset:offset + n]
        h = F.silu(x_e @ w1[e].T) * (x_e @ w3[e].T)
        if scores is not None:
            h = h * scores[offset:offset + n].unsqueeze(-1)
        parts.append(h @ w2[e].T)
        offset += n
    if not parts:
        # No routed tokens: return a grad-connected zero (not ``zeros_like``).
        return x * 0.0
    return torch.cat(parts, dim=0)


def _run_experts_grouped_mm_gpu(

    Returns:
        Expert output of shape ``[total_routed_tokens, dim]``.
    """
    # offs: cumulative split offsets (int32) for torch._grouped_mm.
    offs = torch.cumsum(num_tokens_per_expert[:-1], dim=0).to(torch.int32)
    # w1/w3 stored as [num_experts, hidden_dim, dim]; grouped_mm expects
    # [num_experts, dim, hidden_dim], so transpose the inner two dims.
    w1_t = w1.transpose(1, 2).contiguous()  # [num_experts, dim, hidden_dim]
    w3_t = w3.transpose(1, 2).contiguous()
    w2_t = w2.transpose(1, 2).contiguous()  # [num_experts, hidden_dim, dim]
    h1 = torch._grouped_mm(x, w1_t, offs=offs)  # pylint: disable=protected-access
    h3 = torch._grouped_mm(x, w3_t, offs=offs)  # pylint: disable=protected-access
    h = F.silu(h1) * h3
    if scores is not None:
        h = h * scores.unsqueeze(-1)
    return torch._grouped_mm(h, w2_t, offs=offs)  # pylint: disable=protected-access


def _run_experts_grouped_mm_npu(
    w1: torch.Tensor,

    Returns:
        Expert output of shape ``[total_routed_tokens, dim]``.
    """
    if x.shape[0] == 0:
        parameter_zero = (w1.sum() + w2.sum() + w3.sum()).to(x.dtype) * 0.0
        return x + parameter_zero

    gate_up_weight = torch.cat((w1, w3), dim=1)
    output = npu_grouped_swiglu(
        x,
        gate_up_weight,
        w2,
        num_tokens_per_expert,

        gate_up_weight,
        w2,
        num_tokens_per_expert,
    )
    if scores is not None:
        output = output * scores.to(output.dtype).unsqueeze(-1)
    return output


# ---------------------------------------------------------------------------
# FeedForward — shared expert / standard SwiGLU FFN

            dim: Input embedding dimension.
            hidden_dim: Intermediate hidden dimension.
            bias: Whether to add a learnable bias.  Defaults to ``False``.
        """
        super().__init__()
        self.w1 = nn.Linear(dim, hidden_dim, bias=bias)
        self.w2 = nn.Linear(hidden_dim, dim, bias=bias)
        self.w3 = nn.Linear(dim, hidden_dim, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Compute SwiGLU feed-forward output.

        Returns:
            Output tensor with the same leading shape and last dimension ``dim``.
        """
        return self.w2(F.silu(self.w1(x)) * self.w3(x))


# ---------------------------------------------------------------------------
# GroupedExperts

        w3 = self.w3.to_local() if isinstance(self.w3, DTensor) else self.w3

        if not self.use_grouped_mm:
            return _run_experts_for_loop(w1, w2, w3, x, num_tokens_per_expert, scores)
        if _is_npu_tensor(x):
            return _run_experts_grouped_mm_npu(w1, w2, w3, x, num_tokens_per_expert, scores)
        if x.device.type == "cuda":
            return _run_experts_grouped_mm_gpu(w1, w2, w3, x, num_tokens_per_expert, scores)
        return _run_experts_for_loop(w1, w2, w3, x, num_tokens_per_expert, scores)


# ---------------------------------------------------------------------------
# TokenChoiceTopKRouter

            route_scale: Scalar multiplier applied to routing scores.
        """
        super().__init__()
        if score_func not in ("sigmoid", "softmax"):
            raise ValueError(
                f"score_func must be 'sigmoid' or 'softmax', got '{score_func}'."
            )
        if num_expert_groups is not None and num_limited_groups is None:
            raise ValueError(
                "num_limited_groups must be set when num_expert_groups is not None."
            )
        self.gate = nn.Linear(dim, num_experts, bias=False)
        self.num_experts = num_experts

        if self.score_func == "sigmoid":
            scores = torch.sigmoid(scores)
        else:
            scores = F.softmax(scores, dim=-1)

        if self.route_scale != 1.0:
            scores = scores * self.route_scale

        # Node-limited routing — mask out low-scoring expert groups.
        scores_for_topk = scores
        if self.num_expert_groups is not None:
            scores_for_topk = self._get_node_limited_routing_scores(scores)

        # Add expert bias only for selection; returned scores remain unbiased.
        scores_with_bias = scores_for_topk
        if expert_bias is not None:

        Returns:
            Scores with non-selected groups masked to ``-inf``.
        """
        num_tokens, num_experts = scores.shape
        experts_per_group = num_experts // self.num_expert_groups
        group_scores = scores.view(
            num_tokens, self.num_expert_groups, experts_per_group
        ).max(dim=-1).values  # [num_tokens, num_groups]

        _, selected_groups = group_scores.topk(self.num_limited_groups, dim=-1)

        mask = scores.new_zeros(num_tokens, self.num_expert_groups)
        mask.scatter_(1, selected_groups, 1.0)
        mask = (
            mask.unsqueeze(-1)
            .expand(num_tokens, self.num_expert_groups, experts_per_group)
            .reshape(num_tokens, num_experts)
        )
        return scores.masked_fill(mask == 0, float("-inf"))


# ---------------------------------------------------------------------------
# Load-balance auxiliary loss

    Returns:
        Scalar loss tensor. Returns 0.0 for empty input.
    """
    import torch.distributed as dist  # pylint: disable=C0415

    num_tokens, top_k = top_scores.shape

    if num_tokens == 0:
        return torch.tensor(0.0, dtype=top_scores.dtype, device=top_scores.device)

    flat_experts = selected_experts.flatten()
    flat_scores = top_scores.flatten()

    expert_fraction = torch.zeros(
        num_experts, dtype=top_scores.dtype, device=top_scores.device
    )
    expert_fraction.scatter_add_(0, flat_experts, torch.ones_like(flat_scores))

    num_sub_sequence = 1
    if sequence_partition_group is not None:
        num_sub_sequence = dist.get_world_size(sequence_partition_group)
        dist.all_reduce(expert_fraction, group=sequence_partition_group)

    expert_fraction = expert_fraction / (num_tokens * num_sub_sequence * top_k)

    expert_prob = torch.zeros(
        num_experts, dtype=top_scores.dtype, device=top_scores.device
    )
    expert_prob.scatter_add_(0, flat_experts, flat_scores)
    expert_prob = expert_prob / (num_tokens * num_sub_sequence)

    loss = num_experts * (expert_fraction * expert_prob).sum()

    return loss


# ---------------------------------------------------------------------------
# MoEAuxLossAutoScaler — gradient injection for auxiliary loss

        Returns:
            The ``output`` tensor, identical in value but with an added
            autograd edge to ``aux_loss``.
        """
        ctx.save_for_backward(aux_loss)
        return output

    @staticmethod
    def backward(ctx: Any, grad_output: torch.Tensor) -> tuple:
        """Inject scaled aux_loss gradient into the backward chain.

        Returns:
            Tuple of (grad_output unchanged, scaled aux_loss gradient).
        """
        (aux_loss,) = ctx.saved_tensors
        if MoEAuxLossAutoScaler.main_loss_backward_scale is None:
            MoEAuxLossAutoScaler.main_loss_backward_scale = torch.tensor(
                1.0, device=aux_loss.device,
            )
        aux_loss_backward_scale = MoEAuxLossAutoScaler.main_loss_backward_scale
        scaled_aux_loss_grad = torch.ones_like(aux_loss) * aux_loss_backward_scale
        return grad_output, scaled_aux_loss_grad

    @staticmethod
    def set_loss_scale(scale: torch.Tensor) -> None:
        """Set the gradient scale for auxiliary loss.

        Args:
            scale: Tensor containing the loss scale value, typically
                ``1 / (num_microbatches * dp_size)`` or similar.
        """
        if MoEAuxLossAutoScaler.main_loss_backward_scale is None:
            MoEAuxLossAutoScaler.main_loss_backward_scale = scale
        else:
            MoEAuxLossAutoScaler.main_loss_backward_scale.copy_(scale)


# ---------------------------------------------------------------------------
# MoE orchestrator

        # Compute aux_loss early and attach it to top_scores via
        # MoEAuxLossAutoScaler so that backward through top_scores
        # also triggers aux_loss gradient (injected into router weights).
        if self.load_balance_coeff is not None:
            lb_loss = self.load_balance_coeff * _compute_load_balance_loss(
                top_scores, selected_experts, self.num_experts,
                sequence_partition_group=self.sequence_partition_group,
            )
            # Apply AutoScaler *before* top_scores is used in the forward path

            )
            # Apply AutoScaler *before* top_scores is used in the forward path
            # so that the main backward through top_scores triggers aux_loss
            # gradient injection. Forward values are unchanged.
            top_scores = MoEAuxLossAutoScaler.apply(top_scores, lb_loss)
            self.last_aux_loss = lb_loss.detach()
        else:
            lb_loss = None
            self.last_aux_loss = None

        if self.score_before_experts:
            routed_x = routed_x * top_scores_sorted.unsqueeze(1)
            expert_out = self.experts(routed_x, num_tokens_per_expert, scores=None)
        else:
            expert_out = self.experts(routed_x, num_tokens_per_expert, scores=top_scores_sorted)

        # --- Shared expert (parallel with routed experts) ---
        shared_out = None
        if self.shared_expert is not None:
            shared_out = self.shared_expert(x_flat)


        # --- Scatter expert outputs back to token order ---
        out = self.unpermutation(expert_out, token_indices, num_tokens, dim)

        # --- Scatter expert outputs back to token order ---
        out = self.unpermutation(expert_out, token_indices, num_tokens, dim)

        if shared_out is not None:
            out = out + shared_out

        result = out.view(bs, seq_len, dim)

        # Attach auxiliary loss to the returned tensor for logging.

        result = out.view(bs, seq_len, dim)

        # Attach auxiliary loss to the returned tensor for logging.
        if lb_loss is not None:
            result._load_balance_loss = lb_loss  # pylint: disable=protected-access

        return result

    def update_expert_bias(

            >>> moe_layer.update_expert_bias(lr=1e-3, num_recomputations=2)
        """
        with torch.no_grad():
            if num_recomputations > 1:
                self.tokens_per_expert.div_(num_recomputations)
            avg = self.tokens_per_expert.float().mean()
            delta = lr * (avg - self.tokens_per_expert.float()).sign()
            delta = delta - delta.mean()
            self.expert_bias.data += delta

        >>>
        >>> # With activation checkpoint (forward executed twice):
        >>> update_expert_bias(moe_layer, lr=1e-3, num_recomputations=2)
    """
    moe.update_expert_bias(lr=lr, num_recomputations=num_recomputations)

        if dim_map == "None":
            return 0

        if isinstance(dim_map, str):
            rank = get_rank()
            rank_list = layout.mesh.get_rank_list_along_axis(dim_map)
            if rank in rank_list:
                return rank_list.index(rank)
            return 0

                    f"Seq dim is sharded by multiple axes {non_none_axes}. "
                    f"Using the last axis for split_id calculation."
                )
            axis_name = non_none_axes[-1]
            rank = get_rank()
            rank_list = layout.mesh.get_rank_list_along_axis(axis_name)
            if rank in rank_list:
                return rank_list.index(rank)

        from hyper_parallel.core.dtensor.device_mesh import _DEVICE_MESH_MAP  # pylint: disable=C0415
        from hyper_parallel.core.dtensor.dtensor import _LAYOUT_CACHE  # pylint: disable=C0415
        from hyper_parallel.core.dtensor.tensor_redistribution import _tensor_redistribution  # pylint: disable=C0415
        from hyper_parallel.core.utils.communication import EXISTING_COMM_GROUPS  # pylint: disable=C0415
        from hyper_parallel.core.pipeline_parallel._p2p import (  # pylint: disable=C0415
            _P2P_MULTI_STREAM_GROUPS,
        )

        EXISTING_COMM_GROUPS.clear()