🟡 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.

🟠 The attention_mask parameter is defined but not used in the _ulysses_ring_attention implementation. Other attention kernels in this file use this mask to handle sequence padding or specific conditioning (like I2V). This feature should be implemented to ensure consistency and correctness for models that rely on attention masking.

🟠 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

thanks, fixed it.

🟡 This axis rule shards the cross-attention KV length as None (replicated), which matches SEQUENCE_PARALLEL_AXIS_RULES. Given that ulysses_ring falls back to tokamax_flash for cross-attention, this is consistent and ensures compatibility when the encoder sequence is not sharded across the context mesh.

🟢 It might be helpful to keep or restore the comment explaining the output shape of _unflatten_heads, as it clarifies the tensor layout for anyone reading this section.

🟡 The use of np.arange and tokamax_splash_attention_mask.NumpyMask inside the shard_map function assumes that q_padded_len, total_kv_len, query_seq_len, and key_seq_len are static values at trace time. This is currently true for TPU workloads in this project, but if dynamic shapes are introduced in the future, this would need to transition to jnp and a dynamic mask implementation.

🟢 Excellent and thorough validation of ulysses_shards. Ensuring it divides the context_axis and that heads are divisible prevents difficult-to-debug sharding errors later in the execution.

🟢 These new tests are very well-structured, covering layout round-trips, global KV padding masking, and various validation error cases. This provides high confidence in the new attention kernel's correctness.