Optimizing flux2klein

[amepas]

Building off of PR 456 (so there are overlapping code changes like the new reference images)

Summary

This PR improves inference speeds for Flux2 Klein Models. ~2.58x overall speedup compared to the original implementation and outperform the standard SGLang GB200 baseline.

---

Performance Benchmarks

• Resolution: 1024 × 1024 (Batch=1)
• Hardware: Single Chip (GB200 vs TPU v7-2)

Implementation	Total Latency	Latency (Excl. Image Saving)	vs. Original
Original Implementation	1131.20 ms/img	765.28 ms/img	1.0x (Baseline)
SGLang GB200 Baseline	563.90 ms/img	406.40 ms/img	~2.01x faster
Flux2Klein 9B (PR)	430.00 ms/img	321.69 ms/img	~2.63x faster

--

Optimizations

1. Tokenizer Caching Outside Warmup & Inference Loops

Instantiated and cached self.tokenizer once in FlaxFlux2KleinPipeline.__init__. String inputs are tokenized directly into contiguous PyTorch/NumPy arrays prior to device placement.

---

2. Early-Stopping Qwen3 Text Encoder Execution at Layer 27 (max_layer_to_run=27)

Standard Qwen3 contains 28+ layers, but Flux.2-Klein only extracts cross-attention text embedding representations from intermediate layers 9, 18, and 27. Running layers 28+ performed pure wasted matrix multiplication compute.

---

3. Dedicated Attention Backend for Qwen3 (Backend Isolation)

Passing shared global attention arguments into both Qwen3 and Flux caused Qwen3 to attempt loading ring Ulysses Splash kernels (ulysses_ring_custom), which are slow on the short 512-token text prompts.

---

4. Casual Attention for Text tokens

Configuring other attention backends outside of dot_product to also use causal attention, reducing the density of the text attention in the denoising loop and ensuring correctness. Also avoiding re-sharding of the text tokens in each loop prevents unnecessary copies.

---

5. Fusing 4 Denoising Steps into jax.lax.scan (_jitted_fused_denoise_loop)

Switching from a standard for loop to compile as a single graph.

---

6. Internal Denoising Loop Optimizations

1. Eliminated Redundant put_data_on_devices inside Denoising Loop (~4.2 ms saved)
2. Pre-Staged Timesteps Array on TPU Device Mesh (~0.9 ms saved). Migrated scheduler_state.timesteps and scheduler_state.sigmas sharding onto the device mesh outside the iterative loop.
3. Direct Step Indexing in Flow Match Euler Scheduler (~0.6 ms saved). Added step_index: Optional[int] = None in scheduling_flow_match_flax.py so the denoising loop passes the exact loop index directly (step_idx) to avoid _find_timestep_id searches (jnp.argmin(jnp.abs(...))) during every step.

---

7. Optimal Configuration Settings:

• attention="ulysses_ring_custom_fixed_m":
• ici_context_parallelism=2 & ulysses_shards=2
• ulysses_attention_chunks=1: Single-chunk ring stream execution.
• flash_block_sizes='{"block_q": 4608, "block_kv": 1024, "block_kv_compute": 1024}'.
• text_encoder_attention="flash"

---

8. Image Saving (uint8) Matching SGLang implementation

Implemented vector clamping and conversion directly on TPU arrays to uint8 ([0, 255]) before copying buffers to CPU for PNG encoding.

---

9. 4D Spatial Vectorization in VAE Decoding (_jitted_vae_decode)

The VAE decoder previously packed and flattened latents into intermediate 3D spaces, requiring multiple transpose and reshape operations. Refactored _jitted_vae_decode (donate_argnums=(1,)) to operate directly on 4D spatial tensors (batch_size, 32, height // 8, width // 8) using direct 4D convolution blocks.

Correctness

All smoke tests pass (4B SSIM: 0.98, 9B SSIM: 0.88)

[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 introduces a unified end-to-end inference pipeline for Flux.2-klein models (4B and 9B) on JAX+TPU, including configurations, entry point scripts, and optimizations such as concurrent AOT compilation and memory-efficient resizing. Feedback on the changes highlights several critical issues: several duplicate files were added in incorrect directories and should be removed; Qwen3 attention needs to correctly handle causal and padding masks when using non-dot_product kernels; step_index should not be a static argument in the scheduler to avoid JIT recompilation; potential index errors when text_encoder_max_layer is less than 27 should be addressed; and fractional per_device_batch_size in the smoke test should be fixed to prevent JAX sharding errors.

[Perseus14]

Could you also add SSIM metric against baseline?