[Flux] TPU Optimizations & Speedups for Flux.1-dev, Flux.1-schnell, and Flux.2-klein

- Implemented performance profiling, kernel fusion, instruction scheduling, and latency optimizations for the Flux model family (Flux.1-dev, Flux.1-schnell, Flux.2-klein-4B/9B) on Cloud TPU v6e (Trillium).
- Enabled ulysses_custom_fixed_m attention kernel with base-2 exponential scaling (use_base2_exp: True).
- Tuned Flash Block Sizes and Pallas multi-head tile batching (heads_per_tile: 3/4).
- Enabled LP LLO instruction scheduler in Pallas (use_experimental_scheduler: True).
- Fused Rotary Position Embeddings (RoPE) directly on (B, L, H, D) layout, eliminating 228 redundant transpose/swapaxes ops.
- Fused LayerNorm scale and shift calculations into single multiply-add ops.
- Maintained latents in native bfloat16 for VAE decoding, cutting VAE decode time by 54.2%.
- Fixed null check for image_rotary_emb in FlaxFluxAttention to prevent AttributeError.

Author: Perseus14

src/maxdiffusion/__init__.py

@@ -367,7 +367,7 @@
   _import_structure["models.controlnet_flax"] = ["FlaxControlNetModel"]
   _import_structure["models.modeling_flax_utils"] = ["FlaxModelMixin"]
   _import_structure["models.unet_2d_condition_flax"] = ["FlaxUNet2DConditionModel"]
-  _import_structure["models.flux.transformers.transformer_flux_flax"] = ["FluxTransformer2DModel"]
+  _import_structure["models.flux.transformers.transformer_flux"] = ["FluxTransformer2DModel"]
   _import_structure["models.vae_flax"] = ["FlaxAutoencoderKL"]
   _import_structure["models.ltx_video.transformers.transformer3d"] = ["Transformer3DModel"]
   _import_structure["pipelines"].extend(["FlaxDiffusionPipeline"])
@@ -444,7 +444,7 @@
     from .models.controlnet_flax import FlaxControlNetModel
     from .models.modeling_flax_utils import FlaxModelMixin
     from .models.unet_2d_condition_flax import FlaxUNet2DConditionModel
-    from .models.flux.transformers.transformer_flux_flax import FluxTransformer2DModel
+    from .models.flux.transformers.transformer_flux import FluxTransformer2DModel
     from .models.ltx_video.transformers.transformer3d import Transformer3DModel
     from .models.vae_flax import FlaxAutoencoderKL
     from .pipelines import FlaxDiffusionPipeline

src/maxdiffusion/checkpointing/flux_checkpointer.py

@@ -27,7 +27,7 @@
     FlaxAutoencoderKL,
     max_logging,
 )
-from maxdiffusion.models.flux.transformers.transformer_flux_flax import FluxTransformer2DModel
+from maxdiffusion.models.flux.transformers.transformer_flux import FluxTransformer2DModel
 from ..pipelines.flux.flux_pipeline import FluxPipeline
 
 from transformers import (CLIPTokenizer, FlaxCLIPTextModel, FlaxT5EncoderModel, AutoTokenizer)

src/maxdiffusion/configs/base_flux2klein.yml

@@ -36,7 +36,7 @@ flux_name: "flux2klein"
 scale_shift_order: "scale_shift"
 use_latents: False
 latents_path: ""
-max_sequence_length: 512
+max_sequence_length: 256
 time_shift: True
 base_shift: 0.5
 max_shift: 1.15
@@ -63,18 +63,31 @@ jit_initializers: True
 # Set true to load weights from pytorch
 from_pt: True
 split_head_dim: True
-attention: 'flash' # Supported attention: dot_product, flash, cudnn_flash_te
+attention: 'ulysses_custom_fixed_m' # Supported attention: dot_product, flash, cudnn_flash_te
+use_base2_exp: True
+use_experimental_scheduler: True
 # If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens.
 # Else we do not pass in segment ids and on vpu bound hardware like trillium this is faster.
 # However, when padding tokens are significant, this will lead to worse quality and should be set to True.
-mask_padding_tokens: True 
+mask_padding_tokens: False 
 # Maxdiffusion has 2 types of attention sharding strategies:
 # 1. attention_sharding_uniform = True : same sequence sharding rules applied for q in both (self and cross attention)
 # 2. attention_sharding_uniform = False : Heads are sharded uniformly across devices for self attention while sequence is sharded
 #    in cross attention q.
 attention_sharding_uniform: True 
 
-flash_block_sizes: {}
+flash_block_sizes: {
+  "block_q" : 4352,
+  "block_kv_compute" : 1024,
+  "block_kv" : 1024,
+  "block_kv_compute_in" : 1024,
+  "heads_per_tile" : 3,
+  "block_q_dkv" : 4352,
+  "block_kv_dkv" : 1024,
+  "block_kv_dkv_compute" : 1024,
+  "block_q_dq" : 4352,
+  "block_kv_dq" : 1024
+}
 # GroupNorm groups
 norm_num_groups: 32
 
@@ -148,7 +161,7 @@ logical_axis_rules: [
                       ['out_channels', 'tensor'],
                       ['conv_out', 'fsdp'],
                     ]
-data_sharding: [['data', 'fsdp', 'context', 'tensor']]
+data_sharding: [['data', 'fsdp']]
 
 # One axis for each parallelism type may hold a placeholder (-1)
 # value to auto-shard based on available slices and devices.
@@ -203,7 +216,7 @@ num_train_epochs: 1
 seed: 0
 output_dir: 'output/'
 output_name: "flux2klein_generated_image.png"
-per_device_batch_size: 1
+per_device_batch_size: 1.0
 
 warmup_steps_fraction: 0.1
 learning_rate_schedule_steps: -1 # By default the length of the schedule is set to the number of steps.

src/maxdiffusion/configs/base_flux2klein_9B.yml

@@ -63,18 +63,31 @@ jit_initializers: True
 # Set true to load weights from pytorch
 from_pt: True
 split_head_dim: True
-attention: 'flash' # Supported attention: dot_product, flash, cudnn_flash_te
+attention: 'ulysses_custom_fixed_m' # Supported attention: dot_product, flash, cudnn_flash_te
+use_base2_exp: True
+use_experimental_scheduler: True
 # If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens.
 # Else we do not pass in segment ids and on vpu bound hardware like trillium this is faster.
 # However, when padding tokens are significant, this will lead to worse quality and should be set to True.
-mask_padding_tokens: True 
+mask_padding_tokens: False 
 # Maxdiffusion has 2 types of attention sharding strategies:
 # 1. attention_sharding_uniform = True : same sequence sharding rules applied for q in both (self and cross attention)
 # 2. attention_sharding_uniform = False : Heads are sharded uniformly across devices for self attention while sequence is sharded
 #    in cross attention q.
 attention_sharding_uniform: True 
 
-flash_block_sizes: {}
+flash_block_sizes: {
+  "block_q" : 4864,
+  "block_kv_compute" : 1024,
+  "block_kv" : 1024,
+  "block_kv_compute_in" : 1024,
+  "heads_per_tile" : 3,
+  "block_q_dkv" : 4864,
+  "block_kv_dkv" : 1024,
+  "block_kv_dkv_compute" : 1024,
+  "block_q_dq" : 4864,
+  "block_kv_dq" : 1024
+}
 # GroupNorm groups
 norm_num_groups: 32
 
@@ -148,7 +161,7 @@ logical_axis_rules: [
                       ['out_channels', 'tensor'],
                       ['conv_out', 'fsdp'],
                     ]
-data_sharding: [['data', 'fsdp', 'context', 'tensor']]
+data_sharding: [['data', 'fsdp']]
 
 # One axis for each parallelism type may hold a placeholder (-1)
 # value to auto-shard based on available slices and devices.
@@ -203,7 +216,7 @@ num_train_epochs: 1
 seed: 0
 output_dir: 'output/'
 output_name: "flux2klein_generated_image.png"
-per_device_batch_size: 1
+per_device_batch_size: 1.0
 
 warmup_steps_fraction: 0.1
 learning_rate_schedule_steps: -1 # By default the length of the schedule is set to the number of steps.

src/maxdiffusion/configs/base_flux_dev.yml

@@ -33,7 +33,7 @@ t5xxl_model_name_or_path: 'ariG23498/t5-v1-1-xxl-flax'
 
 # Flux params
 flux_name: "flux-dev"
-max_sequence_length: 512
+max_sequence_length: 256
 time_shift: True
 base_shift: 0.5
 max_shift: 1.15
@@ -62,30 +62,31 @@ jit_initializers: True
 # Set true to load weights from pytorch
 from_pt: True
 split_head_dim: True
-attention: 'flash' # Supported attention: dot_product, flash, cudnn_flash_te
-use_base2_exp: False
-use_experimental_scheduler: False
+attention: 'ulysses_custom_fixed_m' # Supported attention: dot_product, flash, cudnn_flash_te, ulysses_custom_fixed_m
+use_base2_exp: True
+use_experimental_scheduler: True
 # If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens.
 # Else we do not pass in segment ids and on vpu bound hardware like trillium this is faster.
 # However, when padding tokens are significant, this will lead to worse quality and should be set to True.
-mask_padding_tokens: True 
+mask_padding_tokens: False 
 # Maxdiffusion has 2 types of attention sharding strategies:
 # 1. attention_sharding_uniform = True : same sequence sharding rules applied for q in both (self and cross attention)
 # 2. attention_sharding_uniform = False : Heads are sharded uniformly across devices for self attention while sequence is sharded
 #    in cross attention q.
 attention_sharding_uniform: True 
 
-#flash_block_sizes: {}
 # Use the following flash_block_sizes on v6e (Trillium) due to larger vmem.
 flash_block_sizes: {
-  "block_q" : 1536,
-  "block_kv_compute" : 1536,
-  "block_kv" : 1536,
-  "block_q_dkv" : 1536,
-  "block_kv_dkv" : 1536,
-  "block_kv_dkv_compute" : 1536,
-  "block_q_dq" : 1536,
-  "block_kv_dq" : 1536
+  "block_q" : 4864,
+  "block_kv_compute" : 1024,
+  "block_kv" : 1024,
+  "block_kv_compute_in" : 1024,
+  "heads_per_tile" : 3,
+  "block_q_dkv" : 4864,
+  "block_kv_dkv" : 1024,
+  "block_kv_dkv_compute" : 1024,
+  "block_q_dq" : 4864,
+  "block_kv_dq" : 1024
 }
 # GroupNorm groups
 norm_num_groups: 32
@@ -162,7 +163,7 @@ logical_axis_rules: [
                       ['out_channels', 'tensor'],
                       ['conv_out', 'fsdp'],
                     ]
-data_sharding: [['data', 'fsdp', 'context', 'tensor']]
+data_sharding: [['data', 'fsdp']]
 
 # One axis for each parallelism type may hold a placeholder (-1)
 # value to auto-shard based on available slices and devices.
@@ -251,7 +252,7 @@ max_train_steps: 1500
 num_train_epochs: 1
 seed: 0
 output_dir: 'sdxl-model-finetuned'
-per_device_batch_size: 1
+per_device_batch_size: 1.0
 
 warmup_steps_fraction: 0.1
 learning_rate_schedule_steps: -1 # By default the length of the schedule is set to the number of steps.

src/maxdiffusion/configs/base_flux_schnell.yml

@@ -61,38 +61,31 @@ jit_initializers: True
 # Set true to load weights from pytorch
 from_pt: True
 split_head_dim: True
-attention: 'flash' # Supported attention: dot_product, flash, cudnn_flash_te
+attention: 'ulysses_custom_fixed_m' # Supported attention: dot_product, flash, cudnn_flash_te
+use_base2_exp: True
+use_experimental_scheduler: True
 # If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens.
 # Else we do not pass in segment ids and on vpu bound hardware like trillium this is faster.
 # However, when padding tokens are significant, this will lead to worse quality and should be set to True.
-mask_padding_tokens: True 
+mask_padding_tokens: False 
 # Maxdiffusion has 2 types of attention sharding strategies:
 # 1. attention_sharding_uniform = True : same sequence sharding rules applied for q in both (self and cross attention)
 # 2. attention_sharding_uniform = False : Heads are sharded uniformly across devices for self attention while sequence is sharded
 #    in cross attention q.
 attention_sharding_uniform: True 
+# Use the following flash_block_sizes on v6e (Trillium) due to larger vmem.
 flash_block_sizes: {
-  "block_q" : 256,
-  "block_kv_compute" : 256,
-  "block_kv" : 256,
-  "block_q_dkv" : 256,
-  "block_kv_dkv" : 256,
-  "block_kv_dkv_compute" : 256,
-  "block_q_dq" : 256,
-  "block_kv_dq" : 256
+  "block_q" : 4864,
+  "block_kv_compute" : 1024,
+  "block_kv" : 1024,
+  "block_kv_compute_in" : 1024,
+  "heads_per_tile" : 3,
+  "block_q_dkv" : 4864,
+  "block_kv_dkv" : 1024,
+  "block_kv_dkv_compute" : 1024,
+  "block_q_dq" : 4864,
+  "block_kv_dq" : 1024
 }
-
-# Use the following flash_block_sizes on v6e (Trillium).
-# flash_block_sizes: {
-#   "block_q" : 2176,
-#   "block_kv_compute" : 2176,
-#   "block_kv" : 2176,
-#   "block_q_dkv" : 2176,
-#   "block_kv_dkv" : 2176,
-#   "block_kv_dkv_compute" : 2176,
-#   "block_q_dq" : 2176,
-#   "block_kv_dq" : 2176
-# }
 # GroupNorm groups
 norm_num_groups: 32
 
@@ -166,7 +159,7 @@ logical_axis_rules: [
                       ['out_channels', 'tensor'],
                       ['conv_out', 'fsdp'],
                     ]
-data_sharding: [['data', 'fsdp', 'context', 'tensor']]
+data_sharding: [['data', 'fsdp']]
 
 # One axis for each parallelism type may hold a placeholder (-1)
 # value to auto-shard based on available slices and devices.
@@ -227,7 +220,7 @@ max_train_steps: 200
 num_train_epochs: 1
 seed: 0
 output_dir: 'sdxl-model-finetuned'
-per_device_batch_size: 1
+per_device_batch_size: 1.0
 
 warmup_steps_fraction: 0.0
 learning_rate_schedule_steps: -1 # By default the length of the schedule is set to the number of steps.

src/maxdiffusion/generate_flux.py

@@ -32,7 +32,7 @@
 from transformers import (CLIPTokenizer, FlaxCLIPTextModel, T5EncoderModel, FlaxT5EncoderModel, AutoTokenizer)
 
 from maxdiffusion import FlaxAutoencoderKL, pyconfig, max_logging, max_utils
-from maxdiffusion.models.flux.transformers.transformer_flux_flax import FluxTransformer2DModel
+from maxdiffusion.models.flux.transformers.transformer_flux import FluxTransformer2DModel
 from maxdiffusion.train_utils import transformer_engine_context
 from maxdiffusion.max_utils import (
     device_put_replicated,
@@ -78,7 +78,7 @@ def unpack(x: Array, height: int, width: int) -> Array:
 
 
 def vae_decode(latents, vae, state, config):
-  img = unpack(x=latents.astype(jnp.float32), height=config.resolution, width=config.resolution)
+  img = unpack(x=latents.astype(jnp.bfloat16), height=config.resolution, width=config.resolution)
   img = img / vae.config.scaling_factor + vae.config.shift_factor
   img = vae.apply({"params": state.params}, img, deterministic=True, method=vae.decode).sample
   return img
@@ -281,7 +281,7 @@ def run(config):
   devices_array = create_device_mesh(config)
   mesh = Mesh(devices_array, config.mesh_axes)
 
-  global_batch_size = config.per_device_batch_size * jax.local_device_count()
+  global_batch_size = int(round(config.per_device_batch_size * jax.local_device_count()))
 
   # LOAD VAE
   with mesh: