Fix broken Llama4 accuracy in MoE part

[nvpohanh]

Llama4 accuracy is broken by a bug in
#39501 . It forgot to transpose the router_scores before applying it to routed_in, causing Llama4 to generate garbage output.

This PR fixes that issue by adding back the transpose() and adding some comments explaining why the transpose() is needed.

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,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? 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, and
  here are tips on formatting docstrings.
• 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.

@ArthurZucker please review since #39501 was made by you. Thanks!

Accuracy tests

Test script:

from transformers import pipeline
import torch

model_id = "meta-llama/Llama-4-Scout-17B-16E-Instruct"

pipe = pipeline(
    "text-generation",
    model=model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)

output = pipe("Roses are red,", max_new_tokens=200)

print(output)

Before the fix on H200:

[{'generated_text': 'Roses are red, 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 
8 8 8 8 8 8 8 8 8 8 8 8 8 8 8'}]

After the fix on H200:

[{'generated_text': "Roses are red, violets are blue, and here are some Valentine's Day-themed books to read with your boo! \n Celebrate Black History Month by reading books by Black authors and about Black culture!  These books are for kids and teens! \n These picture books celebrate Black History Month and are perfect for reading with your kids!  These books highlight Black history, cu... \n These books are all winners or nominees for major literary awards, including the Pulitzer Prize, the National Book Award, and t... \n 100 Book Challenge 2024: Read a Book Set in Another Country \n Check out these books that are set in different countries around the world!  Read a book set in another country and see what you... \n Explore the diverse world through these books that are set in different countries and cultures!  Read a book set in another coun... \n 100 Book Challenge 2024: Read a Book Written by an Author of Color \n Check out these books written by authors of color!"}]

[Rocketknight1]

cc @ArthurZucker!

[FThompsonAWS]

Confirmed with the given repro that this issue also affects Llama4 CPU execution on transformers v4.54, v4.55, and v4.56. The output is accurate on v4.53.

[nvpohanh]

@ArthurZucker Could you review this? Thanks!

[Cyrilvallez]

Indeed, this got lost! Thanks a lot for the fix!!

[github-actions]

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

run-slow: llama4