fix(llama4): align MoE interface for EP/TP compatibility
Author: Chao1Han
Pull request: yes

# What does this PR do?
Llama4 router and experts use a non-standard interface that is incompatible with the EP/TP hooks in tensor_parallel.py, causing ValueError when running with device_map or tensor parallelism.

- Router now returns (router_logits, top_k_weights, top_k_index) matching what RouterParallel expects
- Experts forward takes (hidden_states, top_k_index, top_k_weights) matching MoeTensorParallelExperts hooks
- MoE forward simplified accordingly
<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

## reproduce the issue
```python
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-4-Scout-17B-16E", tp_plan="auto")
# ValueError at RouterParallel._prepare_output_fn
```

## Code Agent Policy

The Transformers repo is currently being overwhelmed by a large number of PRs and issue comments written by
code agents. We are currently bottlenecked by our ability to review and respond to them. As a result, 
**we ask that new users do not submit pure code agent PRs** at this time. 
You may use code agents in drafting or to help you diagnose issues. We'd also ask autonomous "OpenClaw"-like agents
not to open any PRs or issues for the moment.

PRs that appear to be fully agent-written will probably be closed without review, and we may block users who do this
repeatedly or maliciously. 

This is a rapidly-evolving situation that's causing significant shockwaves in the open-source community. As a result, 
this policy is likely to be updated regularly in the near future. For more information, please read [`CONTRIBUTING.md`](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md).

- [x] I confirm that this is not a pure code agent PR.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker @Cyrilvallez
- vision models: @yonigozlan @molbap
- audio models: @eustlb @ebezzam @vasqu
- multimodal models: @zucchini-nlp
- graph models: @clefourrier

Library:

- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- continuous batching: @remi-or @ArthurZucker @McPatate
- pipelines: @Rocketknight1
- tokenizers: @ArthurZucker and @itazap
- trainer: @SunMarc
- attention: @vasqu @ArthurZucker @CyrilVallez
- model loading (from pretrained, etc): @CyrilVallez
- distributed: @3outeille @ArthurZucker
- CIs: @ydshieh

Integrations:

- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization: @SunMarc
- kernels: @drbh
- peft: @BenjaminBossan @githubnemo

Devices/Backends:

- AMD ROCm: @ivarflakstad
- Intel XPU: @IlyasMoutawwakil
- Ascend NPU: @ivarflakstad 

Documentation: @stevhliu

Research projects are not maintained and should be taken as is.

 -->


Comments:
Hi @Rocketknight1. Could you please take a look at this PR? I do have a requirement to run EP with Llama 4, but the current implementation of the Llama 4 router does not match the assumptions made by `MoeTensorParallelExperts`.
I believe the current test coverage is already sufficient to ensure that this change will not introduce any BC issues. If needed, I can add an EP test case as well.

cc @cyrilvallez for TP I think! 

Hi @Cyrilvallez . Could you please take a look?

cc @ArthurZucker here since you originally introduced it for EP from git history!

run-slow: llama4

[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/26646303691)

This comment contains `run-slow`, running the specified jobs: 

models: ["models/llama4"]
quantizations: []

## CI Results
[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/26646303691)

### Commit Info
| Context | Commit | Description |
|---------|--------|-------------|
| RUN | [47eb83a3](https://github.com/huggingface/transformers/commit/47eb83a305c7339d722b0893e88d78f019191e96) | workflow commit (merge commit) |
| PR | [c2fd2331](https://github.com/Chao1Han/transformers/commit/c2fd2331aacb6961a67846d744fc7d31bd38095e) | branch commit (from PR) |
| main | [96d3ecce](https://github.com/huggingface/transformers/commit/96d3eccef9e156a82e218cadd345158e4bba435e) | base commit (on `main`) |

⚠️ **Model CI failed to report results**

The test failure analysis could not be completed. Please check the [workflow run](https://github.com/huggingface/transformers/actions/runs/26646303691) for details.


Thanks @ArthurZucker  ! I verified the changes don't affect numerics: I ran the new vs original Llama4TextExperts/Router with identical fp32 weights across 3 random seeds — outputs are bit-exact(max abs diff = 0.0). The refactor only moves the routing-weight multiply from the MoE wrapper into the experts forward; the bmm computation is unchanged. 
And for CI failure (exit code 137, "Executing the custom container implementation failed"), looks like an infra OOM on the self-hosted runner — unrelated to this PR. Mind re-triggering?


run-slow: llama4

**Repro** (8x B200, Llama-4-Scout-17B-16E-Instruct in BF16, on this PR's branch):

`repro.py`:
```python
import os, torch, torch.distributed as dist
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.distributed import DistributedConfig

MODEL = "meta-llama/Llama-4-Scout-17B-16E-Instruct"
tok = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(
    MODEL,
    dtype="bfloat16",
    attn_implementation="eager",
    tp_plan="auto",
    distributed_config=DistributedConfig(enable_expert_parallel=True),
).eval()
inputs = tok("The capital of France is", return_tensors="pt").to(f"cuda:{rank}")
with torch.no_grad():
    model(**inputs)
```

```bash
torchrun --nproc-per-node=8 repro.py
```
```
File "transformers/models/llama4/modeling_llama4.py", line 77, in forward
    gate_up = torch.bmm(routed_in, self.gate_up_proj)
RuntimeError: CUDA error: CUBLAS_STATUS_EXECUTION_FAILED when calling cublasGemmStridedBatchedEx
```

Can you investigate? 

Thanks for the repro!
It used self.num_experts (the global count) to size router_scores and the bmm batch dim, but under EP gate_up_proj is sharded to the local expert count. When the module's num_experts attribute isn't updated to the local value, bmm gets mismatched batch dims.
Fixed by deriving the expert count from gate_up_proj.shape[0] (always the physical local count) and adding one extra scatter column to absorb the sentinel before dropping it.
Can you help rerun llama4 ci?

Hi @ArthurZucker , Could you please help kick off the CI and take another look at the review? I’d really appreciate it.

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_46024). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: llama4