Refactor Phase B to use fused jax.lax.fori_loop for optimal multi-host execution

Author: amepas

src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py

@@ -256,43 +256,49 @@ def put_data_on_devices(x, sharding):
       prompt_embeds_jax = put_data_on_devices(prompt_embeds_jax, data_sharding)
 
       # ---------------------------------------------------------------------
-      # PHASE B: Denoising Loop (Flux Transformer)
+      # PHASE B: Denoising Loop (Flux Transformer - Fused JAX fori_loop)
       # ---------------------------------------------------------------------
       max_logging.log(
           f"[PHASE B] Running {num_inference_steps}-step E2E Denoising Loop on a batch of {batch_size} images..."
       )
       t0 = time.perf_counter()
 
-      guidance_vec_val = None
-      vec_val = None
-
-      for step_idx in range(num_inference_steps):
-        timestep = scheduler_state.timesteps[step_idx]
-        t_vec = jnp.array([timestep / 1000.0] * batch_size)
-
-        # Execute transformer forward pass step
-        model_output = self._jitted_transformer_step(
-            params, latents_jax, img_ids_val, prompt_embeds_jax, txt_ids_val, vec_val, t_vec, guidance_vec_val
-        )
-
-        # Update latents using FlowMatch step
-        latents_jax = self.scheduler.step(
-            state=scheduler_state,
-            model_output=model_output.sample,
-            timestep=scheduler_state.timesteps[step_idx],
-            sample=latents_jax,
-        ).prev_sample
-
-        # Print progress
-        sigma_val = scheduler_state.sigmas[step_idx]
-        max_logging.log(
-            f" -> Step {step_idx}: Timestep = {scheduler_state.timesteps[step_idx]:.4f}, Sigma = {sigma_val:.4f}"
-        )
-
+      timesteps_array = scheduler_state.timesteps
+      sigmas_array = scheduler_state.sigmas
+
+      @jax.jit
+      def denoise_loop(t_params, latents, prompt_embeds, txt_ids, img_ids, timesteps, sigmas):
+        def step_fn(step_idx, current_latents):
+          t_val = timesteps[step_idx] / 1000.0
+          t_vec = jnp.full((batch_size,), t_val, dtype=current_latents.dtype)
+          out = self.transformer.apply(
+              {"params": t_params},
+              hidden_states=current_latents,
+              img_ids=img_ids,
+              encoder_hidden_states=prompt_embeds,
+              txt_ids=txt_ids,
+              pooled_projections=None,
+              timestep=t_vec,
+              guidance=None,
+          ).sample
+          dt = sigmas[step_idx + 1] - sigmas[step_idx]
+          return current_latents + dt * out
+
+        return jax.lax.fori_loop(0, num_inference_steps, step_fn, latents)
+
+      latents_jax = denoise_loop(
+          params,
+          latents_jax,
+          prompt_embeds_jax,
+          txt_ids_val,
+          img_ids_val,
+          timesteps_array,
+          sigmas_array,
+      )
       latents_jax.block_until_ready()
 
-    trace["denoise_loop"] = time.perf_counter() - t0
-    max_logging.log(f" -> [TIMING] Denoising Loop (Flux): {trace['denoise_loop']:.4f} seconds ⏱️")
+      trace["denoise_loop"] = time.perf_counter() - t0
+      max_logging.log(f" -> [TIMING] Denoising Loop (Flux): {trace['denoise_loop']:.4f} seconds ⏱️")
 
     # ---------------------------------------------------------------------
     # PHASE C: Decode Latents (VAE Decoder)