From bf0abf8cd1417df270d89150e913ca42dabaf349 Mon Sep 17 00:00:00 2001 From: Mohit Khatwani Date: Thu, 15 Oct 2026 21:14:48 +0000 Subject: [PATCH 1/3] Ulysses attention: pin the attention block's input and output layout Under Ulysses the attention backward wants the hidden-state gradient tokens-minor ({1,2,0}), for the [B, S, H, D] -> [B, H, S, D] transpose ahead of the inverse all-to-all. XLA carried that layout out through the residual add into the neighbouring MoE layers, where it transposed the expert-parallel gathered token buffer (bf16[8,8192,4096], 8x the local tokens) in front of the token reduce-scatter: 8 x 512 MiB copies per scanned 4-layer block on Qwen3.5-397B (128 v7x, fsdp 32, cp 4, ep 8). A custom_vjp identity constrains the attention block's input and output, and their cotangents, to the row-major layout (with_layout_constraint, as #5685 does for in_proj_qkvz), so the relayout stays on the attention-local tensors. It applies only in training under context_parallel_strategy=ulysses. 60-layer AOT (tpu7x-256): all 16 MoE token reduce-scatters per block take {2,1,0} again (8 were transposed), and copy/transpose traffic drops by ~54 GiB per microbatch. --- src/maxtext/layers/attentions.py | 43 ++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/src/maxtext/layers/attentions.py b/src/maxtext/layers/attentions.py index 4b78d71078..1566b4c468 100644 --- a/src/maxtext/layers/attentions.py +++ b/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 From 3d5b76d477574059de144a8d0e51dee1bc0bcba2 Mon Sep 17 00:00:00 2001 From: Mohit Khatwani Date: Thu, 15 Oct 2026 21:14:57 +0000 Subject: [PATCH 2/3] Ulysses attention: let the remat policy save the exchanged Q/K/V (ulysses_qkv) The Splash backward needs the head-sharded Q/K/V that the Ulysses all-to-all produces. Nothing saved them, so the remat backward replayed the forward exchanges to rebuild them: on Qwen3.5-397B that is the Q, K and V all-to-alls of all 15 attention layers again per microbatch, one of them with almost no compute to hide behind. Tag the exchanged tensors "ulysses_qkv" and add the matching remat key (default remat, so nothing changes unless it is set). ulysses_qkv=offload is the setting to use: on the 128-chip v7x trainer (60-layer AOT, tpu7x-256) it keeps all all-to-alls async, leaves only the output exchange to replay, and lowers the device peak, at 4.69 GiB of host memory per device. ulysses_qkv=device costs +8.6 GiB of HBM and, with no replayed work left to overlap them with, XLA runs both backward all-to-alls synchronously. --- src/maxtext/configs/base.yml | 1 + src/maxtext/configs/types.py | 10 ++++++++++ src/maxtext/layers/attention_op.py | 5 +++++ 3 files changed, 16 insertions(+) diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index bba0085bab..4c96afa5d5 100644 --- a/src/maxtext/configs/base.yml +++ b/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') diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 7f7171de76..b6762330d9 100644 --- a/src/maxtext/configs/types.py +++ b/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", diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index 38f93b23fb..dadd81cdb4 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -2469,6 +2469,11 @@ def wrap_flash_attention( 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) + # 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. From 59904c5abd8585e9680b2f6a1c6ef46502737f99 Mon Sep 17 00:00:00 2001 From: Mohit Khatwani Date: Thu, 15 Oct 2026 21:15:05 +0000 Subject: [PATCH 3/3] Ulysses attention: exchange Q, K and V in one all-to-all XLA TPU does not combine the three same-axis all-to-alls (checked in compiled HLO), and the K/V ones (1 KV head per rank after the replication) are small enough to be mostly latency: their forward overlap windows were 4 and 8 instructions. ulysses_all_to_all_qkv lays rank r's block of each head axis 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, and its transpose is one all-to-all in the backward too. It falls back to three exchanges when the tensors can't be concatenated. Bit-identical to the separate exchanges on a CPU mesh (loss and dQ/dK/dV diff 0.0). On the 60-layer 397B AOT the Ulysses all-to-alls per microbatch drop from 183 to 93. The test's HLO count now expects 2 forward all-to-alls (packed Q/K/V, output) instead of 4. --- .../kernels/attention/ulysses_attention.py | 28 +++++++++++++++++++ src/maxtext/layers/attention_op.py | 6 ++-- .../unit/ulysses_kv_head_replication_test.py | 5 ++-- 3 files changed, 34 insertions(+), 5 deletions(-) diff --git a/src/maxtext/kernels/attention/ulysses_attention.py b/src/maxtext/kernels/attention/ulysses_attention.py index aff3a2b1a9..fbb3b41767 100644 --- a/src/maxtext/kernels/attention/ulysses_attention.py +++ b/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) diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index dadd81cdb4..5d063eb132 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -2466,9 +2466,9 @@ 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") diff --git a/tests/unit/ulysses_kv_head_replication_test.py b/tests/unit/ulysses_kv_head_replication_test.py index 36d0e7df99..62d6e7158c 100644 --- a/tests/unit/ulysses_kv_head_replication_test.py +++ b/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