Refactor custom_splash_attention for improved clarity

[Perseus14]

Overview

This PR refactors and cleans up the custom Splash attention kernels by properly encapsulating the inner-VPU tiling step (block_kv_compute_in) and removing redundant parameters. It also aligns the newly introduced custom Ring Attention implementations with these cleaned-up signatures.

Changes

• custom_splash_attention.py:
  • Added block_kv_compute_in directly to the _BlockSizes dataclass, streamlining parameter passing.
  • Removed unused high-level attention wrappers (tpu_custom_attention, make_custom_splash_sdpa) since orchestration is now fully handled in attention_flax.py.
  • Cleaned up _flash_attention_kernel and _splash_attention_forward_ring by stripping out redundant arguments like bq, q_seq_len, and explicit bkv_compute_in passes.
• ring_attention_kernel.py:
  • Updated the function signatures for make_custom_ring_attention, _custom_ring_attention_forward, and _custom_bidirectional_ring_forward to drop the explicit bkv_compute_in argument, correctly extracting it from block_sizes instead.
• attention_flax.py:
  • Updated all instantiations of custom_splash._BlockSizes to pass block_kv_compute_in.
  • Fixed downstream calls to make_splash_mha and make_custom_ring_attention to remove the now-redundant bkv_compute_in kwarg.

Impact

• Reduces parameter bloat across the low-level Pallas kernels.
• Ensures API consistency across standard Ulysses attention and Tokamax Ring Attention.

[github-actions]

e2e testgrid: https://8bcf50593faf4ea38060e236169827e5-dot-us-central1.composer.googleusercontent.com/dags/maxdiffusion_tpu_e2e/grid

[eltsai]

Thanks for cleaning this up @Perseus14 ! LGTM

[Perseus14]

@eltsai Could you check again? I have made minor cleanup to ring attention implementation as well