Diff Coverage

def _local_expert_major_forward(experts, dispatched_states, local_expert_counts):
    """Run grouped experts for an already expert-major token stream."""
    grouped_forward = getattr(experts, "forward_expert_major", None)
    if not callable(grouped_forward):
        raise TypeError(
            f"{type(experts).__name__}: fused EP dispatch requires "
            "experts.forward_expert_major"
        )
    return grouped_forward(dispatched_states, local_expert_counts)


def _get_global_expert_count(module):
    """Return the model-level routed expert count for an MoE module."""

            raise ValueError(
                f"{type(module).__name__}: unsupported expert activation {hidden_act!r}; "
                "provide experts.act_fn or extend the EP activation registry"
            )
    if use_npu_moe_token_dispatch and not use_grouped_gemm:
        raise ValueError("fused EP token dispatch requires grouped GEMM")
    if use_npu_moe_token_dispatch and not callable(
        getattr(module.experts, "forward_expert_major", None)
    ):
        raise TypeError(
            f"{type(module.experts).__name__}: fused EP token dispatch requires "
            "experts.forward_expert_major"
        )
    module.experts.ep_act_fn = activation

    module.experts.ep_use_grouped_gemm = use_grouped_gemm
    # The forward write itself lives in the forward rewriter (05 §15.2.3:
    # the single MethodType/assignment site); this binder only sets the
    # companion attributes above.
    local_forward = (
        _local_expert_major_forward
        if use_npu_moe_token_dispatch
        else _local_swiglu_expert_forward
    )
    _install_bound_forward(module.experts, local_forward)


def _npu_moe_token_permute(
    tokens: torch.Tensor,

    tokens: torch.Tensor,
    indices: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Load the optional NPU fused permutation only when requested."""
    from hyper_parallel.components.functional.moe_token_permute import (  # pylint: disable=C0415
        moe_token_permute,
    )

    return moe_token_permute(tokens, indices)


def _npu_moe_token_unpermute(
    permuted_tokens: torch.Tensor,

    sorted_indices: torch.Tensor,
    probabilities: torch.Tensor,
) -> torch.Tensor:
    """Load the optional NPU fused unpermutation only when requested."""
    from hyper_parallel.components.functional.moe_token_unpermute import (  # pylint: disable=C0415
        moe_token_unpermute,
    )

    return moe_token_unpermute(permuted_tokens, sorted_indices, probabilities)


def _reorder_variable_chunks(
    tensor: torch.Tensor,

    chunk_sizes: list[int],
    order: list[int],
) -> torch.Tensor:
    """Reorder variable-length chunks without changing rows within a chunk."""
    chunks = tensor.split(chunk_sizes, dim=0)
    return torch.cat([chunks[index] for index in order], dim=0)


def _prepare_fused_ep_dispatch(
    hidden_states: torch.Tensor,

    ep_rank: int,
    ep_group: Any,
) -> _FusedEPDispatch:
    """Permute tokens with the NPU fused op and prepare EP split metadata."""
    flattened_states = hidden_states.reshape(-1, hidden_states.shape[-1])
    local_counts = torch.bincount(
        topk_indices.reshape(-1),
        minlength=global_expert_count,
    )
    global_counts = torch.empty(
        (ep_size, global_expert_count),
        dtype=local_counts.dtype,
        device=local_counts.device,
    )
    dist.all_gather_into_tensor(global_counts, local_counts, group=ep_group)

    expert_start = ep_rank * local_expert_count
    expert_end = expert_start + local_expert_count
    received_counts_by_source = global_counts[:, expert_start:expert_end]
    send_counts = local_counts.reshape(ep_size, local_expert_count).sum(dim=1).tolist()
    receive_counts = received_counts_by_source.sum(dim=1).tolist()
    local_expert_counts = received_counts_by_source.sum(dim=0)
    received_chunk_sizes = received_counts_by_source.reshape(-1).tolist()
    source_to_expert_order = [
        source * local_expert_count + expert
        for expert in range(local_expert_count)
        for source in range(ep_size)
    ]
    expert_to_source_order = [
        expert * ep_size + source
        for source in range(ep_size)
        for expert in range(local_expert_count)
    ]
    permuted_states, unpermute_indices = _npu_moe_token_permute(
        flattened_states,
        topk_indices.to(torch.int32),
    )
    return _FusedEPDispatch(
        states=permuted_states,
        unpermute_indices=unpermute_indices,
        send_counts=send_counts,
        receive_counts=receive_counts,

    topk_weights: torch.Tensor,
    ep_group: Any,
) -> torch.Tensor:
    """Run NPU fused token permutation around EP all-to-all and experts."""
    ep_size = ep_group.size()
    ep_rank = dist.get_rank(group=ep_group)
    local_expert_count = module.experts.local_expert_count
    global_expert_count = local_expert_count * ep_size
    dispatch = _prepare_fused_ep_dispatch(
        hidden_states,
        topk_indices,
        local_expert_count=local_expert_count,
        global_expert_count=global_expert_count,

        ep_size=ep_size,
        ep_rank=ep_rank,
        ep_group=ep_group,
    )
    received_states = ep_all_to_all(
        dispatch.states,
        dispatch.send_counts,
        dispatch.receive_counts,
        ep_group,

        dispatch.send_counts,
        dispatch.receive_counts,
        ep_group,
    )
    expert_major_states = _reorder_variable_chunks(
        received_states,
        dispatch.received_chunk_sizes,
        dispatch.source_to_expert_order,
    )
    expert_major_outputs = module.experts(
        expert_major_states,
        dispatch.local_expert_counts,
    )
    source_major_outputs = _reorder_variable_chunks(
        expert_major_outputs,
        [
            dispatch.received_chunk_sizes[index]
            for index in dispatch.source_to_expert_order

            for index in dispatch.source_to_expert_order
        ],
        dispatch.expert_to_source_order,
    )
    combined_outputs = ep_all_to_all(
        source_major_outputs.contiguous(),
        dispatch.receive_counts,
        dispatch.send_counts,
        ep_group,

        dispatch.receive_counts,
        dispatch.send_counts,
        ep_group,
    )
    output = _npu_moe_token_unpermute(
        combined_outputs,
        dispatch.unpermute_indices,
        topk_weights.to(combined_outputs.dtype),
    )
    return output.view_as(hidden_states)


def _prepare_ep_dispatch(
    hidden_states: torch.Tensor,

    expert_offset = ep_rank * local_expert_count

    output_shape = tuple(hidden_states.shape)
    topk_indices, topk_weights = router_fn(module, hidden_states)  # [T, K]
    if use_npu_moe_token_dispatch:
        return _fused_ep_routed_forward(
            module,
            hidden_states,
            topk_indices,
            topk_weights,