Use single random seed across embedding chunks

[Shuwen-Fang]

Description

Use a single random seed (rngs.params()) across all embedding chunks in RoutedMoE to ensure tokens are routed to the same destination for each chunk when random routing is enabled. This ensures tokens are routed to the same destination for each chunk, and removes unnecessary sorts across embedding chunks, reducing step time from 712 ms (xprof trace before) down to 665 ms (xprof trace after).

Tests

Added test_moe_emb_chunking_random_routing in tests/unit/moe_test.py to verify that chunked and non-chunked RoutedMoE with random routing produce identical outputs.

Checklist

• I have performed a self-review of my code.
• 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.

[codecov]

Codecov Report

❌ Patch coverage is 51.11111% with 22 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...xtext/checkpoint_conversion/utils/param_mapping.py	33.33%	4 Missing ⚠️
src/maxtext/multimodal/utils.py	0.00%	4 Missing ⚠️
src/maxtext/common/checkpointing.py	33.33%	1 Missing and 1 partial ⚠️
src/maxtext/input_pipeline/input_pipeline_utils.py	60.00%	1 Missing and 1 partial ⚠️
...xtext/kernels/attention/splash_attention_kernel.py	0.00%	2 Missing ⚠️
src/maxtext/layers/nnx_wrappers.py	0.00%	2 Missing ⚠️
src/maxtext/multimodal/processor.py	0.00%	2 Missing ⚠️
...t/integration/vllm/maxtext_vllm_adapter/adapter.py	0.00%	1 Missing ⚠️
src/maxtext/layers/quantizations.py	66.66%	1 Missing ⚠️
...ners/post_train/distillation/distillation_utils.py	50.00%	1 Missing ⚠️
... and 1 more

📢 Thoughts on this report? Let us know!