Add tokamax megablox kernel

[khatwanimohit]

Description

• adds tokamax as a deps in requirements*.txt files.
• use ragged_dot api from tokamax hidding behind a flag use_tokamax
• this PR does a conditional import for tokamax because maxtext's seed-env is pinned at jax==0.7.0 and tokamax has jax>=0.7.2 in its deps. Added a TODO to remove the condition once maxtext upgrades to jax==0.8.0

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

Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.

Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.

Tests

Perf and convergence runs done by @suexu1025 in b/436335556#comment11

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.

[suexu1025]

do we plan to include latest version in the nightly build?

[khatwanimohit]

our nightly docker image build only supports using jax,jaxlib,libtpu nightly versions. If we want to have a nightly run then we can prepend pip install git+https://github.com/openxla/tokamax.git before the train command which should work

[rdyro]

It might be good to distinguish specific tokamax kernels, use_tokamax_gmm use_tokamax_splash etc