Fix LTX2 cross-attention axis names and ring attention fallback

Author: Perseus14

src/maxdiffusion/models/ltx2/attention_ltx2.py

@@ -436,15 +436,28 @@ def __init__(
           dtype=dtype,
       )
 
+    is_self_attention = context_dim is None
+    if not is_self_attention:
+      if attention_kernel in ("tokamax_ring", "tokamax_ring_custom", "ulysses_ring"):
+        attention_kernel = "tokamax_flash"  # do not use ring attention for cross attention
+      if attention_kernel in ("ulysses_ring_custom", "ulysses_ring_custom_bidir"):
+        attention_kernel = "ulysses_custom"  # plain ulysses (no ring) for cross attention
+
+      axis_names_q = (common_types.BATCH, common_types.CROSS_ATTN_HEAD, common_types.CROSS_ATTN_Q_LENGTH, common_types.D_KV)
+      axis_names_kv = (common_types.BATCH, common_types.CROSS_ATTN_HEAD, common_types.CROSS_ATTN_KV_LENGTH, common_types.D_KV)
+    else:
+      axis_names_q = (common_types.BATCH, common_types.SELF_ATTN_HEAD, common_types.SELF_ATTN_Q_LENGTH, common_types.D_KV)
+      axis_names_kv = (common_types.BATCH, common_types.SELF_ATTN_HEAD, common_types.SELF_ATTN_KV_LENGTH, common_types.D_KV)
+
     self.attention_op = NNXAttentionOp(
         mesh=mesh,
         attention_kernel=attention_kernel,
         scale=dim_head**-0.5,
         heads=heads,
         dim_head=dim_head,
         dtype=dtype,
-        axis_names_q=(common_types.BATCH, common_types.SELF_ATTN_HEAD, common_types.SELF_ATTN_Q_LENGTH, common_types.D_KV),
-        axis_names_kv=(common_types.BATCH, common_types.SELF_ATTN_HEAD, common_types.SELF_ATTN_KV_LENGTH, common_types.D_KV),
+        axis_names_q=axis_names_q,
+        axis_names_kv=axis_names_kv,
         flash_block_sizes=flash_block_sizes,
         flash_min_seq_length=flash_min_seq_length,
         ulysses_shards=ulysses_shards,