[Recipes][MoE] Plumb MoE optimization switches across launcher and recipes

[wenxindongwork]

Summary

Plumb switches for evaluating MoE routing and collective optimizations in end-to-end runs:

• TPU_SMALL_ALL_REDUCE: Pallas low-latency ICI all-reduce for small token batches
• TPU_MOE_SINGLE_SORT: Packed single-sort permutation
• TPU_MOE_ROUTER_TOPK: sort vs rowmax top-k selection
• TPU_MOE_LOWTOK: Streaming few-token EP MoE kernel

All flags default to disabled (0 / sort) to preserve bit-exact baseline behavior.
Variables are propagated to VLLM_RAY_EXTRA_ENV_VARS_TO_COPY, forwarded to the rollout container, and included in the JAX cache hash payload.

Pin

• tpu-inference: bumped to 0524fdf70097f9973a9400115b925736c2f5450e (HEAD of atwigg/mlperf, which merges #3794 and #3793).

Note: Marked as Draft; untested on TPU hardware.

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request enables granular control over MoE (Mixture of Experts) and collective communication optimizations for TPU inference workloads. By plumbing these configuration switches through the infrastructure, users can now experiment with performance-tuning flags while maintaining bit-exact baseline behavior by default. The changes also ensure these settings are correctly synchronized across distributed Ray environments and cache-aware systems.

Highlights

• MoE Optimization Switches: Introduced new environment variables (TPU_SMALL_ALL_REDUCE, TPU_MOE_SINGLE_SORT, TPU_MOE_ROUTER_TOPK, TPU_MOE_LOWTOK) to control MoE routing and collective optimizations.
• Environment Propagation: Updated launcher and recipe scripts to ensure these optimization flags are propagated to Ray workers and included in JAX cache hashing.
• Dependency Update: Bumped tpu-inference dependency to include recent upstream changes from the mlperf branch.

New Features

🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request updates the tpu-inference dependency and introduces several environment variables (TPU_SMALL_ALL_REDUCE, TPU_MOE_SINGLE_SORT, TPU_MOE_ROUTER_TOPK, and TPU_MOE_LOWTOK) to configure MoE and collective optimizations across launcher and recipe scripts. The review feedback suggests double-quoting these variables in k8s_launcher.sh for consistency and robustness, propagating them to start_eval, and adding a deduplication check when appending them to VLLM_RAY_EXTRA_ENV_VARS_TO_COPY in mlperf_base.sh.

[gemini-code-assist]

For consistency and robustness, these environment variables should be double-quoted (with escaped quotes \"...\") just like the other variables above them (e.g., HF_TOKEN). This prevents potential issues if any of these variables contain spaces or special characters.

Additionally, please ensure these same MoE optimization switches are propagated to the worker startup command in start_eval (around line 1192) to maintain consistency between rollout and evaluation runs.

Suggested change

	${TPU_SMALL_ALL_REDUCE:+TPU_SMALL_ALL_REDUCE=${TPU_SMALL_ALL_REDUCE}} \
	${TPU_MOE_SINGLE_SORT:+TPU_MOE_SINGLE_SORT=${TPU_MOE_SINGLE_SORT}} \
	${TPU_MOE_ROUTER_TOPK:+TPU_MOE_ROUTER_TOPK=${TPU_MOE_ROUTER_TOPK}} \
	${TPU_MOE_LOWTOK:+TPU_MOE_LOWTOK=${TPU_MOE_LOWTOK}} \
	${TPU_SMALL_ALL_REDUCE:+TPU_SMALL_ALL_REDUCE=\"${TPU_SMALL_ALL_REDUCE}\"} \
	${TPU_MOE_SINGLE_SORT:+TPU_MOE_SINGLE_SORT=\"${TPU_MOE_SINGLE_SORT}\"} \
	${TPU_MOE_ROUTER_TOPK:+TPU_MOE_ROUTER_TOPK=\"${TPU_MOE_ROUTER_TOPK}\"} \
	${TPU_MOE_LOWTOK:+TPU_MOE_LOWTOK=\"${TPU_MOE_LOWTOK}\"} \

References

1. Adherence to a consistent style across all projects is crucial for predictability and maintainability. ^(link)

[gemini-code-assist]

To prevent duplicate entries in VLLM_RAY_EXTRA_ENV_VARS_TO_COPY if any of these variables are already present, we can use a deduplication check similar to the one used for RAGGED_GATHER_REDUCE_VERSION above.

Suggested change

	for _opt_var in TPU_SMALL_ALL_REDUCE TPU_MOE_SINGLE_SORT TPU_MOE_ROUTER_TOPK TPU_MOE_LOWTOK; do
	if [[ -n "${VLLM_RAY_EXTRA_ENV_VARS_TO_COPY:-}" ]]; then
	export VLLM_RAY_EXTRA_ENV_VARS_TO_COPY="${VLLM_RAY_EXTRA_ENV_VARS_TO_COPY},${_opt_var}"
	else
	export VLLM_RAY_EXTRA_ENV_VARS_TO_COPY="${_opt_var}"
	fi
	done
	for _opt_var in TPU_SMALL_ALL_REDUCE TPU_MOE_SINGLE_SORT TPU_MOE_ROUTER_TOPK TPU_MOE_LOWTOK; do
	if [[ ",${VLLM_RAY_EXTRA_ENV_VARS_TO_COPY:-}," != *",${_opt_var},"* ]]; then
	export VLLM_RAY_EXTRA_ENV_VARS_TO_COPY="${VLLM_RAY_EXTRA_ENV_VARS_TO_COPY:+${VLLM_RAY_EXTRA_ENV_VARS_TO_COPY},}${_opt_var}"
	fi
	done