clean tokamax gmm and add test

[shuningjin]

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
• FIXES: b/523289148

Tests

unit test tokamax_test

Test on v5p-8: log

 pytest -v -m integration_test tests/integration/tokamax_test.py

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

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

📢 Thoughts on this report? Let us know!