Fix Qwen 3.5 35B RL gibberish output issue and show end-to-end RL convergence

[AntonyMei]

Description

This PR fixes the Qwen 3.5 RL gibberish output issue and shows end-to-end RL training convergence with Qwen 3.5 35B + v5p-64 on GSM8K. This is a clean up of the fixes in yixuanm-dev-35b-new branch.

FIXES: b/521604343
FIXES: b/542769108

Tests

On yixuanm-dev-35b-new, Qwen 3.5 35B + v5p-64 on GSM8K shows end-to-end convergence:
gs://yixuanm-maxtext-logs/ConvergenceTesting/ym-qw35-35/tensorboard

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.

[gemini-code-assist]

Code Review

This pull request adds support for Qwen3.5 post-training reinforcement learning (RL) using GRPO on TPU hardware. Key enhancements include a hybrid cache layout helper for Qwen's Gated DeltaNet (GDN) to prevent silent quantization, scanned-weight unrolling for direct MaxText-to-MaxText weight synchronization, support for partial multi-dimensional rotary position embeddings (MRoPE), and dynamic replication of batch dimensions in the Mixture of Experts (MoE) layer to avoid sharding mismatches. Additionally, answer extraction utilities have been updated to handle prefilled reasoning start tokens. The code review feedback suggests several robustness improvements, including converting all path segments to strings during weight unrolling, adding defensive checks for null values in the MoE layer to prevent AttributeError crashes, and making the answer extraction regex robust to empty reasoning blocks.

[NuojCheng]

Sharding and MoE components LGTM if my comments are resolved. Thank you for the code cleanup Yixuan!

[A9isha]

Thanks a lot Yixuan!