Ring + Ulysses 2D context parallelism

[csgoogle]

Description
This PR adds support for a new ulysses_ring attention mode for WAN models. The implementation keeps the public sequence sharding on the existing context mesh axis, then internally reshapes that axis into private ring and ulysses axes so the attention path can combine Ulysses all-to-all head sharding with ring-based KV rotation.

Changes

• Adds ulysses_ring attention kernel registration and routing.
• Introduces internal mesh reshaping from context into hidden ring and ulysses axes.
• Adds ulysses_shards config plumbing through WAN pipeline, WAN transformer blocks, and attention ops.
• Adds logical axis rules for Ulysses ring attention.
• Updates WAN config files to document ulysses_ring support and add ulysses_shards.
• Adds tests covering layout round-trip behavior and validation errors for invalid Ulysses shard settings.

Testing

• Added unit tests in src/maxdiffusion/tests/attention_test.py for ulysses_ring behavior and validation.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[github-actions]

🤖 Hi @csgoogle, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This Pull Request introduces a hybrid Ulysses + Ring 2D context parallelism for attention, which is a significant addition for handling long sequences in large-scale diffusion models. The implementation correctly sets up the internal mesh and handles the necessary all_to_all communications. However, there are some concerns regarding the masking logic and the handling of the attention_mask parameter in the new kernel that should be addressed to ensure correctness across all use cases.

🔍 General Feedback

• Positive Highlights: The refactoring of Attention classes to use a more flexible attention_config dictionary is a great improvement for maintainability and reduces the complexity of method signatures.
• Robustness: The validation checks for ulysses_shards and mesh axes are well-implemented and will help users avoid misconfiguration.
• Testing: New tests for the Ulysses-Ring layout are included, although more comprehensive tests covering the actual attention computation with padding would further improve confidence.

[github-actions]

🟡 block_q and block_kv were increased from 512 to 2048 in this config. Is this change required for the new ulysses_ring attention kernel or is it a general performance optimization? It might be worth documenting why this specific model size received this update, or applying it to others if it's a general improvement.

[github-actions]

🟠 The use of np.tile(kv_shard_valid, num_ring_shards) assumes that every shard in the ring has an identical valid/padding pattern. In sequence parallelism, padding is typically only at the end of the global sequence (affecting only the last shard). Tiling a local mask will incorrectly mask valid tokens in earlier shards or fail to mask padding tokens in the last shard correctly.

Consider calculating the mask based on global sequence positions:

    # Example logic for global masking
    ring_idx = jax.lax.axis_index(ring_axis)
    q_global_start = ring_idx * query_seq_len
    # ... logic to create a mask that correctly accounts for orig_q_seq_len across the ring

[csgoogle]

thanks, fixed it.

[github-actions]

🤖 Hi @csgoogle, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This PR successfully implements the ulysses_ring attention mode, a sophisticated 2D parallelism strategy that combines Ulysses all-to-all head sharding with Ring KV rotation. The implementation is clean, well-integrated into the existing WAN model infrastructure, and includes a comprehensive suite of unit tests verifying both functional correctness and validation logic.

🔍 General Feedback

• Elegant Mesh Reshaping: The internal reshaping of the context mesh axis into private ring and ulysses axes within shard_map is a great architectural choice, keeping the public configuration simple.
• Robust Masking: The global sequence masking logic in the hybrid kernel correctly handles padding, which is critical for accuracy in sequence-parallel settings.
• Thorough Testing: The addition of detailed round-trip and validation tests in attention_test.py is excellent and ensures the new feature is robust against edge cases.

[github-actions]

🤖 Hi @csgoogle, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This PR introduces a robust implementation of the ulysses_ring attention mode, effectively combining Ulysses all-to-all head sharding with Ring-based KV rotation for WAN models. The implementation is well-architected, leveraging internal mesh reshaping and shard_map to provide a clean public interface while handling complex 2D parallelism internally.

🔍 General Feedback

• High Quality Implementation: The use of internal axes (ring, ulysses) within shard_map is a clever and effective way to manage 2D parallelism.
• Comprehensive Testing: The added unit tests are thorough and cover important edge cases, including global sequence padding and validation logic.
• Plumbing Consistency: The configuration plumbing is consistently applied across all WAN model variants and pipelines.
• Safety: The fallback mechanism for cross-attention and the rigorous validation of shard divisibility ensure the stability of the new feature.

[github-actions]

🟡 The logic for prepending `new_rules` to `logical_axis_rules` might result in a redundant or messy list if the rules are already partially present at the end of the list. While functionally correct due to JAX's first-match rule priority, it would be cleaner to ensure a unique set of rules is maintained without duplicates when they are combined at the end of this block.

[github-actions]

🟡 This silent fallback to `tokamax_flash` for cross-attention when `ulysses_ring` is requested is a safe choice but could be confusing for users. It might be helpful to log a warning when this swap occurs to inform the user that ring attention is only being applied to self-attention.

[github-actions]

🟢 Using `np.arange` to create masks inside a function that will be JIT-compiled works fine in JAX because shapes are static, but it results in a concrete numpy array being baked into the graph. For very large sequence lengths, this might increase the graph size. This is acceptable for current usage, but worth keeping in mind if sequence lengths grow significantly.