Ulysses attention: layout pin, saved exchanged Q/K/V, packed Q/K/V all-to-all (on #5779)

src/maxtext/configs/base.yml

@@ -557,6 +557,7 @@ remat_policy: 'full'
 decoder_layer_input: 'device' # this tensor cannot be rematerialized - it serves as periodic checkpoints that act as the remat start points
 indexer_cutoff_threshold: 'remat' # Remat policy for indexer cutoff threshold (can be set to 'remat', 'offload', or 'device')
 context: 'remat' # From https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/jax/attention.py#L581-L583
+ulysses_qkv: 'remat' # Q/K/V after the Ulysses all-to-all (context_parallel_strategy='ulysses'); device/offload skips replaying it in the backward
 gdn: 'remat' # Remat policy for GatedDeltaNet recurrence kernel and residuals (can be set to 'remat', 'device', or 'offload')
 gdn_conv: 'remat' # Remat policy for GatedDeltaNet 1D conv activations (can be set to 'remat', 'device', or 'offload')
 gdn_states: 'remat' # Partial GDN remat: keep only core_attn_out + t_inv + chunk_states (what the bwd Pallas kernel needs) so the fwd kernel is not replayed in bwd, without saving all GDN residuals like 'gdn' ('remat', 'device', or 'offload')

src/maxtext/configs/types.py

@@ -2614,6 +2614,14 @@ class RematAndOffload(BaseModel):
       RematLocation.REMAT, description="Remat policy for the indexer cutoff threshold (shape: [batch, seq_len])."
   )
   context: RematLocation = Field(RematLocation.REMAT, description="Remat policy for the attention context.")
+  ulysses_qkv: RematLocation = Field(
+      RematLocation.REMAT,
+      description=(
+          "Remat policy for Q/K/V after the Ulysses all-to-all (context_parallel_strategy='ulysses'). Saving them"
+          " (device or offload) keeps the backward from replaying the all-to-all; costs the exchanged Q/K/V per"
+          " attention layer."
+      ),
+  )
   gdn: RematLocation = Field(
       RematLocation.REMAT,
       description="Remat policy for GatedDeltaNet (GDN) recurrence kernel and residuals (remat, device, offload).",
@@ -5807,6 +5815,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
           "decoder_layer_input",
           "indexer_cutoff_threshold",
           "context",
+          "ulysses_qkv",
           "gdn",
           "gdn_conv",
           "gdn_states",
@@ -7376,6 +7385,7 @@ def set_derived_values_and_validate(self) -> "RLConfig":
           "decoder_layer_input",
           "indexer_cutoff_threshold",
           "context",
+          "ulysses_qkv",
           "gdn",
           "gdn_conv",
           "gdn_states",

src/maxtext/kernels/attention/ulysses_attention.py

@@ -247,6 +247,34 @@ def ulysses_all_to_all(tensor: Any, ulysses_axis: str):
   return jax.lax.all_to_all(tensor, ulysses_axis, split_axis=1, concat_axis=2, tiled=True)
 
 
+def ulysses_all_to_all_qkv(query: Any, key: Any, value: Any, ulysses_axis: str, ulysses_size: int):
+  """`ulysses_all_to_all` of Q, K and V as one collective.
+
+  XLA does not combine the three exchanges, and the K/V ones are small enough to be mostly latency. Rank r's block
+  of each tensor's head axis is laid out next to the others ([its H_q/U query heads, its H_kv/U key heads, its
+  H_kv/U value heads]), so one tiled all-to-all over the concatenated head axis hands every rank the same heads as
+  three separate exchanges, and the backward is one all-to-all as well. Falls back to three exchanges when the
+  tensors cannot be concatenated (different batch, sequence, head dim or dtype).
+  """
+  b, hq, s, d = query.shape
+  hk = key.shape[1]
+  if key.shape != value.shape or (key.shape[0], key.shape[2], key.shape[3]) != (b, s, d) or not (
+      query.dtype == key.dtype == value.dtype and hq % ulysses_size == 0 and hk % ulysses_size == 0
+  ):
+    return tuple(ulysses_all_to_all(x, ulysses_axis) for x in (query, key, value))
+  nq, nk = hq // ulysses_size, hk // ulysses_size
+  packed = jnp.concatenate(
+      [
+          query.reshape(b, ulysses_size, nq, s, d),
+          key.reshape(b, ulysses_size, nk, s, d),
+          value.reshape(b, ulysses_size, nk, s, d),
+      ],
+      axis=2,
+  ).reshape(b, ulysses_size * (nq + 2 * nk), s, d)
+  packed = ulysses_all_to_all(packed, ulysses_axis)
+  return packed[:, :nq], packed[:, nq : nq + nk], packed[:, nq + nk :]
+
+
 def inverse_ulysses_all_to_all(tensor: Any, ulysses_axis: str):
   """Moves `[B, H/U, S, D]` back to `[B, H, S/U, D]`."""
   return jax.lax.all_to_all(tensor, ulysses_axis, split_axis=2, concat_axis=1, tiled=True)

src/maxtext/layers/attention_op.py

@@ -2466,9 +2466,14 @@ def wrap_flash_attention(
         kv_replication = ulysses_attention.kv_head_replication(key.shape[1], self.mesh.shape[context_axis])
         key = ulysses_attention.repeat_kv_heads_for_ulysses(key, kv_replication)
         value = ulysses_attention.repeat_kv_heads_for_ulysses(value, kv_replication)
-        query = ulysses_attention.ulysses_all_to_all(query, context_axis)
-        key = ulysses_attention.ulysses_all_to_all(key, context_axis)
-        value = ulysses_attention.ulysses_all_to_all(value, context_axis)
+        query, key, value = ulysses_attention.ulysses_all_to_all_qkv(
+            query, key, value, context_axis, self.mesh.shape[context_axis]
+        )
+        # Saving the exchanged Q/K/V (`ulysses_qkv=device|offload` under remat_policy=custom) keeps the backward
+        # from replaying the all-to-all to rebuild them for the Splash backward.
+        query = checkpoint_name(query, "ulysses_qkv")
+        key = checkpoint_name(key, "ulysses_qkv")
+        value = checkpoint_name(value, "ulysses_qkv")
         if decoder_segment_ids_q is not None:
           # Q and KV segment IDs are the same tensor in this train-only
           # self-attention path, so one gather serves both kernel operands.

src/maxtext/layers/attentions.py

@@ -75,6 +75,40 @@
 # pytype: disable=attribute-error
 
 
+def _row_major(x):
+  try:
+    from jax.experimental.layout import Layout, with_layout_constraint  # pylint: disable=import-outside-toplevel
+  except ImportError:  # older JAX
+    return x
+  if not isinstance(x, jax.core.Tracer):
+    return x
+  return with_layout_constraint(x, Layout(major_to_minor=tuple(range(x.ndim))))
+
+
+@jax.custom_vjp
+def _pin_row_major_with_cotangent(x):
+  """Identity that constrains `x` and its cotangent to the default (row-major) XLA layout.
+
+  Under Ulysses the attention backward wants the hidden-state gradient tokens-minor, for the
+  [B, S, H, D] -> [B, H, S, D] transpose ahead of the inverse all-to-all. Without a pin XLA carries that
+  layout out through the residual add into the neighbouring MoE layers, where it transposes the
+  expert-parallel-gathered token buffer (8x the local tokens) in front of the token reduce-scatter. Pinning
+  the attention block's input and output keeps the relayout on the attention-local tensors.
+  """
+  return _row_major(x)
+
+
+def _pin_row_major_fwd(x):
+  return _row_major(x), None
+
+
+def _pin_row_major_bwd(_, g):
+  return (_row_major(g),)
+
+
+_pin_row_major_with_cotangent.defvjp(_pin_row_major_fwd, _pin_row_major_bwd)
+
+
 @dataclasses.dataclass(repr=False)
 class L2Norm(nnx.Module):
   """
@@ -1161,6 +1195,13 @@ def __call__(
           inputs_q, inputs_kv, inputs_positions, decoder_segment_ids, use_shared_kv=use_shared_kv
       )
 
+    pin_layout = (
+        model_mode == MODEL_MODE_TRAIN and self.config.context_parallel_strategy.lower() == "ulysses" and not use_shared_kv
+    )
+    if pin_layout:
+      same_kv = inputs_kv is inputs_q
+      inputs_q = _pin_row_major_with_cotangent(inputs_q)
+      inputs_kv = inputs_q if same_kv else _pin_row_major_with_cotangent(inputs_kv)
     inputs_q = self._maybe_shard_with_logical(inputs_q, input_axis_names)
     inputs_kv = self._maybe_shard_with_logical(inputs_kv, input_axis_names)
     qkv_sharding = create_sharding(self.mesh, input_axis_names)
@@ -1284,6 +1325,8 @@ def __call__(
       out = self.attention_op.reorder_for_attention_load_balance(out, to_natural=True)
     if getattr(self.config, "distill_beta", 0.0) > 0.0:
       self.sow(nnx.Intermediate, "out_projection_activations", out)
+    if pin_layout:
+      out = _pin_row_major_with_cotangent(out)
     out = checkpoint_name(out, "out_proj")
     return out, kv_cache
 

tests/unit/ulysses_kv_head_replication_test.py

@@ -314,8 +314,9 @@ def dense(q, k, v):
 
   with nn_partitioning.axis_rules(cfg.logical_axis_rules), jax.set_mesh(mesh):
     hlo = jax.jit(flash("ulysses")).lower(*args).compile().as_text()
-  # q, k, v in and the output back; no K/V all-gather (the one all-gather is the int32 segment IDs).
-  assert hlo.count(" all-to-all(") == 4, hlo.count(" all-to-all(")
+  # q, k, v in as one packed exchange and the output back; no K/V all-gather (the one all-gather is the int32
+  # segment IDs).
+  assert hlo.count(" all-to-all(") == 2, hlo.count(" all-to-all(")
   assert hlo.count(" all-gather(") + hlo.count(" all-gather-start(") == 1