Carry rollout MoE routing from the vLLM sampler to the maxtext trainer

[khatwanimohit]

Lets an RL run replay the MoE expert routing a rollout actually used instead of re-running the router during training. Both ends already existed -- vLLM and tpu-inference capture routing (enable_return_routed_experts -> CompletionOutput.routed_experts, [length, num_layers, top_k]) and MaxText consumes it (forced_routed_experts) -- so this is the middle.

Sampler side:

• VllmConfig.return_routed_experts sets vLLM's engine arg; VllmSampler.detokenize collects the arrays and SamplerOutput carries them.
• The in-process adapter populates SamplingResponse.routed_experts, and warns when a request asks for routing the engine was not built to capture, since capture is engine-level and the alternative is silently handing the trainer None -- indistinguishable from a dense model.

Pipeline side:

• routed_experts on TokenSegment, TrajectoryItem and RLTrainerPayload, plus UNSET_ROUTED_EXPERT = -1 for slots the trainer must not replay.
• The GRPO adapter pads short captures with the sentinel rather than expert 0.
• The padded assembler lays routing out over [prompt | completion] the same way the token ids are padded: prompts right-aligned keeping the tail (as _left_pad does), completions left-aligned. Getting this wrong detaches every replayed expert from its token -- silent, not a crash.

Demo side (math_gsm8k_dist), all opt-in:

• trainer node: --trainer_engine=maxtext swaps PeftTrainer for MaxText's engine wired for replay, importing MaxText lazily and skipping the tunix actor model, which only knows the demo's dense architectures. --maxtext_load_parameters_path points it at the checkpoint; without one it warns, since random trainer weights make replayed routing meaningless.
• rollout node: --return_routed_experts, and --load_rollout_weights to load the real checkpoint instead of vLLM's dummy weights. With both sides on the same weights a first step is meaningful with no weight transfer at all.
• launcher.sh and k8s_launcher.sh: ROUTER_REPLAY=1 sets the above; slices and meshes are now configurable rather than hardcoded.
• Dockerfile: opt-in INSTALL_MAXTEXT build arg, which also drops triton (GPU-only, and its native module aborts the process when MaxText imports on a TPU host).

Tests pin layout rather than mere presence, since the realistic failure is misalignment. The loss assertions are differential -- two different replays must disagree -- because a "loss is finite" check passes just as happily when the routing is dropped. Each was verified to fail against a mutation that disables what it covers.

Checklist

• I have added all the necessary unit tests for my change.
• I have verified that my change does not break existing code and all unit tests pass.
• I have added all appropriate doc-strings/documentation.
• My PR is based on the latest changes of the main branch (if unsure, rebase the code).
• I have signed the Contributor License Agreement.
• I have followed Contribution Guidelines.

[tianshub]

Thank you for the update, looks much cleaner.