feat(flux2klein): optimize model loading and XLA compilation latency by 2 minutes

Author: amepas

src/maxdiffusion/generate_flux2klein.py

@@ -341,6 +341,7 @@ def qwen3_init_fn():
       def unbox_fn(x):
         return x.unbox() if isinstance(x, flax_spmd.LogicallyPartitioned) else x
 
+      t_sub0 = time.time()
       params = jax.tree_util.tree_map(
           unbox_fn, abstract_transformer_vars["params"], is_leaf=lambda k: isinstance(k, flax_spmd.LogicallyPartitioned)
       )
@@ -356,17 +357,23 @@ def unbox_fn(x):
       )
       qwen3_params = flax.core.unfreeze(qwen3_params)
 
+      max_logging.log(f" -> [SUB-TIMING 1/3] PyTree unboxing template setup: {time.time() - t_sub0:.2f}s")
+      t_sub1 = time.time()
+
       params = load_and_convert_flux_klein_weights(safetensors_path, params, num_double_layers, depth)
       vae_params, vae_bn_mean, vae_bn_std = load_and_convert_vae_weights(vae_safetensors_path, vae_params)
       qwen3_params = load_and_convert_qwen3_weights(text_encoder_path, qwen3_params, qwen3_config)
+      max_logging.log(f" -> [SUB-TIMING 2/3] Safetensors loading & key mapping: {time.time() - t_sub1:.2f}s")
 
+      t_sub2 = time.time()
       if config.weights_dtype == "bfloat16":
         max_logging.log("Casting JAX parameters to bfloat16 in-place...")
         cast_dict_to_bfloat16_inplace(params, exclude_keywords=("norm",))
         cast_dict_to_bfloat16_inplace(vae_params, exclude_keywords=("norm",))
         cast_dict_to_bfloat16_inplace(qwen3_params, exclude_keywords=("norm",))
         vae_bn_mean = vae_bn_mean.astype(jnp.bfloat16)
         vae_bn_std = vae_bn_std.astype(jnp.bfloat16)
+        max_logging.log(f" -> [SUB-TIMING 2b/3] bfloat16 in-place casting: {time.time() - t_sub2:.2f}s")
 
       params = flax.core.freeze(params)
       vae_params = flax.core.freeze(vae_params)
@@ -375,6 +382,7 @@ def unbox_fn(x):
       max_logging.log("\n" + "=" * 80)
       max_logging.log("🚀 Pinning all parameters to TPU HBM permanently...")
       max_logging.log("=" * 80 + "\n")
+      t_sub3 = time.time()
       max_logging.log("Putting params on TPU HBM...")
       with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
         try:
@@ -393,6 +401,7 @@ def unbox_fn(x):
         vae_params = jax.tree_util.tree_map(max_utils.device_put_replicated, vae_params, vae_shardings)
         max_logging.log("Putting qwen3_params on TPU HBM...")
         qwen3_params = jax.tree_util.tree_map(max_utils.device_put_replicated, qwen3_params, qwen3_shardings)
+      max_logging.log(f" -> [SUB-TIMING 3/3] TPU HBM device_put placement: {time.time() - t_sub3:.2f}s")
       max_logging.log("All parameters placed on TPU HBM successfully!")
       gc.collect()
       jax.effects_barrier()
@@ -482,6 +491,7 @@ def unbox_fn(x):
     max_logging.log("\n" + "=" * 80)
     max_logging.log("🚀 Running initial dry run (Warmup Pass) to compile XLA graphs...")
     max_logging.log("=" * 80)
+    pipeline.compile_aot_async(params, vae_params, qwen3_params, batch_size=config.batch_size, height=config.height, width=config.width)
     _, warmup_trace = pipeline(
         prompt=active_prompts,
         params=params,
@@ -498,6 +508,7 @@ def unbox_fn(x):
         batch_size=config.batch_size,
         use_latents=use_latents_flag,
         latents=latents_to_use,
+        warmup=True,
         output_dir=config.output_dir,
         output_name="flux2klein_warmup.png",
     )

src/maxdiffusion/models/flux/util.py

@@ -398,11 +398,12 @@ def cast_dict_to_bfloat16_inplace(d, device=None, exclude_keywords=None, parent_
       is_excluded = exclude_keywords and any(kw.lower() in current_key.lower() for kw in exclude_keywords)
       target_dtype = jnp.float32 if is_excluded else jnp.bfloat16
 
-      d[k] = v.astype(target_dtype)
-      if hasattr(d[k], "block_until_ready"):
-        d[k].block_until_ready()
-      del v
-      gc.collect()
+      if v.dtype != target_dtype:
+        d[k] = v.astype(target_dtype)
+        if hasattr(d[k], "block_until_ready"):
+          d[k].block_until_ready()
+        del v
+        gc.collect()
 
 
 # -----------------------------------------------------------------------------
@@ -571,12 +572,16 @@ def load_and_convert_vae_weights(safetensors_path, jax_params):
   max_logging.log(f"Loading VAE weights from: {safetensors_path}")
   pt_state_dict = load_file(safetensors_path)
 
-  def get_pytorch_weight_tensor(key):
-    return pt_state_dict[key]
-
   # Unfreeze JAX params so we can load the weights
   jax_params = flax.core.unfreeze(jax_params)
 
+  first_leaf = jax.tree_util.tree_leaves(jax_params)[0]
+  target_dtype = first_leaf.dtype
+
+  def get_pytorch_weight_tensor(key, dtype=target_dtype):
+    tensor = pt_state_dict[key]
+    return jnp.array(tensor, dtype=dtype)
+
   # Map weights
   max_logging.log("Mapping VAE decoder weights to JAX parameters...")
 

src/maxdiffusion/models/qwen3_flax.py

@@ -648,15 +648,18 @@ def load_and_convert_qwen3_weights(safetensors_path: str, jax_params: dict, conf
     torch_weights = load_file(safetensors_path)
   max_logging.log("Safetensors weights loaded successfully. Starting JAX parameter mapping...")
 
-  # Helper to transpose and cast weight
-  def get_w(name: str, transpose: bool = True) -> np.ndarray:
+  first_leaf = jax.tree_util.tree_leaves(jax_params)[0]
+  target_dtype = first_leaf.dtype
+
+  # Helper to transpose and cast weight directly to target_dtype
+  def get_w(name: str, transpose: bool = True):
     nonlocal torch_weights
     if name not in torch_weights:
       raise KeyError(f"Weight '{name}' not found in safetensors!")
     t = torch_weights[name]
     if len(t.shape) == 2 and transpose:
       t = t.T
-    return t
+    return jnp.array(t, dtype=target_dtype)
 
   # Create mutable copy of JAX params to populate
   import flax

src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py

@@ -97,14 +97,65 @@ def transformer_step(t_params, latents, img_ids, prompt_embeds, txt_ids, vec, ti
           guidance=guidance,
       )
 
-    @jax.jit
+    @jax.jit(donate_argnums=(1,))
     def vae_decode(v_params, latents_unpatched):
       return self.vae.apply({"params": v_params}, latents=latents_unpatched, method=self.vae.decode)
 
     self._jitted_qwen3_forward = qwen3_forward
     self._jitted_transformer_step = transformer_step
     self._jitted_vae_decode = vae_decode
 
+  def compile_aot_async(self, params, vae_params, qwen3_params, batch_size=1, height=1024, width=1024):
+    """Triggers AOT compilation for Qwen3, Flux Transformer, and VAE concurrently using ThreadPoolExecutor."""
+    self._setup_jit_functions()
+    max_logging.log("🚀 Pre-compiling XLA graphs for Qwen3, Flux Transformer, and VAE concurrently...")
+    from concurrent.futures import ThreadPoolExecutor
+
+    seq_len_img = (height // 16) * (width // 16)
+    seq_len_txt = self._config.max_sequence_length
+
+    dummy_ids = jnp.zeros((batch_size, seq_len_txt), dtype=jnp.int32)
+    dummy_mask = jnp.ones((batch_size, seq_len_txt), dtype=jnp.int32)
+
+    dummy_latents = jnp.zeros((batch_size, seq_len_img, 128), dtype=jnp.float32)
+    dummy_img_ids = jnp.zeros((batch_size, seq_len_img, 4), dtype=jnp.int32)
+    dummy_prompt_embeds = jnp.zeros((batch_size, seq_len_txt, 12288), dtype=jnp.bfloat16)
+    dummy_txt_ids = jnp.zeros((batch_size, seq_len_txt, 4), dtype=jnp.float32)
+    dummy_t_vec = jnp.zeros((batch_size,), dtype=jnp.float32)
+
+    dummy_unpacked_latents = jnp.zeros((batch_size, 32, height // 8, width // 8), dtype=jnp.float32)
+
+    def compile_qwen3():
+      t0 = time.perf_counter()
+      with self.mesh, nn_partitioning.axis_rules(self._config.logical_axis_rules):
+        self._jitted_qwen3_forward.lower(qwen3_params, dummy_ids, dummy_mask).compile()
+      max_logging.log(f" -> [AOT COMPILED] Qwen3 Text Encoder in {time.perf_counter() - t0:.2f}s")
+
+    def compile_transformer():
+      t0 = time.perf_counter()
+      with self.mesh, nn_partitioning.axis_rules(self._config.logical_axis_rules):
+        self._jitted_transformer_step.lower(
+            params, dummy_latents, dummy_img_ids, prompt_embeds=dummy_prompt_embeds, txt_ids=dummy_txt_ids, vec=None, timestep=dummy_t_vec, guidance=None
+        ).compile()
+      max_logging.log(f" -> [AOT COMPILED] Flux Transformer Step in {time.perf_counter() - t0:.2f}s")
+
+    def compile_vae():
+      t0 = time.perf_counter()
+      with self.mesh, nn_partitioning.axis_rules(self._config.logical_axis_rules):
+        self._jitted_vae_decode.lower(vae_params, dummy_unpacked_latents).compile()
+      max_logging.log(f" -> [AOT COMPILED] VAE Decoder in {time.perf_counter() - t0:.2f}s")
+
+    t_start = time.perf_counter()
+    with ThreadPoolExecutor(max_workers=3) as executor:
+      futures = [
+          executor.submit(compile_qwen3),
+          executor.submit(compile_transformer),
+          executor.submit(compile_vae),
+      ]
+      for future in futures:
+        future.result()
+    max_logging.log(f"⚡ [AOT CONCURRENT COMPILATION COMPLETE] Total AOT compile time: {time.perf_counter() - t_start:.2f}s")
+
   def _prepare_latents(self, config, batch_size, height, width):
     num_channels_latents = 32
     latent_height = height // 8
@@ -147,6 +198,7 @@ def __call__(
       use_latents: bool = False,
       latents: Optional[Any] = None,
       measure_time: bool = False,
+      warmup: bool = False,
       output_dir: str = "output/",
       output_name: str = "flux2klein_generated_image.png",
   ):
@@ -293,8 +345,9 @@ def put_data_on_devices(x, sharding):
       # ---------------------------------------------------------------------
       # PHASE B: Denoising Loop (Flux Transformer - Standalone Step JIT)
       # ---------------------------------------------------------------------
+      steps_to_run = 1 if warmup else num_inference_steps
       print(
-          f"{host_prefix} [PHASE B] Running {num_inference_steps}-step E2E Denoising Loop on a batch of {batch_size} images...",
+          f"{host_prefix} [PHASE B] Running {steps_to_run}-step E2E Denoising Loop on a batch of {batch_size} images (warmup={warmup})...",
           flush=True,
       )
       t0 = time.perf_counter()
@@ -303,7 +356,7 @@ def put_data_on_devices(x, sharding):
         guidance_vec_val = None
         vec_val = None
 
-        for step_idx in range(num_inference_steps):
+        for step_idx in range(steps_to_run):
           timestep = scheduler_state.timesteps[step_idx]
           t_vec = jnp.full((batch_size,), timestep / 1000.0, dtype=latents_jax.dtype)