diff --git a/src/maxdiffusion/__init__.py b/src/maxdiffusion/__init__.py
index e9addadcc..deb9d7e04 100644
--- a/src/maxdiffusion/__init__.py
+++ b/src/maxdiffusion/__init__.py
@@ -16,6 +16,15 @@
 
 __version__ = "0.22.0.dev0"
 
+# Monkeypatch jax.numpy.clip to support legacy a_min and a_max arguments (fixed for newer JAX 0.4.31+)
+import jax.numpy as jnp
+orig_clip = jnp.clip
+def patched_clip(a, a_min=None, a_max=None, *args, **kwargs):
+  c_min = kwargs.pop("min", a_min)
+  c_max = kwargs.pop("max", a_max)
+  return orig_clip(a, min=c_min, max=c_max, *args, **kwargs)
+jnp.clip = patched_clip
+
 from typing import TYPE_CHECKING
 
 from .utils import (
@@ -367,7 +376,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 +453,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
diff --git a/src/maxdiffusion/checkpointing/flux_checkpointer.py b/src/maxdiffusion/checkpointing/flux_checkpointer.py
index 70b54d08d..c7f427de6 100644
--- a/src/maxdiffusion/checkpointing/flux_checkpointer.py
+++ b/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)
diff --git a/src/maxdiffusion/configs/base_flux2klein.yml b/src/maxdiffusion/configs/base_flux2klein.yml
index f2813c8fd..7162164e4 100644
--- a/src/maxdiffusion/configs/base_flux2klein.yml
+++ b/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' # Supported attention: dot_product, flash, cudnn_flash_te, ulysses_custom
+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.
diff --git a/src/maxdiffusion/configs/base_flux2klein_9B.yml b/src/maxdiffusion/configs/base_flux2klein_9B.yml
index a6c670a69..0515f61d8 100644
--- a/src/maxdiffusion/configs/base_flux2klein_9B.yml
+++ b/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' # Supported attention: dot_product, flash, cudnn_flash_te, ulysses_custom
+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.
diff --git a/src/maxdiffusion/configs/base_flux_dev.yml b/src/maxdiffusion/configs/base_flux_dev.yml
index 20f1fcc8f..dbf62969f 100644
--- a/src/maxdiffusion/configs/base_flux_dev.yml
+++ b/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' # Supported attention: dot_product, flash, cudnn_flash_te, ulysses_custom
+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.
diff --git a/src/maxdiffusion/configs/base_flux_schnell.yml b/src/maxdiffusion/configs/base_flux_schnell.yml
index 5d3f7f2fa..2426434cd 100644
--- a/src/maxdiffusion/configs/base_flux_schnell.yml
+++ b/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' # Supported attention: dot_product, flash, cudnn_flash_te, ulysses_custom
+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.
diff --git a/src/maxdiffusion/generate_flux.py b/src/maxdiffusion/generate_flux.py
index 8d2153c65..820c2bc97 100644
--- a/src/maxdiffusion/generate_flux.py
+++ b/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:
@@ -482,8 +482,9 @@ def validate_inputs(latents, latent_image_ids, prompt_embeds, text_ids, timestep
     max_logging.log(f"Inference time: {t1 - t0:.1f}s.")
     imgs = np.array(imgs)
     imgs = (imgs * 0.5 + 0.5).clip(0, 1)
+    imgs = np.nan_to_num(imgs, nan=0.0)
     imgs = np.transpose(imgs, (0, 2, 3, 1))
-    imgs = np.uint8(imgs * 255)
+    imgs = np.uint8(np.clip(imgs * 255, 0, 255))
     for i, image in enumerate(imgs):
       Image.fromarray(image).save(f"flux_{i}.png")
 
diff --git a/src/maxdiffusion/generate_flux2klein.py b/src/maxdiffusion/generate_flux2klein.py
index 7956c850d..647c17104 100644
--- a/src/maxdiffusion/generate_flux2klein.py
+++ b/src/maxdiffusion/generate_flux2klein.py
@@ -35,7 +35,7 @@
 from maxdiffusion.max_utils import create_device_mesh
 from maxdiffusion.train_utils import transformer_engine_context
 
-from maxdiffusion.models.flux.transformers.transformer_flux_flax import Flux2KleinTransformer2DModel
+from maxdiffusion.models.flux.transformers.transformer_flux import Flux2KleinTransformer2DModel
 from maxdiffusion.models.vae_flax import FlaxAutoencoderKL
 from maxdiffusion.models.qwen3_flax import FlaxQwen3Config, FlaxQwen3Model
 from maxdiffusion.models.qwen3_utils import load_and_convert_qwen3_weights
@@ -141,8 +141,24 @@ def main(argv):
   config = pyconfig.config
   os.makedirs(config.output_dir, exist_ok=True)
 
+  # Derive global batch size from per_device_batch_size
+  if getattr(config, "per_device_batch_size", None) is not None:
+    calculated_batch_size = max(1, int(round(config.per_device_batch_size * jax.device_count())))
+    has_explicit_batch_size = any(arg.startswith("batch_size=") for arg in custom_overrides)
+    if not has_explicit_batch_size:
+      config.batch_size = calculated_batch_size
+      pyconfig._config.keys["batch_size"] = calculated_batch_size
+      max_logging.log(
+          f"ℹ️ Set global batch_size={config.batch_size} from per_device_batch_size={config.per_device_batch_size} across {jax.device_count()} TPU devices."
+      )
+
   # 2. Setup device mesh
-  if config.batch_size == 1 and config.ici_tensor_parallelism == 1 and jax.device_count() > 1:
+  if (
+      config.batch_size == 1
+      and config.ici_tensor_parallelism == 1
+      and config.ici_context_parallelism == 1
+      and jax.device_count() > 1
+  ):
     max_logging.log(
         f"ℹ️ Auto-configuring Tensor Parallelism: ici_tensor_parallelism={jax.device_count()}, ici_fsdp_parallelism=1 for batch_size=1 on {jax.device_count()} TPU devices."
     )
@@ -255,6 +271,9 @@ def main(argv):
       dtype=jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32,
       weights_dtype=jnp.bfloat16 if config.weights_dtype == "bfloat16" else jnp.float32,
       attention_kernel=config.attention,
+      flash_block_sizes=getattr(config, "flash_block_sizes", None),
+      use_base2_exp=getattr(config, "use_base2_exp", False),
+      use_experimental_scheduler=getattr(config, "use_experimental_scheduler", False),
       scale_shift_order=getattr(config, "scale_shift_order", "shift_scale"),
   )
 
@@ -511,25 +530,35 @@ def unbox_fn(x):
     max_logging.log("\n" + "=" * 80)
     max_logging.log("⏱️ Running timed pass at full TPU speed...")
     max_logging.log("=" * 80)
-    _, main_trace = pipeline(
-        prompt=active_prompts,
-        params=params,
-        vae_params=vae_params,
-        qwen3_params=qwen3_params,
-        vae_bn_mean=vae_bn_mean,
-        vae_bn_std=vae_bn_std,
-        transformer_shardings=transformer_shardings,
-        vae_shardings=vae_shardings,
-        qwen3_shardings=qwen3_shardings,
-        height=config.height,
-        width=config.width,
-        num_inference_steps=config.num_inference_steps,
-        batch_size=config.batch_size,
-        use_latents=use_latents_flag,
-        latents=latents_to_use,
-        output_dir=config.output_dir,
-        output_name=config.output_name,
+    import contextlib
+
+    profile_path = getattr(config, "profiler_dir", "/tmp/profile_flux2klein")
+    profiler_context = (
+        jax.profiler.trace(profile_path) if getattr(config, "enable_profiler", False) else contextlib.nullcontext()
     )
+    if getattr(config, "enable_profiler", False):
+      max_logging.log(f"Profiling enabled! Recording trace to {profile_path}...")
+
+    with profiler_context:
+      _, main_trace = pipeline(
+          prompt=active_prompts,
+          params=params,
+          vae_params=vae_params,
+          qwen3_params=qwen3_params,
+          vae_bn_mean=vae_bn_mean,
+          vae_bn_std=vae_bn_std,
+          transformer_shardings=transformer_shardings,
+          vae_shardings=vae_shardings,
+          qwen3_shardings=qwen3_shardings,
+          height=config.height,
+          width=config.width,
+          num_inference_steps=config.num_inference_steps,
+          batch_size=config.batch_size,
+          use_latents=use_latents_flag,
+          latents=latents_to_use,
+          output_dir=config.output_dir,
+          output_name=config.output_name,
+      )
     main_time = (
         main_trace.get("prompt_encoding", 0.0) + main_trace.get("denoise_loop", 0.0) + main_trace.get("vae_decode", 0.0)
     )
diff --git a/src/maxdiffusion/generate_flux_multi_res.py b/src/maxdiffusion/generate_flux_multi_res.py
index c74533d42..938f3503e 100644
--- a/src/maxdiffusion/generate_flux_multi_res.py
+++ b/src/maxdiffusion/generate_flux_multi_res.py
@@ -31,7 +31,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.max_utils import (
     device_put_replicated,
     get_memory_allocations,
@@ -327,7 +327,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
 
diff --git a/src/maxdiffusion/models/__init__.py b/src/maxdiffusion/models/__init__.py
index d82305aa3..9ee64270f 100644
--- a/src/maxdiffusion/models/__init__.py
+++ b/src/maxdiffusion/models/__init__.py
@@ -30,7 +30,7 @@
   from .vae_flax import FlaxAutoencoderKL
   from .z_image.transformer_z_image import ZImageTransformer2DModel
   from .lora import *
-  from .flux.transformers.transformer_flux_flax import FluxTransformer2DModel
+  from .flux.transformers.transformer_flux import FluxTransformer2DModel
   from .ltx_video.transformers.transformer3d import Transformer3DModel
 
 else:
diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py
index 8b84ae057..efaac7def 100644
--- a/src/maxdiffusion/models/attention_flax.py
+++ b/src/maxdiffusion/models/attention_flax.py
@@ -915,7 +915,7 @@ def wrap_ulysses_attention(query, key, value, attention_mask):
           block_sizes=bsizes,
           orig_q_seq_len=query_seq_len,
           orig_kv_seq_len=key_seq_len,
-          heads_per_tile=heads_per_tile,
+          heads_per_tile=1 if use_fixed_m else heads_per_tile,
           use_base2_exp=use_base2_exp,
           use_experimental_scheduler=use_experimental_scheduler,
           vmem_limit_bytes=vmem_limit_bytes,
@@ -1424,7 +1424,7 @@ def wrap_ulysses_ring_attention(query, key, value):
           block_sizes=bsizes,
           orig_q_seq_len=query_seq_len,
           orig_kv_seq_len=key_seq_len,
-          heads_per_tile=heads_per_tile,
+          heads_per_tile=1 if use_fixed_m else heads_per_tile,
           use_base2_exp=use_base2_exp,
           use_experimental_scheduler=use_experimental_scheduler,
           vmem_limit_bytes=vmem_limit_bytes,
@@ -2992,14 +2992,24 @@ def __call__(
       # key_proj = nn.with_logical_constraint(key_proj, self.key_axis_names)
       # value_proj = nn.with_logical_constraint(value_proj, self.value_axis_names)
 
-    if not isinstance(image_rotary_emb, (tuple, list)):
-      image_rotary_emb = rearrange(image_rotary_emb, "n d (i j) -> n d i j", i=2, j=2)
-
-    query_proj = query_proj.swapaxes(1, 2)
-    key_proj = key_proj.swapaxes(1, 2)
-    query_proj, key_proj = apply_rope(query_proj, key_proj, image_rotary_emb)
-    query_proj = query_proj.swapaxes(1, 2)
-    key_proj = key_proj.swapaxes(1, 2)
+    if image_rotary_emb is not None:
+      if not isinstance(image_rotary_emb, (tuple, list)):
+        image_rotary_emb_reordered = rearrange(image_rotary_emb, "n d (i j) -> 1 n 1 d i j", i=2, j=2)
+        B_q, L_q, H_q, D_q = query_proj.shape
+        q_ = query_proj.reshape(B_q, L_q, H_q, D_q // 2, 1, 2)
+        k_ = key_proj.reshape(B_q, L_q, H_q, D_q // 2, 1, 2)
+        query_proj = (
+            (image_rotary_emb_reordered[..., 0] * q_[..., 0] + image_rotary_emb_reordered[..., 1] * q_[..., 1])
+            .reshape(B_q, L_q, H_q, D_q)
+            .astype(query_proj.dtype)
+        )
+        key_proj = (
+            (image_rotary_emb_reordered[..., 0] * k_[..., 0] + image_rotary_emb_reordered[..., 1] * k_[..., 1])
+            .reshape(B_q, L_q, H_q, D_q)
+            .astype(key_proj.dtype)
+        )
+      else:
+        query_proj, key_proj = apply_rope(query_proj, key_proj, image_rotary_emb)
 
     query_proj = query_proj.reshape(B, -1, H * D)
     key_proj = key_proj.reshape(B, -1, H * D)
diff --git a/src/maxdiffusion/models/flux/__init__.py b/src/maxdiffusion/models/flux/__init__.py
index 217c0ac82..e24a6e690 100644
--- a/src/maxdiffusion/models/flux/__init__.py
+++ b/src/maxdiffusion/models/flux/__init__.py
@@ -14,4 +14,4 @@
 limitations under the License.
 """
 
-from .transformers.transformer_flux_flax import FluxTransformer2DModel
+from .transformers.transformer_flux import FluxTransformer2DModel
diff --git a/src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py b/src/maxdiffusion/models/flux/transformers/transformer_flux.py
similarity index 96%
rename from src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py
rename to src/maxdiffusion/models/flux/transformers/transformer_flux.py
index af8e3763a..02a75c673 100644
--- a/src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py
+++ b/src/maxdiffusion/models/flux/transformers/transformer_flux.py
@@ -187,20 +187,31 @@ def __call__(self, hidden_states, temb, image_rotary_emb=None):
 
     qkv_proj = qkv.reshape(B, L, K, H, D)
     q, k, v = jnp.split(qkv_proj, 3, axis=2)
-    q = q.squeeze(2).swapaxes(1, 2)
-    k = k.squeeze(2).swapaxes(1, 2)
-    v = v.squeeze(2).swapaxes(1, 2)
+    q = q.squeeze(2)
+    k = k.squeeze(2)
+    v = v.squeeze(2)
 
     q = self.attn.query_norm(q)
     k = self.attn.key_norm(k)
 
     if image_rotary_emb is not None:
-      image_rotary_emb_reordered = rearrange(image_rotary_emb, "n d (i j) -> n d i j", i=2, j=2)
-      q, k = apply_rope(q, k, image_rotary_emb_reordered)
+      image_rotary_emb_reordered = rearrange(image_rotary_emb, "n d (i j) -> 1 n 1 d i j", i=2, j=2)
+      q_ = q.reshape(B, L, H, D // 2, 1, 2)
+      k_ = k.reshape(B, L, H, D // 2, 1, 2)
+      q = (
+          (image_rotary_emb_reordered[..., 0] * q_[..., 0] + image_rotary_emb_reordered[..., 1] * q_[..., 1])
+          .reshape(B, L, H, D)
+          .astype(q.dtype)
+      )
+      k = (
+          (image_rotary_emb_reordered[..., 0] * k_[..., 0] + image_rotary_emb_reordered[..., 1] * k_[..., 1])
+          .reshape(B, L, H, D)
+          .astype(k.dtype)
+      )
 
-    q = q.transpose(0, 2, 1, 3).reshape(q.shape[0], q.shape[2], -1)
-    k = k.transpose(0, 2, 1, 3).reshape(k.shape[0], k.shape[2], -1)
-    v = v.transpose(0, 2, 1, 3).reshape(v.shape[0], v.shape[2], -1)
+    q = q.reshape(B, L, H * D)
+    k = k.reshape(B, L, H * D)
+    v = v.reshape(B, L, H * D)
 
     attn_output = self.attn.attention_op.apply_attention(q, k, v)
     attn_output = checkpoint_name(attn_output, "attn_output")
@@ -605,13 +616,12 @@ def __call__(
       train: bool = False,
   ):
     hidden_states = self.img_in(hidden_states)
-    timestep = self.timestep_embedding(timestep, 256)
-    timestep = nn.with_logical_constraint(timestep, ("activation_batch", None))
+    timestep = nn.with_logical_constraint(timestep, ("activation_batch",))
 
-    if self.guidance_embeds:
-      guidance = self.timestep_embedding(guidance, 256)
-    else:
+    if not self.guidance_embeds:
       guidance = None
+    else:
+      guidance = nn.with_logical_constraint(guidance, ("activation_batch",))
     temb = (
         self.time_text_embed(timestep, pooled_projections)
         if guidance is None
@@ -767,6 +777,8 @@ class Flux2KleinSingleTransformerBlock(nn.Module):
   precision: float = None
   use_global_modulation: bool = False
   use_swiglu: bool = True
+  use_base2_exp: bool = False
+  use_experimental_scheduler: bool = False
 
   def setup(self):
     mlp_hidden_dim = int(self.dim * self.mlp_ratio)
@@ -814,18 +826,15 @@ def setup(self):
         attention_kernel=self.attention_kernel,
         mesh=self.mesh,
         flash_block_sizes=self.flash_block_sizes,
+        use_base2_exp=self.use_base2_exp,
+        use_experimental_scheduler=self.use_experimental_scheduler,
     )
 
   def __call__(self, hidden_states, temb=None, image_rotary_emb=None, temb_mod=None):
     residual = hidden_states
     if self.use_global_modulation:
       shift_msa, scale_msa, gate = jnp.split(temb_mod, 3, axis=-1)
-      shift_msa = jnp.expand_dims(shift_msa, axis=1)
-      scale_msa = jnp.expand_dims(scale_msa, axis=1)
-      gate = jnp.expand_dims(gate, axis=1)
-
-      norm_hidden_states = self.norm(hidden_states)
-      norm_hidden_states = (1 + scale_msa) * norm_hidden_states + shift_msa
+      norm_hidden_states = self.norm(hidden_states) * (1.0 + scale_msa) + shift_msa
     else:
       norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
 
@@ -835,22 +844,21 @@ def __call__(self, hidden_states, temb=None, image_rotary_emb=None, temb_mod=Non
 
     B, L = hidden_states.shape[:2]
     H, D, K = self.num_attention_heads, qkv.shape[-1] // (self.num_attention_heads * 3), 3
-    qkv_proj = qkv.reshape(B, L, K, H, D).transpose(2, 0, 3, 1, 4)
-    q, k, v = qkv_proj
+    qkv_proj = qkv.reshape(B, L, K, H, D)
+    q, k, v = jnp.split(qkv_proj, 3, axis=2)
+    q = q.squeeze(2)
+    k = k.squeeze(2)
+    v = v.squeeze(2)
 
     q = self.attn.query_norm(q)
     k = self.attn.key_norm(k)
 
     if image_rotary_emb is not None:
-      if isinstance(image_rotary_emb, (tuple, list)):
-        image_rotary_emb_reordered = image_rotary_emb
-      else:
-        image_rotary_emb_reordered = rearrange(image_rotary_emb, "n d (i j) -> n d i j", i=2, j=2)
-      q, k = apply_rope(q, k, image_rotary_emb_reordered)
+      q, k = apply_rope(q, k, image_rotary_emb)
 
-    q = q.transpose(0, 2, 1, 3).reshape(q.shape[0], q.shape[2], -1)
-    k = k.transpose(0, 2, 1, 3).reshape(k.shape[0], k.shape[2], -1)
-    v = v.transpose(0, 2, 1, 3).reshape(v.shape[0], v.shape[2], -1)
+    q = q.reshape(B, L, H * D)
+    k = k.reshape(B, L, H * D)
+    v = v.reshape(B, L, H * D)
 
     attn_output = self.attn.attention_op.apply_attention(q, k, v)
 
@@ -884,6 +892,8 @@ class Flux2KleinTransformerBlock(nn.Module):
   mlp_ratio: float = 4.0
   qkv_bias: bool = True
   use_global_modulation: bool = False
+  use_base2_exp: bool = False
+  use_experimental_scheduler: bool = False
 
   def setup(self):
     if self.use_global_modulation:
@@ -924,6 +934,8 @@ def setup(self):
         weights_dtype=self.weights_dtype,
         precision=self.precision,
         qkv_bias=self.qkv_bias,
+        use_base2_exp=self.use_base2_exp,
+        use_experimental_scheduler=self.use_experimental_scheduler,
     )
     self.ff = FlaxSwiGluFeedForward(
         self.dim,
@@ -956,20 +968,6 @@ def __call__(
       shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = jnp.split(temb_mod_img, 6, axis=-1)
       c_shift_msa, c_scale_msa, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = jnp.split(temb_mod_txt, 6, axis=-1)
 
-      shift_msa = jnp.expand_dims(shift_msa, axis=1)
-      scale_msa = jnp.expand_dims(scale_msa, axis=1)
-      gate_msa = jnp.expand_dims(gate_msa, axis=1)
-      shift_mlp = jnp.expand_dims(shift_mlp, axis=1)
-      scale_mlp = jnp.expand_dims(scale_mlp, axis=1)
-      gate_mlp = jnp.expand_dims(gate_mlp, axis=1)
-
-      c_shift_msa = jnp.expand_dims(c_shift_msa, axis=1)
-      c_scale_msa = jnp.expand_dims(c_scale_msa, axis=1)
-      c_gate_msa = jnp.expand_dims(c_gate_msa, axis=1)
-      c_shift_mlp = jnp.expand_dims(c_shift_mlp, axis=1)
-      c_scale_mlp = jnp.expand_dims(c_scale_mlp, axis=1)
-      c_gate_mlp = jnp.expand_dims(c_gate_mlp, axis=1)
-
       norm1_hidden_states = self.norm1(hidden_states) * (1.0 + scale_msa) + shift_msa
       norm1_encoder_hidden_states = self.norm1_context(encoder_hidden_states) * (1.0 + c_scale_msa) + c_shift_msa
 
@@ -1039,6 +1037,8 @@ class Flux2KleinTransformer2DModel(nn.Module, FlaxModelMixin, ConfigMixin):
   dtype: jnp.dtype = jnp.float32
   weights_dtype: jnp.dtype = jnp.float32
   precision: float = None
+  use_base2_exp: bool = False
+  use_experimental_scheduler: bool = False
 
   def setup(self):
     self.inner_dim = self.num_attention_heads * self.attention_head_dim
@@ -1050,6 +1050,8 @@ def setup(self):
         dtype=self.dtype,
         weights_dtype=self.weights_dtype,
         precision=self.precision,
+        use_base2_exp=self.use_base2_exp,
+        use_experimental_scheduler=self.use_experimental_scheduler,
     )
 
     if self.use_global_modulation:
@@ -1109,6 +1111,8 @@ def setup(self):
           mlp_ratio=self.mlp_ratio,
           qkv_bias=self.qkv_bias,
           use_global_modulation=self.use_global_modulation,
+          use_base2_exp=self.use_base2_exp,
+          use_experimental_scheduler=self.use_experimental_scheduler,
       )
       double_blocks.append(double_block)
     self.double_blocks = double_blocks
@@ -1128,6 +1132,8 @@ def setup(self):
           precision=self.precision,
           mlp_ratio=self.mlp_ratio,
           use_global_modulation=self.use_global_modulation,
+          use_base2_exp=self.use_base2_exp,
+          use_experimental_scheduler=self.use_experimental_scheduler,
       )
       single_blocks.append(single_block)
     self.single_blocks = single_blocks
@@ -1187,9 +1193,9 @@ def __call__(
 
     if self.use_global_modulation:
       temb_silu = nn.silu(temb)
-      double_stream_mod_img = self.double_stream_modulation_img(temb_silu)
-      double_stream_mod_txt = self.double_stream_modulation_txt(temb_silu)
-      single_stream_mod = self.single_stream_modulation(temb_silu)
+      double_stream_mod_img = jnp.expand_dims(self.double_stream_modulation_img(temb_silu), axis=1)
+      double_stream_mod_txt = jnp.expand_dims(self.double_stream_modulation_txt(temb_silu), axis=1)
+      single_stream_mod = jnp.expand_dims(self.single_stream_modulation(temb_silu), axis=1)
     else:
       double_stream_mod_img, double_stream_mod_txt, single_stream_mod = None, None, None
 
diff --git a/src/maxdiffusion/models/flux/util.py b/src/maxdiffusion/models/flux/util.py
index 952519776..0ef694be9 100644
--- a/src/maxdiffusion/models/flux/util.py
+++ b/src/maxdiffusion/models/flux/util.py
@@ -249,7 +249,7 @@ def load_flow_model(name: str, eval_shapes: dict, device: str, hf_download: bool
           renamed_pt_key = renamed_pt_key.replace("out_layer", "linear_2")
         elif "final_layer" in renamed_pt_key:
           renamed_pt_key = renamed_pt_key.replace("final_layer.linear", "proj_out")
-          renamed_pt_key = renamed_pt_key.replace("final_layer.adaLN_modulation_1", "norm_out.Dense_0")
+          renamed_pt_key = renamed_pt_key.replace("final_layer.adaLN_modulation_1", "norm_out.linear")
 
         pt_tuple_key = tuple(renamed_pt_key.split("."))
         flax_key, flax_tensor = rename_key_and_reshape_tensor(pt_tuple_key, tensor, eval_shapes)
@@ -301,16 +301,20 @@ def unpack_latents(latents, batch_size, num_channels_latents, height, width):
   back to the unpacked spatial grid shape (batch_size, channels, height//8, width//8).
   """
   import numpy as np
+  import jax
+  import jax.numpy as jnp
+
+  xp = jnp if isinstance(latents, jax.Array) else np
 
   h_latent = height // 8
   w_latent = width // 8
 
   # 1. Reshape to split spatial grid and packed channel blocks
-  latents = np.reshape(latents, (batch_size, h_latent // 2, w_latent // 2, num_channels_latents, 2, 2))
+  latents = xp.reshape(latents, (batch_size, h_latent // 2, w_latent // 2, num_channels_latents, 2, 2))
   # 2. Permute dimensions back to unpacked order
-  latents = np.transpose(latents, (0, 3, 1, 4, 2, 5))
+  latents = xp.transpose(latents, (0, 3, 1, 4, 2, 5))
   # 3. Flatten back to 4D unpacked latent shape
-  latents = np.reshape(latents, (batch_size, num_channels_latents, h_latent, w_latent))
+  latents = xp.reshape(latents, (batch_size, num_channels_latents, h_latent, w_latent))
   return latents
 
 
diff --git a/src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py b/src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py
index 634ec8d9e..9b97e819b 100644
--- a/src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py
+++ b/src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py
@@ -30,7 +30,7 @@
 from maxdiffusion import max_logging
 from maxdiffusion.max_utils import device_put_replicated
 from ..pipeline_flax_utils import FlaxDiffusionPipeline
-from ...models.flux.transformers.transformer_flux_flax import Flux2KleinTransformer2DModel
+from ...models.flux.transformers.transformer_flux import Flux2KleinTransformer2DModel
 from ...models.vae_flax import FlaxAutoencoderKL
 from ...models.qwen3_flax import FlaxQwen3Model
 from ...schedulers.scheduling_flow_match_flax import FlaxFlowMatchScheduler, compute_empirical_mu
@@ -345,12 +345,12 @@ def put_data_on_devices(x, sharding):
     # Apply Channel-wise Batch Normalization Scaling in packed sequence format (denormalize)
     vae_bn_mean_seq = vae_bn_mean.reshape(1, 1, 128)
     vae_bn_std_seq = vae_bn_std.reshape(1, 1, 128)
-    latents_bn = latents_jax * vae_bn_std_seq + vae_bn_mean_seq
+    latents_bn = (latents_jax * vae_bn_std_seq + vae_bn_mean_seq).astype(jnp.bfloat16)
 
-    # Unpack packed latents back to spatial grid
+    # Unpack packed latents back to spatial grid on TPU device
     latents_unpacked = unpack_latents(latents_bn, batch_size, 32, height, width)
 
-    # Decode VAE latents to RGB pixels
+    # Decode VAE latents to RGB pixels natively in BF16
     decoded_out = self._jitted_vae_decode(vae_params, latents_unpacked)
     # VAE output is in decoded_out.sample
     images_rgb = decoded_out.sample
diff --git a/src/maxdiffusion/pipelines/flux/flux_pipeline.py b/src/maxdiffusion/pipelines/flux/flux_pipeline.py
index 15b2c4f5e..f05a18606 100644
--- a/src/maxdiffusion/pipelines/flux/flux_pipeline.py
+++ b/src/maxdiffusion/pipelines/flux/flux_pipeline.py
@@ -30,7 +30,7 @@
 from ...models import FlaxAutoencoderKL
 from ...schedulers import (FlaxEulerDiscreteScheduler)
 from ..pipeline_flax_utils import FlaxDiffusionPipeline
-from maxdiffusion.models.flux.transformers.transformer_flux_flax import FluxTransformer2DModel
+from maxdiffusion.models.flux.transformers.transformer_flux import FluxTransformer2DModel
 
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
diff --git a/src/maxdiffusion/pyconfig.py b/src/maxdiffusion/pyconfig.py
index d1121ca3f..a0274227e 100644
--- a/src/maxdiffusion/pyconfig.py
+++ b/src/maxdiffusion/pyconfig.py
@@ -127,9 +127,14 @@ def __init__(self, argv: list[str], **kwargs):
         raw_keys[k] = raw_data_from_cmd_line[k]  # take the raw data, no type conversion
       elif k in raw_data_from_cmd_line:
         try:
-          raw_keys[k] = _yaml_types_to_parser[type(raw_data_from_yaml[k])](
-              raw_data_from_cmd_line[k]
-          )  # take the command line value, but type it like the config value.
+          parser = _yaml_types_to_parser[type(raw_data_from_yaml[k])]
+          try:
+            raw_keys[k] = parser(raw_data_from_cmd_line[k])
+          except ValueError:
+            if parser is int:
+              raw_keys[k] = float(raw_data_from_cmd_line[k])
+            else:
+              raise
         except ValueError as e:
           raise ValueError(f"Couldn't parse value from command line '{raw_data_from_cmd_line[k]}' for key '{k}'") from e
       else:
diff --git a/src/maxdiffusion/tests/nnx_flux2klein_test.py b/src/maxdiffusion/tests/nnx_flux2klein_test.py
index 5ae88d8eb..59c38e836 100644
--- a/src/maxdiffusion/tests/nnx_flux2klein_test.py
+++ b/src/maxdiffusion/tests/nnx_flux2klein_test.py
@@ -23,7 +23,7 @@
 import jax.numpy as jnp
 from flax import nnx
 
-from maxdiffusion.models.flux.transformers.transformer_flux_flax import NNXFluxTransformer2DModel
+from maxdiffusion.models.flux.transformers.transformer_flux import NNXFluxTransformer2DModel
 from maxdiffusion.models.qwen3_flax import FlaxQwen3Config, NNXFlaxQwen3Model
 from maxdiffusion.models.vae_flax import NNXFlaxAutoencoderKL
 from maxdiffusion.models.embeddings_flax import NNXCombinedTimestepGuidanceTextProjEmbeddings