Add Tokamax ring MHA attention for TPU by huytransformer · Pull Request #4266 · AI-Hypercomputer/maxtext

huytransformer · 2026-06-25T05:29:48Z

Description

Adds TPU support for context_parallel_strategy=ring in Flash attention using Tokamax Splash.

The ring path keeps Q sequence-sharded, rotates K/V blocks across the context axis, and combines the partial attention results using online softmax. This PR uses the Tokamax Splash components and includes the MaxText integration, config validation, and sharding checks.

Sequence packing and CP load balancing are not enabled in this PR.

Tests

Ran:

python3 -m pytest tests/unit/tokamax_ring_attention_test.py
python3 -m pytest tests/unit/tokamax_splash_attention_mask_test.py
python3 -m pytest tests/unit/configs_value_test.py -k tpu_tokamax_ring
python3 -m pytest tests/unit/attention_test.py -k tpu_flash_attention_ring_context_parallel

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 · 2026-06-25T05:34:32Z

Codecov Report

❌ Patch coverage is 61.33614% with 735 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...okamax_splash_attention/splash_attention_kernel.py	67.34%	205 Missing and 82 partials ⚠️
.../tokamax_splash_attention/splash_attention_mask.py	40.54%	129 Missing and 3 partials ⚠️
...max_splash_attention/splash_attention_mask_info.py	60.49%	71 Missing and 25 partials ⚠️
...c/maxtext/kernels/tokamax_splash_attention/base.py	18.62%	83 Missing ⚠️
...axtext/kernels/attention/tokamax_ring_attention.py	57.97%	31 Missing and 27 partials ⚠️
.../tokamax_splash_attention/ring_attention_kernel.py	78.04%	22 Missing and 23 partials ⚠️
src/maxtext/layers/attention_op.py	65.21%	11 Missing and 13 partials ⚠️
...s/tokamax_splash_attention/ring_attention_utils.py	76.74%	6 Missing and 4 partials ⚠️

📢 Thoughts on this report? Let us know!

# Conflicts: # src/maxtext/layers/attention_op.py

NuojCheng · 2026-07-06T17:22:48Z

+  if keys["use_qk_clip"]:
+    raise ValueError("TPU Tokamax ring attention does not support QK-Clip statistics yet.")
+  if keys["enable_dropout"] and keys["dropout_rate"] > 0.0:
+    raise ValueError("TPU Tokamax ring attention does not support dropout yet.")


pyconfig_deprecated.py is no longer in use.

NuojCheng · 2026-07-06T17:41:25Z

+      if self.use_qk_clip:
+        raise ValueError("TPU Tokamax ring attention does not support QK-Clip statistics yet.")
+      if self.enable_dropout and self.dropout_rate > 0.0:
+        raise ValueError("TPU Tokamax ring attention does not support dropout yet.")


This is a great, exhaustive list of early stopping conditions (we need more of this!). As a quick refactoring idea: what do you think about mapping these out in a dictionary instead of using a long chain of if statements?

That's a neat idea.
I saw people mostly use if....raise individually in this file so I didn't bring in more abstractions. I will make changes though if you insist but it's a P2 for now : )

NuojCheng · 2026-07-06T17:44:55Z

+          indexer_mask=indexer_mask,
+          use_ragged_attention=use_ragged_attention,
+          record_max_logits=record_max_logits,
+      )


I would move the validation logic to initi of the module

NuojCheng · 2026-07-06T17:45:47Z

+          axis_names_kv=axis_names_kv,
+          dkv_dim_q=3,
+          dkv_dim_kv=3,
+      )


I would move validation logic to init instead of call

NuojCheng · 2026-07-06T17:54:13Z

+    mha_generic_flash_cp_output = jax.device_get(mha_generic_flash_cp_output)
+
+    self.assertTrue(
+        jax.numpy.allclose(mha_generic_output, mha_generic_flash_cp_output, rtol=1e-01, atol=1e-01, equal_nan=False),


are 1e-01 the best we can do? maybe try using dtype=float32 and see if we can reduce them

1e-1 was due to following other similar instances for comparing cross kernels. Some examples: attention_test.py line 674 (dot_product vs flash), 881 (dot_product vs flash_cp) , 1872 (MLA CP dot_product vs flash+cp).

Edited:
This is the best we can go afaik

# Forward rtol=1e-2, atol=1e-2 # Backward rtol=1e-2, atol=1e-7

NuojCheng · 2026-07-06T17:55:32Z

    hooks:
      - id: pyink
-        exclude: src/maxtext/input_pipeline/protos/
+        exclude: src/maxtext/input_pipeline/protos/|src/maxtext/kernels/tokamax_splash_attention/(base|ring_attention_kernel|splash_attention_kernel|splash_attention_mask|splash_attention_mask_info)\.py


not necessary, please do not change this file

NuojCheng · 2026-07-06T17:55:50Z

+  - "src/maxtext/kernels/tokamax_splash_attention/ring_attention_kernel.py"
+  - "src/maxtext/kernels/tokamax_splash_attention/splash_attention_kernel.py"
+  - "src/maxtext/kernels/tokamax_splash_attention/splash_attention_mask.py"
+  - "src/maxtext/kernels/tokamax_splash_attention/splash_attention_mask_info.py"


why making this change?

I just wanted codecov to ignore the bulk from tokamax but ya we don't have to touch this

NuojCheng · 2026-07-06T18:01:57Z

+from maxtext.kernels.attention import tokamax_ring_attention
+
+
+class TokamaxRingAttentionTest(absltest.TestCase):


I think many of these tests are ran on CPU instead of TPU. For shape checks we could use AOT instead of real runs. We could chat more about details.

for TPU previously we only had forward test and I added backward test.

NuojCheng · 2026-07-10T23:07:29Z

    )
-    if context_parallel_size > 1 and self.context_parallel_strategy.lower() == "ring":
-      if "gpu" not in self.hardware:
+    context_parallel_strategy = self.context_parallel_strategy.lower()


nit: alternatively we can create Enum for context parallel strategy, e.g. in https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/common/common_types.py, to avoid mis-spell.

that would touch several paths so perhaps a separate PR?

NuojCheng · 2026-07-10T23:09:29Z

+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tokamax ring attention helpers."""


please add references in tokamax

NuojCheng · 2026-07-10T23:17:14Z

+# limitations under the License.
+"""Tokamax SplashAttention runtime.
+
+These modules are adapted from OpenXLA Tokamax SplashAttention.


provide link

NuojCheng · 2026-07-10T23:27:56Z

+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================


Is this file necessary to clone? Looks like it is mask-unrelated? I am inclined to minimize cloning as much as possible.

I'll respond below

NuojCheng · 2026-07-10T23:28:30Z

+  return f"splash_{attention_type}_{phase}{segments}{residuals}"
+
+
+# Splash attention implementation


emmm empty?

not empty. I didn't change anything about this part from tokamax 0.0.12. It just means it owns the stuff below

NuojCheng · 2026-07-10T23:31:11Z

+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================


is this file necessary to clone?

I'll respond below

NuojCheng · 2026-07-10T23:31:49Z

+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================


is this file necessary to clone?

I'll respond below

NuojCheng

Overall LGTM. I am a little scared by the scale of cloned codes. I have two suggestions on editing the PR description:

Please specify reasons why cloning tokamax ring/splash and making local changes are necessary;
Please provide example maxtext command for running ring attention, with performance comparison, see example description of PR#2783, also add link to your doc.

…kamax-ring-tpu # Conflicts: # src/maxtext/configs/pyconfig_deprecated.py # src/maxtext/layers/attention_op.py # tests/unit/configs_value_test.py

huytransformer force-pushed the htn/tokamax-ring-tpu branch from 050abe5 to fb361d0 Compare June 25, 2026 07:09

Add Tokamax ring attention for TPU

9b610fb

huytransformer force-pushed the htn/tokamax-ring-tpu branch from fb361d0 to 9b610fb Compare June 25, 2026 16:17

huytransformer marked this pull request as ready for review June 26, 2026 04:24

huytransformer force-pushed the htn/tokamax-ring-tpu branch from ce68360 to f02b0bc Compare June 26, 2026 20:33

Merge branch 'main' into htn/tokamax-ring-tpu

98393b6

huytransformer force-pushed the htn/tokamax-ring-tpu branch from f02b0bc to 98393b6 Compare June 28, 2026 03:27

Merge remote-tracking branch 'upstream/main' into htn/tokamax-ring-tpu

a54462b

# Conflicts: # src/maxtext/layers/attention_op.py

huytransformer force-pushed the htn/tokamax-ring-tpu branch from 8219873 to a54462b Compare June 30, 2026 03:32

NuojCheng reviewed Jul 6, 2026

View reviewed changes

Comment thread src/maxtext/layers/attention_op.py

NuojCheng reviewed Jul 6, 2026

View reviewed changes

huytransformer requested a review from xibinliu as a code owner July 9, 2026 17:14

huytransformer force-pushed the htn/tokamax-ring-tpu branch 3 times, most recently from 45e5f89 to 37b508c Compare July 10, 2026 21:30

NuojCheng reviewed Jul 10, 2026

View reviewed changes

NuojCheng approved these changes Jul 10, 2026

View reviewed changes

NuojCheng changed the title ~~Add Tokamax ring attention for TPU~~ Add Tokamax ring MLA attention for TPU Jul 10, 2026

NuojCheng changed the title ~~Add Tokamax ring MLA attention for TPU~~ Add Tokamax ring MHA attention for TPU Jul 10, 2026

Merge remote-tracking branch 'refs/remotes/upstream/main' into htn/to…

0eccf59

…kamax-ring-tpu # Conflicts: # src/maxtext/configs/pyconfig_deprecated.py # src/maxtext/layers/attention_op.py # tests/unit/configs_value_test.py

huytransformer force-pushed the htn/tokamax-ring-tpu branch from 37b508c to 0eccf59 Compare July 11, 2026 07:52

		from maxtext.kernels.attention import tokamax_ring_attention


		class TokamaxRingAttentionTest(absltest.TestCase):

		return f"splash_{attention_type}_{phase}{segments}{residuals}"


		# Splash attention implementation

Uh oh!

Conversation

huytransformer commented Jun 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Tests

Checklist

Uh oh!

codecov Bot commented Jun 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

huytransformer Jul 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NuojCheng left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

huytransformer commented Jun 25, 2026 •

edited

Loading

codecov Bot commented Jun 25, 2026 •

edited

Loading

huytransformer Jul 9, 2026 •

edited

Loading