[TPU][Feature]: Request Fix for TPU MoE Pallas Kernel

[yuanfz98]

🚀 The feature, motivation and pitch

Hello,
The temp fix for moe_pallas will produce a long IR as dynamo will unroll the for loop in current implementation of MOE.
Could anyone give a robust moe pallas ?
#10963

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[yuanfz98]

Ref to #11764

[mgoin]

I added a basic kernel test to show that the moe_pallas kernel is valid for simple cases: #17496

However when I try to enable the kernel and run a full evaluation through a model, I start seeing correctness issues #18025