Device VAE on the 1xN mesh path (reclaims ~5s; 2-card 1024 E2E 1.68x->1.75x)

ttnn_moss/pipeline_mesh.py

@@ -20,7 +20,11 @@ from .dit_tp import TtMossImageTransformerTP
 class MeshDiTPipeline(MossImagePipeline):
     def __init__(self, mesh_device, weights_root, **kw):
         self.mesh = mesh_device
-        kw.setdefault("vae_device", False)   # host VAE on the mesh path
+        # VAE decode runs ON-DEVICE by default, REPLICATED across the mesh dies
+        # (the decoder isn't the compute bulk and needs no sharding; from_torch
+        # replicates the input, vae.py's _mesh_to_torch takes die-0's result).
+        # This reclaims the host-VAE penalty on the mesh path. MOSS_VAE_DEVICE=0
+        # still forces the host diffusers decode (the validated fallback).
         super().__init__(mesh_device, weights_root=weights_root, **kw)
 
     def _build_dit(self):

ttnn_moss/vae.py

@@ -49,6 +49,20 @@ GN_EPS = 1e-6
 _CORE = None  # set per-device in decoder init
 
 
+def _mesh_to_torch(device, t):
+    """to_torch that is mesh-aware. On a 1xN mesh the decoder runs REPLICATED
+    (from_torch replicates), so its output tensor is distributed and a plain
+    ttnn.to_torch fails ('supply a mesh composer'). Compose the identical shards
+    and keep the first copy; on a single device, unchanged."""
+    try:
+        multi = device.get_num_devices() > 1
+    except Exception:
+        multi = False
+    if multi:
+        return ttnn.to_torch(t, mesh_composer=ttnn.ConcatMeshToTensor(device, dim=0))[0:1]
+    return ttnn.to_torch(t)
+
+
 def _num_out_blocks(hw: int) -> int:
     """L1 tiling for DRAM-interleaved group_norm; larger HW -> more blocks."""
     if hw <= 4096:
@@ -350,6 +364,6 @@ class TtFlux2VaeDecoder:
         x = ttnn.to_memory_config(x, ttnn.DRAM_MEMORY_CONFIG)
         x, H, W = self.conv_out(x, H, W)
 
-        out = ttnn.to_torch(x).reshape(1, H, W, self.conv_out.cout).permute(0, 3, 1, 2)
+        out = _mesh_to_torch(self.device, x).reshape(1, H, W, self.conv_out.cout).permute(0, 3, 1, 2)
         ttnn.deallocate(x)
         return out.float().clamp(-1, 1)