Add support for Tokamax GMM v2 in MaxText.

[CaptainO5]

Description

This change temporarily forks and integrates a new Pallas Mosaic TPU v2 ragged dot implementation from Tokamax. A new configuration option use_gmm_v2 is added to enable this feature. GMM v2 is currently not supported when used with quantization. The forward and backward passes in the grouped matrix multiplication (GMM) operation are updated to use the new v2 implementation when enabled.

BUGS: b/505011787

Tests

I added a unit test in integration/tokamax_test and also verified that the e2e test ran successfully.

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

Files with missing lines	Patch %	Lines
...ernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py	60.96%	178 Missing and 32 partials ⚠️
...rnels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py	72.51%	63 Missing and 9 partials ⚠️
src/maxtext/kernels/megablox/ops.py	81.25%	1 Missing and 2 partials ⚠️

📢 Thoughts on this report? Let us know!

[shuningjin]

Thank you for integrating this performant new kernel! I see we are able to use full v2 with bf16 now, which is great.

• Had a question on transpose_rhs
• As mentioned by others, let's move the check to types.py (for use_gmm_v2, use_tokamax_gmm, quantization).

[gobbleturk]

Awesome, thank you!

[shuningjin]

LGTM. Thank you!