Fix: Gemma 3 & 4 Base Model and Flax Linen Decoding

[RexBearIU]

Description

This PR introduces architectural configurations for Gemma3 and Gemma4, and implements critical decoding layer KV-cache propagation fixes for JAX/NNX
(nnx_decoders.py) and Flax Linen (decoders.py) pipelines.

Before this change, Flax Linen scanned blocks (_apply_gemma3_scanned_blocks and _apply_gemma4_scanned_blocks) lacked mechanisms to propagate
intermediate KV-caches across scanned layers under nn.scan, causing end-to-end decodes to fail.

Key Changes:
- Flax Linen Decoder (src/maxtext/layers/decoders.py): Added kv_caches and attention_metadata propagation across Gemma3/4 scanned blocks
via a stack/unstack mapping over nn.scan.
- JAX/NNX Decoder (src/maxtext/layers/nnx_decoders.py): Unified dynamic KV-cache carry updates within scanned block execution.
- Model Scaffolding (gemma3.py, gemma4.py): Added kv_cache routing through scannable blocks and layer inputs.

Tests

• Validated via end-to-end decoding with Gemma3-4B on an 8-device TPU mesh.
• Run unit tests:

python3 -m pytest tests/post_training/unit/lora_utils_test.py

• Decode test:
• command and output

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.

[codecov]

Codecov Report

❌ Patch coverage is 34.81481% with 88 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/decoders.py	0.00%	36 Missing ⚠️
src/maxtext/layers/nnx_decoders.py	41.07%	22 Missing and 11 partials ⚠️
src/maxtext/utils/maxtext_utils.py	66.66%	6 Missing and 2 partials ⚠️
src/maxtext/models/gemma4.py	45.45%	3 Missing and 3 partials ⚠️
src/maxtext/models/gemma3.py	37.50%	2 Missing and 3 partials ⚠️

📢 Thoughts on this report? Let us know!

[github-actions]

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

[NuojCheng]

is kv_caches a python list, not jax array? jax array is not mutable. If so, use if isinstance(kv_caches) List to make sure kv_caches is a list

[RexBearIU]

Yes, good catch. kv_caches is indeed a Python list containing JAX PyTrees. Since JAX arrays/PyTrees are immutable but we mutate the outer list structure in update_kv_caches_after_scan, I've added explicit isinstance(kv_caches, list) checks and validation tests to ensure it behaves correctly. I just pushed the fixes.