LTX-2 minor fixes

[prishajain1]

This PR introduces the following changes for the LTX2 pipeline:

• Added Context parallelism replication for audio stream in transformer (this resolved some audio related issues occurring in particular sharding configurations on ironwood and seems to reduce the latency for other cases. Since audio stream is only 126 tokens, this shouldn't be an issue even if replicated)
• Replicated VAE params and latents before call to decode method. With this change, single video latency dropped by 3x (observed no OOM in different configs caused by this change, tested on trillium and ironwood 8 chips)
• Added flash_min_seq_length parameter to config file, now the choice of attention is controlled by this parameter. Experimentally observed that having dot attention for audio to audio self attn and video to audio cross attn and flash for audio to video provided more speed up. We had to hardcode "flash" as the attention choice for audio to video, because the way the attention mechanism is currently decided bases on flash min seq length compares the seq length for all q, k/v and flash is chosen only if all have seq length > flash min seq length. Experimentally we observed more speed ups if audio to video was kept in flash instead of dot.
• Added JAX annotations for profiling
• jitted the step() function in scheduling_flow_match_flax.py
• Converted the timesteps ‘t’ from static Python integer/float to a dynamic JAX tensor (to prevent compilation of a program for each and every timestep)

Results

We performed the experiments below for dp = 2, cp = 2, tp = 2, fsdp = 1, per device batch size = 0.125

Hardware	Status	Time (s)	Video Link
v7x-8	Current main	105.14	View Video
	After fix	28.63	View Video
v6e-8	Current main	138.24	View Video
	After fix	58.74	View Video

[github-actions]

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

[Perseus14]

Overall looks good. Few minor comments.

Please add few results comparing baseline and this branch to get a better idea of the improvements.