feat: extend BSHD all-to-all layout to ulysses_ring attention

Applies the same BSHD-native layout used in _ulysses_attention to
_ulysses_ring_attention: all-to-all operates on the smaller post-shard
head dimension, deferring the BHSD transpose until after the collective.
Also routes ulysses_ring through the BSHD RoPE path in FlaxWanAttention.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: sagarchapara

src/maxdiffusion/models/attention_flax.py

@@ -723,20 +723,22 @@ def _ulysses_ring_attention(
   num_ring_shards = mesh.shape[ring_axis]
   num_sequence_shards = num_ulysses_shards * num_ring_shards
 
-  query, orig_q_seq_len = _reshape_data_for_flash(query, heads, num_sequence_shards)
-  key, _ = _reshape_data_for_flash(key, heads, num_sequence_shards)
-  value, _ = _reshape_data_for_flash(value, heads, num_sequence_shards)
+  query, orig_q_seq_len = _reshape_data_for_ulysses(query, heads, num_sequence_shards)
+  key, _ = _reshape_data_for_ulysses(key, heads, num_sequence_shards)
+  value, _ = _reshape_data_for_ulysses(value, heads, num_sequence_shards)
 
-  num_heads = query.shape[1]
+  num_heads = query.shape[2]
   if num_heads % num_ulysses_shards != 0:
     raise ValueError(
         "Ulysses ring attention requires the number of heads to be divisible by the Ulysses shard count, "
         f"got heads={num_heads} and ulysses_shards={num_ulysses_shards}."
     )
-  block_sizes = _select_flash_block_sizes(query, key, flash_block_sizes, dtype, "tokamax_ring")
+  block_sizes = _select_flash_block_sizes(
+      _bshd_as_bhsd_shape(query), _bshd_as_bhsd_shape(key), flash_block_sizes, dtype, "tokamax_ring"
+  )
 
-  q_axis_names = nn.logical_to_mesh_axes(axis_names_q)
-  kv_axis_names = nn.logical_to_mesh_axes(axis_names_kv)
+  q_axis_names = nn.logical_to_mesh_axes(_bshd_axis_names(axis_names_q))
+  kv_axis_names = nn.logical_to_mesh_axes(_bshd_axis_names(axis_names_kv))
 
   @functools.partial(
       jax.shard_map,
@@ -746,9 +748,13 @@ def _ulysses_ring_attention(
       check_vma=False,
   )
   def wrap_ulysses_ring_attention(query, key, value):
-    query = jax.lax.all_to_all(query, axis_name=ulysses_axis, split_axis=1, concat_axis=2, tiled=True)
-    key = jax.lax.all_to_all(key, axis_name=ulysses_axis, split_axis=1, concat_axis=2, tiled=True)
-    value = jax.lax.all_to_all(value, axis_name=ulysses_axis, split_axis=1, concat_axis=2, tiled=True)
+    query = jax.lax.all_to_all(query, axis_name=ulysses_axis, split_axis=2, concat_axis=1, tiled=True)
+    key = jax.lax.all_to_all(key, axis_name=ulysses_axis, split_axis=2, concat_axis=1, tiled=True)
+    value = jax.lax.all_to_all(value, axis_name=ulysses_axis, split_axis=2, concat_axis=1, tiled=True)
+
+    query = _bshd_to_bhsd(query)
+    key = _bshd_to_bhsd(key)
+    value = _bshd_to_bhsd(value)
 
     uses_fused_kernel = block_sizes.use_fused_bwd_kernel
     block_q_sizes = (block_sizes.block_q, block_sizes.block_q_dkv)
@@ -809,11 +815,12 @@ def wrap_ulysses_ring_attention(query, key, value):
     attention_output = vmapped_splash(query, key, value, segment_ids)
     attention_output = attention_output[:, :, :query_seq_len, :kv_size].astype(query.dtype)
 
+    attention_output = _bhsd_to_bshd(attention_output)
     return jax.lax.all_to_all(
         attention_output,
         axis_name=ulysses_axis,
-        split_axis=2,
-        concat_axis=1,
+        split_axis=1,
+        concat_axis=2,
         tiled=True,
     )
 
@@ -824,8 +831,8 @@ def wrap_ulysses_ring_attention(query, key, value):
         f" axis, batch dimension: {query.shape[0]}, devices_in_batch_sharding: {devices_in_batch_sharding}"
     )
   x = wrap_ulysses_ring_attention(query, key, value)
-  x = x[:, :, :orig_q_seq_len, :]
-  x = _reshape_heads_to_head_dim(x)
+  x = x[:, :orig_q_seq_len, :, :]
+  x = x.reshape(x.shape[0], x.shape[1], -1)
 
   return x
 
@@ -950,7 +957,7 @@ def _apply_attention(
   """Routes to different attention kernels."""
   _check_attention_inputs(query, key, value)
   seq_len_idx = 1
-  if query.ndim == 4 and attention_kernel != "ulysses":
+  if query.ndim == 4 and attention_kernel != "ulysses" and attention_kernel not in ULYSSES_RING_ATTENTION_KERNELS:
     seq_len_idx = 2
   if attention_kernel in ["flash", "tokamax_flash", "ulysses"] or attention_kernel in ULYSSES_RING_ATTENTION_KERNELS:
     can_use_flash_attention = (
@@ -1628,7 +1635,7 @@ def __call__(
 
       if rotary_emb is not None:
         with self.conditional_named_scope("attn_rope"):
-          if self.attention_op.attention_kernel == "ulysses":
+          if self.attention_op.attention_kernel == "ulysses" or self.attention_op.attention_kernel in ULYSSES_RING_ATTENTION_KERNELS:
             query_proj = _unflatten_heads_bshd(query_proj, self.heads)
             key_proj = _unflatten_heads_bshd(key_proj, self.heads)
             value_proj = _unflatten_heads_bshd(value_proj, self.heads)