Add seq parallelism for attention and MoE MLP

[suexu1025]

Description

1. Add seq_parallelism + exp_parallelism for attention + followed MLP module
   with sp+ep, moe customer 2k seq inference improved by 20%
2. Fix prefill_KV_cache sharding mismatch during seq_parallelism
3. decode improved by 10%
4. Enable inference auto layout in mistral model

FIXES: b/374773995

Tests

tested on v6e/v5p:
SEQ=2048

python MaxText/inference_microbenchmark.py MaxText/configs/inference.yml max_prefill_predict_length=$SEQ max_target_length=6144 model_name=mixtral-8x7b ici_fsdp_parallelism=1 ici_autoregressive_parallelism=1 ici_expert_parallelism=1 ici_context_parallelism=4 ici_tensor_parallelism=1 scan_layers=false per_device_batch_size=1 attention=dot_product megablox=False quantization=int8 checkpoint_is_quantized=True quantize_kvcache=True capacity_factor=1 tokenizer_path=assets/tokenizer.mistral-v3 compute_axis_order=0,2,1,3 ar_cache_axis_order=0,2,1,3 enable_jax_profiler=True inference_microbenchmark_prefill_lengths="$SEQ" base_output_directory=$OUT_DIR run_name=$RUN_NAME profiler=xplane model_call_mode=inference inference_microbenchmark_stages=prefill

Checklist

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

• 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.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[RissyRan]

Thanks Qinwen! Would be great if you could help run a few steps of training tests with profiles. Previously we met issues in optimization of inference, and some changes have degragation of training performance.

[suexu1025]

address comments