Fix host slowdown regression in SenCache by reverting to pure python control flow

Author: Perseus14

src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py

@@ -374,22 +374,19 @@ def run_inference_2_2(
         continue
 
       # Accumulate deltas since last full compute
-      dx_norm = jnp.sqrt(jnp.mean((latents - ref_latent) ** 2))
-      dt = jnp.abs(t_float - ref_timestep)
+      dx_norm = float(jnp.sqrt(jnp.mean((latents - ref_latent) ** 2)))
+      dt = abs(t_float - ref_timestep)
       accum_dx += dx_norm
       accum_dt += dt
 
       # Sensitivity score (Eq. 9)
       score = alpha_x * accum_dx + alpha_t * accum_dt
 
-      cond = jnp.logical_and(score <= sen_epsilon, reuse_count < max_reuse)
-
-      def hit_fn(args):
-        ref_noise_pred, ref_latent, ref_timestep, accum_dx, accum_dt, reuse_count, cache_count = args
-        return ref_noise_pred, ref_latent, ref_timestep, accum_dx, accum_dt, reuse_count + 1, cache_count + 1
-
-      def miss_fn(args):
-        _, _, _, _, _, _, cache_count = args
+      if score <= sen_epsilon and reuse_count < max_reuse:
+        noise_pred = ref_noise_pred
+        reuse_count += 1
+        cache_count += 1
+      else:
         latents_doubled = jnp.concatenate([latents] * 2)
         timestep = jnp.broadcast_to(t, bsz * 2)
         noise_pred, _, _ = transformer_forward_pass_full_cfg(
@@ -404,19 +401,12 @@ def miss_fn(args):
             rotary_emb=rotary_emb,
             encoder_attention_mask=encoder_attention_mask,
         )
-        return (
-            noise_pred,
-            latents,
-            jnp.array(t_float, dtype=jnp.float32),
-            jnp.array(0.0, dtype=jnp.float32),
-            jnp.array(0.0, dtype=jnp.float32),
-            jnp.array(0, dtype=jnp.int32),
-            cache_count,
-        )
-
-      noise_pred, ref_latent, ref_timestep, accum_dx, accum_dt, reuse_count, cache_count = jax.lax.cond(
-          cond, hit_fn, miss_fn, (ref_noise_pred, ref_latent, ref_timestep, accum_dx, accum_dt, reuse_count, cache_count)
-      )
+        ref_noise_pred = noise_pred
+        ref_latent = latents
+        ref_timestep = t_float
+        accum_dx = 0.0
+        accum_dt = 0.0
+        reuse_count = 0
 
       latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple()
 

src/maxdiffusion/pipelines/wan/wan_pipeline_i2v_2p2.py

@@ -491,21 +491,18 @@ def run_inference_2_2_i2v(
         latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple()
         continue
 
-      dx_norm = jnp.sqrt(jnp.mean((latents - ref_latent) ** 2))
-      dt = jnp.abs(t_float - ref_timestep)
+      dx_norm = float(jnp.sqrt(jnp.mean((latents - ref_latent) ** 2)))
+      dt = abs(t_float - ref_timestep)
       accum_dx += dx_norm
       accum_dt += dt
 
       score = alpha_x * accum_dx + alpha_t * accum_dt
 
-      cond = jnp.logical_and(score <= sen_epsilon, reuse_count < max_reuse)
-
-      def hit_fn(args):
-        ref_noise_pred, ref_latent, ref_timestep, accum_dx, accum_dt, reuse_count, cache_count = args
-        return ref_noise_pred, ref_latent, ref_timestep, accum_dx, accum_dt, reuse_count + 1, cache_count + 1
-
-      def miss_fn(args):
-        _, _, _, _, _, _, cache_count = args
+      if score <= sen_epsilon and reuse_count < max_reuse:
+        noise_pred = ref_noise_pred
+        reuse_count += 1
+        cache_count += 1
+      else:
         latents_doubled = jnp.concatenate([latents, latents], axis=0)
         latents_doubled = jnp.transpose(latents_doubled, (0, 4, 1, 2, 3))
         latent_model_input = jnp.concatenate([latents_doubled, condition_doubled], axis=1)
@@ -524,19 +521,12 @@ def miss_fn(args):
             encoder_attention_mask=encoder_attention_mask,
         )
         noise_pred = jnp.transpose(noise_pred, (0, 2, 3, 4, 1))
-        return (
-            noise_pred,
-            latents,
-            jnp.array(t_float, dtype=jnp.float32),
-            jnp.array(0.0, dtype=jnp.float32),
-            jnp.array(0.0, dtype=jnp.float32),
-            jnp.array(0, dtype=jnp.int32),
-            cache_count,
-        )
-
-      noise_pred, ref_latent, ref_timestep, accum_dx, accum_dt, reuse_count, cache_count = jax.lax.cond(
-          cond, hit_fn, miss_fn, (ref_noise_pred, ref_latent, ref_timestep, accum_dx, accum_dt, reuse_count, cache_count)
-      )
+        ref_noise_pred = noise_pred
+        ref_latent = latents
+        ref_timestep = t_float
+        accum_dx = 0.0
+        accum_dt = 0.0
+        reuse_count = 0
 
       latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple()