[LingBot World] Eleven bit-identical removals of repeated work in the realtime path (#7648)

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py

@@ -505,22 +505,8 @@ class WanDistCausalConv3d(nn.Conv3d):
         self._trim_cache: dict[int, tuple[int, int, int]] = {}
 
     def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor:
-        padding = list(self._padding)
-        if cache_x is not None and padding[4] > 0:
-            cache_x = cache_x.to(x.device)
-            x = torch.cat([cache_x, x], dim=2)
-            padding[4] -= cache_x.shape[2]
-
-        x = F.pad(x, padding)
-        x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange(
-            x,
-            group=self.group,
-            halo_size=self.halo_size,
-            split_dim=self.split_dim,
-            recv_top_buf=self._halo_recv_top_buf,
-            recv_bottom_buf=self._halo_recv_bottom_buf,
-        )
-        shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim])
+        x_padded, local_extent = self._assemble_input(x, cache_x)
+        shift, start, upper_bound = self._get_trim_params(local_extent)
         if shift:
             x_padded = _narrow_along_dim(
                 x_padded,
@@ -532,6 +518,83 @@ class WanDistCausalConv3d(nn.Conv3d):
         out = _trim_local_conv_output(out, self.halo_size, start, upper_bound, split_dim=self.split_dim)
         return _zero_invalid_extent(out, split_dim=self.split_dim)
 
+    def _assemble_input(self, x: torch.Tensor, cache_x: torch.Tensor | None) -> tuple[torch.Tensor, int]:
+        """Build the conv input -- cached frames, causal/spatial zero padding and the halo columns -- in one pass.
+
+        The straightforward version wrote the activation three times per call: ``torch.cat`` of the cached
+        frames, ``F.pad`` for the causal and spatial zeros, then another ``torch.cat`` to attach the halos the
+        neighbours sent. Every causal conv of the decoder runs this on every latent frame, so at 480x832 that
+        was the largest single source of memory traffic in the decode. Here the final tensor is allocated once
+        and the cached frames and the new frames are copied straight into their slots; only the padding slices
+        are zeroed. The halo rows are exchanged through contiguous send and receive buffers (a P2P op needs a
+        contiguous tensor and the halo slot of a 5D tensor is a strided view) and copied into their slots. The
+        conv sees the same tensor as before, byte for byte.
+
+        Returns the assembled tensor and the extent along the split dimension the trim parameters are keyed by
+        (the extent before halos, as ``halo_exchange`` saw it).
+        """
+        pad_w0, pad_w1, pad_h0, pad_h1, pad_t0, pad_t1 = self._padding
+        cache_frames = 0
+        if cache_x is not None and pad_t0 > 0:
+            cache_x = cache_x.to(x.device)
+            cache_frames = int(cache_x.shape[2])
+            pad_t0 -= cache_frames
+        assert pad_t0 >= 0, "Wan temporal cache exceeds the causal padding"
+        rank, world_size = _rank_world(self.group)
+        halo = self.halo_size if world_size > 1 else 0
+        halo_h = halo if self.split_dim == "height" else 0
+        halo_w = halo if self.split_dim == "width" else 0
+
+        batch, channels, frames, height, width = x.shape
+        top, left = pad_h0 + halo_h, pad_w0 + halo_w
+        total_frames = pad_t0 + cache_frames + frames + pad_t1
+        total_height = top + height + pad_h1 + halo_h
+        total_width = left + width + pad_w1 + halo_w
+        buf = torch.empty((batch, channels, total_frames, total_height, total_width), dtype=x.dtype, device=x.device)
+        # Zero only what the conv must read as zeros: the causal frames in
+        # front (and any trailing ones) and the spatial padding rows/columns.
+        # The halo columns are filled below.
+        if pad_t0:
+            buf[:, :, :pad_t0].zero_()
+        if pad_t1:
+            buf[:, :, total_frames - pad_t1 :].zero_()
+        if pad_h0:
+            buf[:, :, :, halo_h:top].zero_()
+        if pad_h1:
+            buf[:, :, :, total_height - halo_h - pad_h1 : total_height - halo_h].zero_()
+        if pad_w0:
+            buf[:, :, :, :, halo_w:left].zero_()
+        if pad_w1:
+            buf[:, :, :, :, total_width - halo_w - pad_w1 : total_width - halo_w].zero_()
+        interior = buf[:, :, pad_t0 : pad_t0 + cache_frames + frames, top : top + height, left : left + width]
+        if cache_frames:
+            interior[:, :, :cache_frames].copy_(cache_x)
+            interior[:, :, cache_frames:].copy_(x)
+        else:
+            interior.copy_(x)
+
+        dim = self.split_tensor_dim
+        local_extent = buf.shape[dim] - 2 * halo
+        if halo:
+            # What halo_exchange sends: the padded tensor's first and last
+            # ``halo`` rows along the split dimension, padding included.
+            top_row_ref = _narrow_along_dim(buf, dim, halo, halo)
+            bottom_row_ref = _narrow_along_dim(buf, dim, buf.shape[dim] - 2 * halo, halo)
+            self._halo_recv_top_buf = _ensure_recv_buf(self._halo_recv_top_buf, top_row_ref)
+            self._halo_recv_bottom_buf = _ensure_recv_buf(self._halo_recv_bottom_buf, bottom_row_ref)
+            _halo_exchange_p2p(
+                rank=rank,
+                world_size=world_size,
+                group=self.group,
+                top_row_ref=top_row_ref,
+                bottom_row_ref=bottom_row_ref,
+                recv_top_buf=self._halo_recv_top_buf,
+                recv_bottom_buf=self._halo_recv_bottom_buf,
+            )
+            _narrow_along_dim(buf, dim, 0, halo).copy_(self._halo_recv_top_buf)
+            _narrow_along_dim(buf, dim, buf.shape[dim] - halo, halo).copy_(self._halo_recv_bottom_buf)
+        return buf, local_extent
+
     def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]:
         trim_params = self._trim_cache.get(local_extent)
         if trim_params is None: