[darisoy]

Summary

Updates the FSDP weight prefetching guide (docs/guides/optimization/fsdp_prefetch.md) with recent optimizations, visual forward and backward timelines, and TPU7x benchmark results:

• Visual forward and backward timelines: Adds annotated SVG timeline diagrams (docs/_static/fsdp_*_overview.svg) comparing prefetch_fsdp_weights=False vs. prefetch_fsdp_weights=True (with prefetch_fsdp_defer_grad_reduce_scatter=True in the backward pass) across TensorCore, Dedicated FSDP SparseCore 0, and SparseCore 1.
• Deferred backward reduce-scatter (prefetch_fsdp_defer_grad_reduce_scatter=True): Explains why gathering weights with to="reduced" (returning unreduced per-device weight gradients from jax.vjp) and pinning scatter_dim as major-most (_pin_scatter_dim_major) allow a single SparseCore reduce-scatter, and documents how carrying step k's unreduced weight gradients in the backward scan carry and issuing their xla_explicit_fsdp-tagged psum_scatter at the start of step k−1 hides backward FSDP reduce-scatters alongside the prefetch all-gathers on the dedicated FSDP SparseCore.
• Inhomogeneous nested-scan boundary sub-layer peeling: Documents how boundary sub-layers whose input or output carry holds weights or deferred gradients from an adjacent group are peeled into single-trip inner while loops so interior sub-layers share an unroll=2 double-buffered inner scan ([GDN 0, GDN 1] → [GDN 2] → [attention] in forward, and [attention] → [GDN 2] → [GDN 1] → [GDN 0] with deferred backward reduce-scatter on Qwen3.5).
• MoE Tokamax gmm_v2 flags (moe_gmm_v2_dlhs_transpose_rhs=True and moe_use_direct_token_gather=True): Documents eliminating explicit backward dlhs RHS weight-transpose copies and token-replication data formatting once FSDP collectives are hidden.
• Updated TPU7x benchmark tables: Adds a top-level summary table across all three configurations (DeepSeek-V3 FSDP-only at −34.2%, DeepSeek-V3 with EP=4 at −38.1%, and Qwen3.5-35B-A3B inhomogeneous nested scan at −25.1%) alongside detailed ablation tables.

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 and provided workload links above if applicable.
• I have added or updated tests accordingly.
• I have updated documentation accordingly.
• I have included unit tests for new features or bug fixes.
• I have added unit test and integration test coverage for all modified code paths.

[github-actions]

🤖 Hi @darisoy, I'm running an automatic MaxText review of this pull request. You can track my progress in the logs. To request another review later, add the gemini-review label or comment @gemini-cli /review.

[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 updates the FSDP weight prefetching documentation to reflect recent architectural improvements. It provides detailed guidance on handling complex, inhomogeneous model blocks, configuring expert and context parallelism, and applying optimal XLA flags to maximize performance on TPU hardware. The changes include updated usage examples and performance tables to help users achieve better throughput and HBM efficiency.

Highlights

• Inhomogeneous Block Prefetching: Added support for prefetching inhomogeneous blocks (e.g., Qwen3.5) by processing sub-layers individually while maintaining nested scan structures.
• Expert and Context Parallelism: Documented co-sharding support for expert and context parallelism, ensuring efficient FSDP gathering across compound axes.
• Performance Optimization Flags: Introduced recommended configurations including shard_mode=explicit, moe_flat_fsdp_weights=True, and updated LIBTPU_INIT_ARGS for improved latency hiding.
• Benchmark Updates: Updated documentation with comprehensive performance benchmarks for DeepSeek-V3 and Qwen3.5, highlighting significant step-time reductions.

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

[github-actions]

📋 Review Summary

Verdict: No significant issues found. The documentation update is highly precise, accurate, and comprehensively details the recent FSDP prefetching improvements and best practices.

Inferred spec & validation

• Sources: PR description, code before/after.
• Problem: FSDP weight prefetching documentation needs updating to reflect new features (symmetric scan, inhomogeneous block prefetching, co-sharded axes, and moe_flat_fsdp_weights).

Spec item	Implementation	Validating test
Document symmetric unroll=2 scan	✅	docs/guides/optimization/fsdp_prefetch.md ✅
Document inhomogeneous sub-layer scans	✅	docs/guides/optimization/fsdp_prefetch.md ✅
Document expert & context parallelism sharding	✅	docs/guides/optimization/fsdp_prefetch.md ✅
Provide recommended MaxText and XLA flags	✅	docs/guides/optimization/fsdp_prefetch.md ✅

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!