[Do Not Merge] Optimizations on Qwen3-Next GatedDeltaNet w/ Kernel & XProf Agent

[Rohan-Bierneni]

Description

Using the suggestions from kernel & xprof agent, try to improve the Gated Delta Net implementation in qwen3.py. We test our changes using the script added as part of this pr. The script tests the forward pass, backward pass, overall train step, and memory consumption between the baseline implementation of the GDN versus our optimized version in qwen3.py. This allowed us to test out changes iteratively and quickly.

To test the script, please use the command:

python3 /src/maxtext/scratch_code/benchmark_gdn_optimization.py

Note: run this script on a TPU/GPU vm since on CPU it will take a while.

So far, total improvements on the Gated Delta Rule using Q3-Next configs & 4k Seq len are:

• https://paste.googleplex.com/5438820566827008
• Forward pass speedup: 2.27x
• Train step speedup: 3.75x
• Memory reduction: 76.01%

If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456

Tests

Tested our changes using the benchmarking script and pr unit tests (train_compile test for qwen3 next)

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 85.52632% with 11 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/MaxText/layers/qwen3.py	84.93%	8 Missing and 3 partials ⚠️

📢 Thoughts on this report? Let us know!