Minor improvement to Flux models

[Perseus14]

Summary

This PR implements comprehensive performance profiling, kernel fusion, instruction scheduling, and latency optimizations for the Flux model family (FLUX.1-dev, FLUX.1-schnell, and FLUX.2-klein-4B/9B) on Google Cloud TPU v6e (Trillium).

Across cumulative optimization passes on TPU v6e-8 (1024×1024 resolution):

• FLUX.1-dev (50 steps): End-to-end steady-state latency dropped from 8.904 s → 7.780 s (-1.124 s per image, 152.84 ms/step, 6.54 iterations/sec).
• FLUX.1-schnell (4 steps): End-to-end latency reduced from 0.860 s → 0.749 s (< 0.75 s total / image).
• FLUX.2-klein-4B (4 steps): Achieved sub-400ms generation at 0.394 s (63.92 ms/step, 15.64 iterations/sec).
• FLUX.2-klein-9B (4 steps): Achieved sub-700ms generation at 0.686 s (137.05 ms/step, 7.30 iterations/sec).

---

Key Changes

1. Cauchy-Schwarz Fixed-$m$ Attention Kernel (ulysses_custom_fixed_m):

   • Replaced default attention with ulysses_custom_fixed_m utilizing Cauchy-Schwarz bound softmax and base-2 exponential scaling (use_base2_exp: True).
   • Eliminates the in-kernel online running-max rescale loop and reciprocal spills on TPU v6e VMEM.

2. Pallas Multi-Head Tile Batching (heads_per_tile: 3 / 4):

   • Batched all local attention heads per TPU chip into a single unified Pallas grid invocation.
   • Eliminates grid launch dispatch overhead and optimizes VMEM memory bandwidth.

3. Pallas LP LLO Instruction Scheduler (use_experimental_scheduler: True):

   • Activated XLA_TPU_FORCE_LP_LLO_SCHEDULER via Pallas compiler parameters.
   • Uses an exact Integer/Linear Programming scheduler to optimally interleave VMEM tile loads/stores, vector operations, and MXU matrix multiplications.

4. Fused Direct RoPE & Transpose Elimination:

   • Simplified Rotary Position Embedding (apply_rope) to run directly on (B, L, H, D) sequence-first tensors in attention_flax.py and transformer_flux.py.
   • Eliminated 228 redundant transpose and swapaxes ops across 57 layers.

5. Fused AdaLN-Zero Math:

   • Reformulated $(x - \mu) \cdot \text{inv_std} \cdot (1 + s) + b$ into a single fused multiply-add $x \cdot \text{scale} + \text{shift}$ in normalization_flax.py and transformer_flux.py.
   • Removes redundant elementwise subtractions across all 57 double and single stream blocks.

6. Native BF16 VAE Latent Decode:

   • Maintained latents in native bfloat16 during vae_decode in generate_flux.py and flux2klein_pipeline.py.
   • Executes VAE convolution blocks directly on TPU MXUs in BF16, cutting VAE decode time from 52.6 ms → 24.1 ms (+54.2% faster).

7. Dynamic Text Truncation & Pre-Broadcasted Global Modulation:

   • Shortened default prompt max sequence length to 256 tokens in base_flux_dev.yml and base_flux2klein.yml for standard prompts.
   • Pre-expanded modulation embeddings along sequence dimension at top-level model scope, removing 120 redundant jnp.expand_dims calls per pass.

8. Module Renaming:

   • Renamed transformer_flux_flax.py → transformer_flux.py for module naming consistency across MaxDiffusion and updated all imports.

---

Performance Comparison Matrix (Cloud TPU v6e-8, 1024×1024 Resolution)

Workload / Model	Baseline	Optimized	Speedup / Gain	Step Latency & Rate
FLUX.1-dev (50 steps)	8.904 s	7.780 s	+12.6% faster (-1.124 s)	152.84 ms/step (6.54 it/s)
FLUX.1-schnell (4 steps)	0.860 s	0.749 s	+12.9% faster (-0.111 s)	152.84 ms/step (6.54 it/s)
FLUX.2-klein-4B (4 steps)	--	0.394 s	Sub-400ms E2E	63.92 ms/step (15.64 it/s)
FLUX.2-klein-9B (4 steps)	--	0.686 s	Sub-700ms E2E	137.05 ms/step (7.30 it/s)

Component-Level Benchmarks

• 50-Step Diffusion Loop (FLUX.1-dev): 8.663 s → 7.642 s (-1.021 s)
• 4-Step Diffusion Loop (FLUX.1-schnell): 0.693 s → 0.611 s (-82 ms)
• VAE Latent Decode: 52.6 ms → 24.1 ms (+54.2% faster)
• Text Encoding (T5 + CLIP): 0.203 s → 0.114 s (+43.8% faster)

---

Verification & Testing

• Verified full end-to-end generation and output quality on Cloud TPU v6e-8.
• Passed formatting (pyink src/maxdiffusion --pyink-indentation=2 --line-length=125).
• Passed linter checks (ruff check src/).

[github-actions]

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

[gemini-code-assist]

Code Review

This pull request renames the Flux transformer module, updates configurations to use Ulysses attention with custom block sizes, optimizes layer normalization calculations, and adds profiling support. A critical issue was identified in attention_flax.py where a missing null check on image_rotary_emb could lead to an AttributeError when it is None.