Simplify kv_cache initialization

src/exo/worker/engines/mflux/pipefusion/distributed_denoising.py

@@ -198,7 +198,7 @@ class DistributedDenoising:
         text_seq_len = prompt_embeds.shape[1]
         hidden_dim = self.transformer.x_embedder.weight.shape[0]
 
-        if self._joint_kv_caches is None:
+        if t == 0:
             self._initialize_kv_caches(
                 batch_size=batch_size,
                 num_img_tokens=num_img_tokens,
@@ -493,10 +493,6 @@ class DistributedDenoising:
         controlnet_single_block_samples: list[mx.array] | None = None,
         kontext_image_ids: mx.array | None = None,
     ) -> mx.array:
-        if t == 0:
-            self._joint_kv_caches = None
-            self._single_kv_caches = None
-
         if t < self.num_sync_steps:
             latents = self._sync_pipeline(
                 t,