[Stacked PR 5/5] Add granular GDN remat policy and 8k/64k latency/memory benchmark suite

[Rohan-Bierneni]

Stacked PR Chain

Stack	Status	Branch	Base	PR
PR 1	🔗 Stack PR 1	rbierneni-gdn-1-ci-hygiene	main	#5151 - [Stacked PR 1/5] Vendor pristine Tokamax GDN forward kernel
PR 2	🔗 Stack PR 2	rbierneni-gdn-2-fwd-kernel	rbierneni-gdn-1-ci-hygiene	#5152 - [Stacked PR 2/5] Import local Tokamax GDN forward kernel with custom remat for backward pass support
PR 3	🔗 Stack PR 3	rbierneni-gdn-3-bwd-kernel	rbierneni-gdn-2-fwd-kernel	#5153 - [Stacked PR 3/5] Add decoupled GDN Pallas backward pass kernel and parity tests
PR 4	🔗 Base PR	rbierneni-gdn-4-model-integration	rbierneni-gdn-3-bwd-kernel	#5154 - [Stacked PR 4/5] Enable GDN backward pass kernel and hybrid precision in Qwen3 model
PR 5	🚀 This PR (#5098)	rbierneni-gdnv3-bwd	rbierneni-gdn-4-model-integration	#5098 - [Stacked PR 5/5] Add granular GDN remat policy and 8k/64k latency/memory benchmark suite

---

Description

This is PR 5 of 5 (the capstone PR #5098) in the stacked series establishing the canonical Pallas Gated Delta Net (GDN) implementation in MaxText.

This PR introduces:

1. Granular Native GDN Rematerialization Policy:

   • Adds native remat toggles gdn: RematLocation and gdn_conv: RematLocation ('remat', 'device', 'offload') to base.yml, types.py, and maxtext_utils.py.
   • Adds cross-field configuration validation (validate_gdn_remat_requires_kernel) enforcing that setting gdn or gdn_conv to 'device' or 'offload' requires use_gdn_kernel=True. This prevents silent misconfiguration where Pure JAX would otherwise ignore the setting and redundantly recompute activations.
   • Adds _expand_gdn_remat_names in maxtext_utils.py expanding "gdn" to individual named residuals: ["gdn_core_attn_out", "gdn_chunk_states", "gdn_t_inv", "gdn_qkv", "gdn_b", "gdn_a"] and "gdn_conv" to ["gdn_conv_out", "gdn_fwd_conv"].
   • Tags core GDN outputs and recurrence residuals with jax.ad_checkpoint.checkpoint_name in gdn_bwd_pallas.py and qwen3.py.
   • Implements _get_gdn_aware_remat_policy in qwen3.py ensuring that when custom remat is enabled (tensors_on_device: ["decoder_layer_input", "gdn"]), forward recurrence residuals are preserved in device memory, completely eliminating redundant forward kernel recomputation (~17 ms) during backward autodiff while keeping all Feed-Forward/MoE layers 100% rematerialized.
   • Adds unit tests in gdn_bwd_pallas_test.py and pyconfig_test.py covering gradient invariance, residual preservation, and configuration validation.

2. Cloud TPU Latency & Memory Benchmark Suite (tests/unit/gdn_benchmark_test.py):

   • Provides standalone 8k and 64k Cloud TPU profiling scripts comparing Pure JAX against the Decoupled GDN Pallas kernel.
   • Measures isolated Pallas hardware calls (forward and backward) as well as the full monolithic decoder layer.

64k Ghostfish Cloud TPU Hardware Performance (1,024 Chunks, S=65,536, B=1, H_k=16, H_v=64, D=128):

• Isolated Pallas Hardware Kernel Calls:
  • Core Pallas Forward Kernel: 37.55 ms
  • Core Pallas Backward Kernel: 93.59 ms
• Full Monolithic 1-Chip Decoder Layer (Projections + Gating + Conv1D + Kernels):
  • Forward Pass: 68.18 ms
  • Backward Pass: 185.99 ms (includes ~150 ms linear GEMM backprop on a single chip)
  • Total Step Latency: 254.17 ms
  • Peak HBM Memory: 32.83 GB (vs Pure JAX OOM)

Files Changed

• src/maxtext/configs/base.yml: Added gdn: 'remat' and gdn_conv: 'remat' custom remat options.
• src/maxtext/configs/types.py: Added gdn and gdn_conv fields in RematAndOffload, included in custom remat lists, and added cross-field validator requiring use_gdn_kernel=True.
• src/maxtext/utils/maxtext_utils.py: Added _expand_gdn_remat_names and integrated with get_save_and_offload_names.
• src/maxtext/models/kernels/gdn/gdn_bwd_pallas.py: Tagged output and residuals with checkpoint_name.
• src/maxtext/models/qwen3.py: Added _GDN_SAVED_NAMES, _get_gdn_aware_remat_policy, and tagged gdn_core_attn_out and out_proj.
• tests/unit/gdn_bwd_pallas_test.py: Added test_gdn_custom_remat_policy_preserves_residuals and test_gdn_granular_remat_requires_gdn_kernel.
• tests/unit/pyconfig_test.py: Added test_gdn_granular_remat_requires_gdn_kernel and test_gdn_granular_remat_accepts_gdn_kernel.
• tests/unit/gdn_benchmark_test.py: Standalone 8k and 64k Cloud TPU latency and memory benchmark suite with isolated kernel timings.

Tests

• Verified with pre-commit run --files ... (all hooks: pylint, pyink, codespell, yamllint passed cleanly).
• Google3 unit tests passed with -c opt: //third_party/py/maxtext/tests/unit:gdn_bwd_pallas_test, //third_party/py/maxtext/tests/unit:pyconfig_test, and //third_party/py/maxtext/tests/unit:gdn_benchmark_ghostfish_test.

Checklist

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests and verified pre-commit linters pass.

[gemini-code-assist]

Code Review

This pull request introduces a fused analytical Pallas Gated Delta Net (GDN) kernel and backward pipeline, integrating it into the Qwen3 model. Key changes include the addition of the hybrid_bwd_analytical_pipeline and a suite of supporting local GDN kernel files to handle causal Conv1D and Gated Delta Rule operations with cached triangular inverse matrices. Additionally, sublane tiling logic is adjusted in the GMM and TGMM kernels. Feedback on the changes highlights a numerical inconsistency where an epsilon of 1e-12 is used for L2 normalization in the backward pass instead of the 1e-6 used in the forward pass, as well as an opportunity to simplify the sublane alignment assignment in the TGMM kernel.

[gemini-code-assist]

The epsilon value 1e-12 used here for calculating the norm of q is inconsistent with the value 1e-6 used in the forward pass (l2norm at line 521). For numerical stability and correctness of the gradient, the same epsilon value should be used in both forward and backward computations.

Suggested change

	r_q = jnp.sqrt(jnp.sum(q_orig**2, axis=-1, keepdims=True) + 1e-12)
	r_q = jnp.sqrt(jnp.sum(q_orig**2, axis=-1, keepdims=True) + 1e-6)

[gemini-code-assist]

Similar to the calculation for r_q, the epsilon value 1e-12 used here for r_k is inconsistent with the 1e-6 used in the forward pass. This should be corrected to 1e-6 to ensure numerical consistency.

Suggested change

	r_k = jnp.sqrt(jnp.sum(k_orig**2, axis=-1, keepdims=True) + 1e-12)
	r_k = jnp.sqrt(jnp.sum(k_orig**2, axis=-1, keepdims=True) + 1e-6)

[gemini-code-assist]

The logic to ensure size_lhs_sublane and size_rhs_sublane are equal can be simplified into a single line, improving readability.

Suggested change

	common_sublane = min(size_lhs_sublane, size_rhs_sublane)
	size_lhs_sublane = common_sublane
	size_rhs_sublane = common_sublane
	size_lhs_sublane = size_rhs_sublane = min(size_lhs_sublane, size_rhs_sublane)

[codecov]

Codecov Report

❌ Patch coverage is 21.62162% with 29 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/models/qwen3.py	22.72%	17 Missing ⚠️
src/maxtext/utils/maxtext_utils.py	10.00%	9 Missing ⚠️
src/maxtext/models/kernels/gdn/gdn_bwd_pallas.py	40.00%	3 Missing ⚠️

📢 Thoughts on this report? Let us know!