[Feature]: moe_align_block_size optimization opportunity for Expert Parallel case

[varun-sundar-rabindranath]

🚀 The feature, motivation and pitch

In the case of expert_parallel, moe_align_block_size initially considers all experts as valid and aligns all tokens appropriately. Before the function returns it marks the experts_ids that are not in the current GPU rank as -1 so the MoE matmuls could skip those blocks.
This is sub-optimal in memory and performance. The solution is to recognize/apply expert_map before or inside moe_align_block_size so we allocate less memory do less work.

Alternatives

No response

Additional context

Related bugfix PR - #19515

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.

[github-actions]

This issue has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this issue should remain open. Thank you!

[github-actions]

This issue has been automatically closed due to inactivity. Please feel free to reopen if you feel it is still relevant. Thank you!