Add Gemma 4 FLOPs & fix sliding window flops computations

[gagika]

Description

Adds TFLOPs calculations for the Gemma 4 architecture (including MoE) and fixes several inaccuracies in existing FLOPs math (sliding window overlap, vision encoder scaling, and shared KV projections).

• Gemma 4 & MoE: Added Gemma 4 support. Fixed fallback MoE calculations to correctly use moe_mlp_dim and generalized MoE layer detection (num_experts > 1).
• Sliding Window: Corrected local causal FLOPs to account for triangular overlap (max_target_length * window - 0.5 * window**2).
• Vision Encoders: Fixed backward pass scaling (x3) for Gemma 3 and Llama 4 when parameters are unfreezed.
• KV Projections: Factored in share_kv_projections for accurate QKV FLOPs.

Tests

Added maxtext_utils_flops_test.py to validate FLOPs calculations across 12 model architectures.

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.

[github-actions]

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

[github-actions]

🤖 I'm sorry @gagika, but I was unable to process your request. Please see the logs for more details.

[codecov]

Codecov Report

❌ Patch coverage is 95.74468% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/utils/maxtext_utils.py	95.74%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!

[github-actions]

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

[github-actions]

## 📋 Review Summary

This Pull Request introduces Gemma 4 FLOPs calculations and significantly improves the accuracy of existing FLOPs math, particularly for sliding window attention and mixed attention architectures. The addition of a comprehensive test suite covering multiple model families is a major highlight and ensures the reliability of these critical metrics.

🔍 General Feedback

• Great Test Coverage: The new maxtext_utils_flops_test.py is excellent. It uses a robust 6 * params * tokens verification strategy that provides high confidence in the computed TFLOPs across various architectures.
• Improved Accuracy: The fixes for sliding window area and vision encoder scaling (backward pass) are well-timed and correct.
• Inconsistency in Shared KV Projections: There is a potential logic error in how share_kv_projections is applied to mixed attention models in the main caller. One unit test specifically assumes local layers do not share KV projections even when the flag is True, but the code currently applies it to both.
• MoE Fallback Logic: The fallback for MoE layer detection is now more generalized, which is good, but might be too broad for future hybrid architectures.

[github-actions]

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

[github-actions]

## 📋 Review Summary

This pull request significantly improves the accuracy of FLOPs and MFU (Model Flops Utilization) calculations across multiple architectures, with a focus on Gemma 4 and corrected sliding window logic. The implementation is thorough, including a new comprehensive test suite that validates calculations for 12 different model configurations.

🔍 General Feedback

• Accuracy Improvements: The switch to a precise triangular overlap formula for sliding window attention and the inclusion of backward pass FLOPs for vision encoders are excellent updates that prevent MFU over-estimation.
• Architectural Coverage: The addition of Gemma 4 specific logic and the generalization of MoE layer detection make the utilities much more robust for future model support.
• Testing: The new maxtext_utils_flops_test.py is a great addition, providing clear manual-calculation-based verification for various architectures.
• Suggestions: I've provided a few suggestions to further generalize the MoE layer detection and ensure consistent dimension usage in MoE FFN calculations.

[gobbleturk]

I love the test! I am not sure how we have gotten this far without testing our tflops calculation...