Fusing Pallas Layout Transposes and Output Projections

[prishajain1]

This PR optimizes the custom sequence-parallel attention path for WAN models by fusing layout transposes and avoiding post-attention copy overhead.

Changes:

• Pallas Kernels (custom_splash_attention.py): Modified forward kernels to write output registers directly in transposed [H, S, D] layout using SRAM transposes.
• Attention Wrappers (attention_flax.py): Bypassed post-attention jnp.swapaxes and _reshape_heads_to_head_dim transpose copy operations for custom kernels.
• Output Projection (FlaxWanAttention): Implemented a custom multi-axis dot_general contraction contracting the H (heads) and D (head features) axes directly, avoiding transposing activations before matrix multiplication.

Performance and Results:

Eliminates layout transpose overhead after the Ulysses all_to_all collective:

Hardware	Model Name	Before fixes	After fix	Gain
TPU v6e-8	WAN 2.1 14B
TPU v7x-8	WAN 2.1 14B

[github-actions]

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