Support packed TPU Tokamax ring attention

src/maxtext/configs/types.py

@@ -3450,8 +3450,6 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
         raise ValueError("TPU ring context parallelism requires use_jax_splash=False.")
       if self.attention_type != "global":
         raise ValueError("TPU Tokamax ring attention is initially supported only for global causal attention.")
-      if self.packing:
-        raise ValueError("TPU Tokamax ring attention does not support packing yet.")
       if self.context_parallel_load_balance:
         if context_parallel_size % 2 != 0:
           raise ValueError("TPU Tokamax ring load balancing requires an even context_parallel_size.")
@@ -3474,9 +3472,8 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
       if self.enable_dropout and self.dropout_rate > 0.0:
         raise ValueError("TPU Tokamax ring attention does not support dropout yet.")
     # STRIPED reorder strategy is a Transformer Engine feature and is GPU-only.
-    # The AUTO + packing case, which training resolves to STRIPED, is not
-    # validated here because test code paths may load the same config but use a
-    # different reorder path. Training's runtime path enforces this.
+    # AUTO is resolved in training because test code paths may load the same
+    # config but use a different reorder path.
     if (
         context_parallel_size > 1
         and "gpu" not in self.hardware

src/maxtext/layers/attention_op.py

@@ -534,8 +534,6 @@ def __init__(
           raise ValueError("TPU Tokamax ring attention requires use_tokamax_splash=True.")
         if self.config.use_jax_splash:
           raise ValueError("TPU Tokamax ring attention requires use_jax_splash=False.")
-        if self.config.packing:
-          raise ValueError("TPU Tokamax ring attention does not support packing yet.")
         if self.attention_type != AttentionType.GLOBAL:
           raise ValueError("TPU Tokamax ring attention is initially supported only for global causal attention.")
 

src/maxtext/utils/train_utils.py

@@ -281,11 +281,15 @@ def create_train_state_fn():
     with jax.set_mesh(mesh):
       if context_parallel_size > 1 and config.context_parallel_load_balance:
 
-        # Determine load balancing reorder strategy based on whether packing is enabled
+        # Determine load balancing reorder strategy.
         if config.context_parallel_reorder_strategy == ReorderStrategy.AUTO:
           reorder_strategy = (
               ReorderStrategy.STRIPED
-              if config.packing and context_parallel_strategy == "ring"
+              if (
+                  config.packing
+                  and context_parallel_strategy == "ring"
+                  and config.hardware in ("gpu", "gpu_multiprocess")
+              )
               else ReorderStrategy.DUAL_CHUNK_SWAP
           )
         else:

tests/unit/attention_test.py

@@ -656,6 +656,20 @@ def get_data(self, dtype):
 
     return lnx, decoder_segment_ids, decoder_positions
 
+  def get_packed_data(self, dtype):
+    """get packed data"""
+    lnx, _, _ = self.get_data(dtype)
+    # Uneven segment lengths so boundaries don't line up with splash blocks, reorder chunks, or the CP shard split.
+    segment_lengths = (80, 240, 112, 80)
+    segment_ids = jnp.concatenate(
+        [jnp.full((length,), segment, dtype=jnp.int32) for segment, length in enumerate(segment_lengths, start=1)]
+    )
+    positions = jnp.concatenate([jnp.arange(length, dtype=jnp.int32) for length in segment_lengths])
+    decoder_segment_ids = jnp.broadcast_to(segment_ids, (self.global_batch_size, self.max_target_length))
+    decoder_positions = jnp.broadcast_to(positions, (self.global_batch_size, self.max_target_length))
+
+    return lnx, decoder_segment_ids, decoder_positions
+
   def get_structured_data(self, dtype):
     """get structured data"""
     lnx = jax.random.normal(
@@ -1143,11 +1157,13 @@ def test_tpu_flash_attention_packed_all_gather_context_parallel(self, context_pa
     )
 
   @parameterized.named_parameters(
-      {"testcase_name": "no_load_balance", "context_parallel_load_balance": False},
-      {"testcase_name": "load_balance", "context_parallel_load_balance": True},
+      {"testcase_name": "no_load_balance", "context_parallel_load_balance": False, "packing": False},
+      {"testcase_name": "load_balance", "context_parallel_load_balance": True, "packing": False},
+      {"testcase_name": "packed", "context_parallel_load_balance": False, "packing": True},
+      {"testcase_name": "packed_load_balance", "context_parallel_load_balance": True, "packing": True},
   )
   @pytest.mark.tpu_only
-  def test_tpu_flash_attention_ring_context_parallel(self, context_parallel_load_balance):
+  def test_tpu_flash_attention_ring_context_parallel(self, context_parallel_load_balance, packing):
     """Test equivalence between dot_product and flash attention + ring context parallelism"""
 
     cfg_cp = pyconfig.initialize(
@@ -1159,12 +1175,15 @@ def test_tpu_flash_attention_ring_context_parallel(self, context_parallel_load_b
         ici_context_parallelism=2,
         use_tokamax_splash=True,
         use_jax_splash=False,
-        packing=False,
+        packing=packing,
         dtype="float32",
     )
     devices_array_cp = maxtext_utils.create_device_mesh(cfg_cp)
     mesh_cp = Mesh(devices_array_cp, cfg_cp.mesh_axes)
-    lnx, decoder_segment_ids, decoder_positions = self.get_data(cfg_cp.dtype)
+    if packing:
+      lnx, decoder_segment_ids, decoder_positions = self.get_packed_data(cfg_cp.dtype)
+    else:
+      lnx, decoder_segment_ids, decoder_positions = self.get_data(cfg_cp.dtype)
     attention_as_mha_generic = Attention(
         config=self.cfg,
         num_query_heads=cfg_cp.num_query_heads,
@@ -1224,15 +1243,17 @@ def test_tpu_flash_attention_ring_context_parallel(self, context_parallel_load_b
     self.assertTrue(
         jax.numpy.allclose(mha_generic_output, mha_generic_flash_cp_output, rtol=1e-02, atol=1e-02, equal_nan=False),
         msg="Logits from generic dot product and flash attention + ring context parallelism are not close. "
-        f"context_parallel_load_balance={context_parallel_load_balance}.",
+        f"context_parallel_load_balance={context_parallel_load_balance}, packing={packing}.",
     )
 
   @parameterized.named_parameters(
-      {"testcase_name": "no_load_balance", "context_parallel_load_balance": False},
-      {"testcase_name": "load_balance", "context_parallel_load_balance": True},
+      {"testcase_name": "no_load_balance", "context_parallel_load_balance": False, "packing": False},
+      {"testcase_name": "load_balance", "context_parallel_load_balance": True, "packing": False},
+      {"testcase_name": "packed", "context_parallel_load_balance": False, "packing": True},
+      {"testcase_name": "packed_load_balance", "context_parallel_load_balance": True, "packing": True},
   )
   @pytest.mark.tpu_only
-  def test_tpu_flash_attention_ring_context_parallel_grad(self, context_parallel_load_balance):
+  def test_tpu_flash_attention_ring_context_parallel_grad(self, context_parallel_load_balance, packing):
     """Test gradient equivalence between dot_product and flash attention + ring context parallelism"""
 
     cfg_cp = pyconfig.initialize(
@@ -1244,12 +1265,15 @@ def test_tpu_flash_attention_ring_context_parallel_grad(self, context_parallel_l
         ici_context_parallelism=2,
         use_tokamax_splash=True,
         use_jax_splash=False,
-        packing=False,
+        packing=packing,
         dtype="float32",
     )
     devices_array_cp = maxtext_utils.create_device_mesh(cfg_cp)
     mesh_cp = Mesh(devices_array_cp, cfg_cp.mesh_axes)
-    lnx, decoder_segment_ids, decoder_positions = self.get_data(cfg_cp.dtype)
+    if packing:
+      lnx, decoder_segment_ids, decoder_positions = self.get_packed_data(cfg_cp.dtype)
+    else:
+      lnx, decoder_segment_ids, decoder_positions = self.get_data(cfg_cp.dtype)
     attention_as_mha_generic = Attention(
         config=self.cfg,
         num_query_heads=cfg_cp.num_query_heads,
@@ -1325,7 +1349,7 @@ def ring_loss(lnx):
     self.assertTrue(
         jax.numpy.allclose(generic_grad, ring_grad, rtol=1e-02, atol=1e-07, equal_nan=False),
         msg="Input gradients from generic dot product and flash attention + ring context parallelism are not close. "
-        f"context_parallel_load_balance={context_parallel_load_balance}.",
+        f"context_parallel_load_balance={context_parallel_load_balance}, packing={packing}.",
     )
 
   @pytest.mark.tpu_only

tests/unit/configs_value_test.py

@@ -143,6 +143,49 @@ def test_tpu_tokamax_ring_config_validation_accepts_load_balance(self):
 
     self.assertTrue(config.context_parallel_load_balance)
 
+  def test_tpu_tokamax_ring_config_validation_accepts_packing(self):
+    argv = [
+        "",
+        _BASE_CONFIG_PATH,
+        "run_name=test",
+        "attention=flash",
+        "use_tokamax_splash=True",
+        "use_jax_splash=False",
+        "context_parallel_strategy=ring",
+        "context_parallel_load_balance=False",
+        "ici_context_parallelism=2",
+        "hardware=tpu",
+        "packing=True",
+        "skip_jax_distributed_system=True",
+    ]
+    mock_devices = [unittest.mock.MagicMock(slice_index=0) for _ in range(8)]
+    with unittest.mock.patch("jax.devices", return_value=mock_devices):
+      config = pyconfig.initialize(argv)
+
+    self.assertTrue(config.packing)
+
+  def test_tpu_tokamax_ring_config_validation_accepts_packed_load_balance(self):
+    argv = [
+        "",
+        _BASE_CONFIG_PATH,
+        "run_name=test",
+        "attention=flash",
+        "use_tokamax_splash=True",
+        "use_jax_splash=False",
+        "context_parallel_strategy=ring",
+        "context_parallel_load_balance=True",
+        "ici_context_parallelism=2",
+        "hardware=tpu",
+        "packing=True",
+        "skip_jax_distributed_system=True",
+    ]
+    mock_devices = [unittest.mock.MagicMock(slice_index=0) for _ in range(8)]
+    with unittest.mock.patch("jax.devices", return_value=mock_devices):
+      config = pyconfig.initialize(argv)
+
+    self.assertTrue(config.context_parallel_load_balance)
+    self.assertTrue(config.packing)
+
   def test_tpu_tokamax_ring_config_validation_rejects_unsupported_configs(self):
     base_args = [
         "",
@@ -168,7 +211,6 @@ def test_tpu_tokamax_ring_config_validation_rejects_unsupported_configs(self):
         (["use_tokamax_splash=False"], ["use_tokamax_splash=True"], "use_tokamax_splash"),
         (["use_jax_splash=True"], ["use_jax_splash=False"], "use_jax_splash"),
         (["attention_type=full"], [], "global causal"),
-        (["packing=True"], ["packing=False"], "packing"),
         (
             [
                 "context_parallel_load_balance=True",