From 7ca0bc5b557e5bc1801fb74dc34715bdf1ff6529 Mon Sep 17 00:00:00 2001
From: ciaranbor <ciaranborourke-dev@proton.me>
Date: Sat, 13 Dec 2025 16:44:58 +0000
Subject: [PATCH] Only send encoder_hidden_states with the first patch (once
 per timestep)

---
 .../mflux/pipefusion/distributed_denoising.py | 23 ++++++++++++-------
 1 file changed, 15 insertions(+), 8 deletions(-)

diff --git a/src/exo/worker/engines/mflux/pipefusion/distributed_denoising.py b/src/exo/worker/engines/mflux/pipefusion/distributed_denoising.py
index 3c7c5b44..81fab4bf 100644
--- a/src/exo/worker/engines/mflux/pipefusion/distributed_denoising.py
+++ b/src/exo/worker/engines/mflux/pipefusion/distributed_denoising.py
@@ -348,8 +348,10 @@ class DistributedDenoising:
         single_kv_caches = self._single_kv_caches
 
         # === Process each patch ===
+        # Encoder hidden states are the same for all patches in a timestep,
+        # so we only need to receive them once (with the first patch)
         output_patches = []
-        for start_token, end_token in token_indices:
+        for patch_idx, (start_token, end_token) in enumerate(token_indices):
             # Extract current patch from full hidden states
             patch_hidden = full_hidden[:, start_token:end_token, :]
 
@@ -360,9 +362,11 @@ class DistributedDenoising:
                     patch_hidden = mx.distributed.recv_like(
                         patch_hidden, self.rank - 1, group=self.group
                     )
-                    encoder_hidden_states = mx.distributed.recv_like(
-                        encoder_hidden_states, self.rank - 1, group=self.group
-                    )
+                    # Only receive encoder_hidden_states once per timestep (with first patch)
+                    if patch_idx == 0:
+                        encoder_hidden_states = mx.distributed.recv_like(
+                            encoder_hidden_states, self.rank - 1, group=self.group
+                        )
 
                     mx.eval(patch_hidden, encoder_hidden_states)
 
@@ -400,12 +404,15 @@ class DistributedDenoising:
 
             elif self.has_joint_blocks and not self.is_last_stage:
                 # Send joint block outputs to next stage
-                mx.eval(
-                    mx.distributed.send(patch_hidden, self.rank + 1, group=self.group),
+                # Only send encoder_hidden_states once per timestep (with first patch)
+                if patch_idx == 0:
                     mx.distributed.send(
                         encoder_hidden_states, self.rank + 1, group=self.group
-                    ),
-                )
+                    )
+                else:
+                    mx.distributed.send(patch_hidden, self.rank + 1, group=self.group)
+
+                mx.eval(patch_hidden, encoder_hidden_states)
 
             # === PHASE 4: Single Blocks with KV Cache ===
             if self.has_single_blocks: