Add back is_batch_shard_by_expert

[NuojCheng]

Description

This PR recover the is_batch_shard_by_expert logic fixing previous full EP decoding issue b/501537579. We introduced a decode_batch_moe logical rule to specifically deal with the decoding case in moe component when batch size is one.

Tests

• Performance test: https://paste.googleplex.com/4690855089799168. We get the same training performance (~81 TFLOP/s/dev compared to Update config correcting rule and log_config #3645)

• Inference vllm test: https://paste.googleplex.com/5221471321456640

• Inference decode test: https://paste.googleplex.com/6557626352664576

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

❌ Patch coverage is 66.66667% with 7 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/moe.py	66.66%	4 Missing and 3 partials ⚠️

📢 Thoughts on this report? Let us know!

[gobbleturk]

longer term this logic should be separated into its own function following go/small-functions, but for now this is probably safer as a rollback and I am currently working on refactoring/cleaning up this file

[RissyRan]

Thanks Nuojin for the change! I was comparing the before/after, and noticed this change. Could you help cross check?

• https://screenshot.googleplex.com/XuuC5bCYnWpgzd9
• Diff: https://diff.googleplex.com/#key=B3n5CBidn5Wc

[NuojCheng]

Thanks Nuojin for the change! I was comparing the before/after, and noticed this change. Could you help cross check?

• https://screenshot.googleplex.com/XuuC5bCYnWpgzd9
• Diff: https://diff.googleplex.com/#key=B3n5CBidn5Wc

The previous hint ("activation_batch_no_exp", "activation_length_no_exp", None) was not correct. It always makes the output replicated along expert physical axis. The new implementation is a more memory-efficient way.

This change of logic only impacts performance under explicit shard mode, which has not been included in benchmarking yet. One exception is batch split, which has their own code path deepseek_batchsplit.py and should not get impacted.