[hengtaoguo]

Description

Adds calculate_hy4_tflops_training_per_device so MFU reporting for hy4-preview follows the same per-model format as DeepSeek-V4/Gemma4.

• FFN: HY4 is added to the routed+shared MoE path, and get_dense_moe_layers returns 1 dense + 77 MoE layers.
• Gated MLA: Q/KV LoRA, the output projection, and the elementwise sigmoid gate_proj (D → H·D_v).
• DSA indexer: projections and scoring run only on the 21/78 "full" indexer layers (layer 0 plus the first layer of every 4-layer MoE cycle); shared layers reuse the previous mask. Attention is sparse (top-k mask ratio) when seq_len > indexer_topk and dense causal otherwise.
• Indexer counted forward-only (1×): its inputs are stop-gradient, the top-k mask is non-differentiable, and HY4 has no indexer loss, so it gets no gradient. This adds about 0.2% of total FLOPs at 4K and about 3.5% at 128K.
• iHC: pre/post mixing projection, stream contraction and post-expansion, twice per layer, plus the final HC head. Unlike mHC there is no k² residual-mixing term.

Tests

• New test_hy4_preview_flops checks the dense/MoE layer split and the weight/attention TFLOPs at seq 4096, and the indexer-bypass path at seq_len = indexer_topk.
• Cross-checked against the abstract model's parameter shapes: 769.9B total / 48.7B active params (HF model card: 770B / 49B). Matmul and indexer-projection FLOPs match 2·tokens·params exactly.
• flop_calculation_test.py: 53/53 pass. pyink clean.
• Training loop functional on a TPU v5p-8 for a mini model:

python -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml model_name=hy4-preview run_name=hy4_mini base_output_directory=/home/hengtaoguo_google_com/projects/runs override_model_config=true base_num_decoder_layers=3 first_num_dense_layers=1 dataset_type=synthetic per_device_batch_size=3 max_target_length=4096 steps=10 enable_checkpointing=false weight_dtype=bfloat16 mu_dtype=bfloat16 attention=dot_product

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. (Note: New PRs get skip-ci-tests so only CPU/lint/docs checks run during review; a bot removes it after 2 approvals on a 1-commit PR to run TPU/GPU suites, or remove it yourself / ask a maintainer on fork PRs to run them sooner.)
• 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.

[github-actions]

🤖 Hi @hengtaoguo, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[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 implements the TFLOPs calculation logic for the Hy4-preview model architecture within the MaxText framework. It ensures that MFU reporting correctly accounts for the specific computational overheads of Hy4, including its gated MLA projections, periodic DSA indexer, and identity hyper-connections. The changes provide a robust way to estimate training performance for this model type.

Highlights

• New TFLOPs Calculation: Added calculate_hy4_tflops_training_per_device to accurately compute training TFLOPs for the Hy4 architecture, accounting for Gated MLA, DSA indexer, and iHC components.
• Architecture Support: Updated get_dense_moe_layers and the main TFLOPs calculation logic to include support for the HY4 decoder block type.
• Validation: Introduced test_hy4_preview_flops in the unit test suite to verify layer counts and TFLOPs calculations against expected values for various sequence lengths.

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]

🤖 Hi @hengtaoguo, 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.

[github-actions]

📋 Review Summary

Verdict: Spec fully met; logic is sound and tests correctly validate both the dense and sparse flop formulas. No significant issues found.

Inferred spec & validation

• Sources: PR description, code before/after.
• Problem: hy4-preview lacks a dedicated per-model TFLOPs calculation format for accurate MFU reporting.

Spec item	Implementation	Validating test
R1 (stated): Calculate HY4 specific TFLOPs including MLA, DSA indexer and iHC.	✅	flop_calculation_test.py::test_hy4_preview_flops ✅
R2 (stated): Register HY4 in get_dense_moe_layers helper.	✅	flop_calculation_test.py::test_hy4_preview_flops ✅
C1 (stated): Indexer FLOPs count forward-only and are bypassed when seq_len <= indexer_topk.	✅ preserved	flop_calculation_test.py::test_hy4_preview_flops ✅

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!