[deprecated] clean tokamax gmm and add test

[copybara-service]

clean tokamax gmm and add test

Description

1. Refactor manual_axis_type conditional specification.

• It was previously introduced by PR#3770 and PR#3869.
• It requires tokamax > 0.12.0 (unreleased yet) with commit
• It is used for gmm shard_map with check_vma=True. This is currently only activated for experimental run (deepseek_batchsplit and use_manual_quantization). Tests won't blocked by dependency.

2. Add unit test for use_tokamax_gmm=True for smoke train. Both bf16 and fp8.

• Smoke train instead of AOT as blocked by b/489205940
• Run in github CI/CD as tpu_only

3. Reset group offset to None: PR#4082 added group_offset to tokamax.ragged_dot and tokamax.ragged_dot_general. However, group_offset is not yet supported by tokamax:

• see tokamax code
• the added unit test meet error

Tests

unit test tokamax_gmm_test

Test on v5p-8:

python3 -m pytest tests/integration/tokamax_gmm_test.py

takes 51.50s, log

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

Files with missing lines	Patch %	Lines
src/maxtext/kernels/megablox/ops.py	0.00%	9 Missing ⚠️
src/maxtext/layers/quantizations.py	33.33%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!

[shuningjin]

Closing this. Created a new PR instead: #4152