allow expert_parallel wrapper to handel kwargs

[rakkit]

Currently, all MoE models rely on the same forward logic (_run_experts_for_loop or _run_experts_grouped_mm), which is hardcoded to use Swiglu.

This PR allows expert_parallel to accept args, allowing users flexibility to define custom expert models. For example, users could specify a different activation function and implement their own forward function, while still reusing the upstream expert_parallel logic:

@expert_parallel
def cunstom_experts_grouped_mm(
    w1: torch.Tensor,
    w2: torch.Tensor,
    w3: torch.Tensor,
    x: torch.Tensor,
    num_tokens_per_expert: torch.Tensor,
    act_fn: Callable | nn.Module,
) -> torch.Tensor:
    

[tianyu-l]

I'm doing a relatively major refactor in #1569
Would appreciate if you can check the new indices_permutation_wrapper is still OK for you to extend.

[rakkit]

thanks @tianyu-l ! Yes, conceptually I think new indices_permutation_wrapper also works for this extension.