[AMD] Fix ROCm VAE Conv2D fast path breaking spatial-parallel decode (#34424)

python/sglang/multimodal_gen/runtime/layers/parallel_conv.py

@@ -183,6 +183,11 @@ def _maybe_contiguous_for_sp_gather(x: torch.Tensor) -> torch.Tensor:
         and not x.is_contiguous()
     ):
         return x.contiguous()
+    # Permuted views (e.g. the platform Conv2D fast path's NTCHW -> NCTHW
+    # permute) are neither contiguous nor channels-last; NCCL still needs a
+    # contiguous buffer for the height gather.
+    if not x.is_contiguous():
+        return x.contiguous()
     return x
 
 
@@ -634,6 +639,9 @@ class SpatialParallelCausalConv3d(nn.Conv3d):
         self.padding = (0, 0, 0)
         self._halo_recv_top_buf: torch.Tensor | None = None
         self._halo_recv_bottom_buf: torch.Tensor | None = None
+        # Set only by the ROCm Conv3D->Conv2D fast path, to swap the inner
+        # conv without displacing the halo exchange and output trim.
+        self._halo_conv_forward = None
         self.rank = get_decode_parallel_rank()
         self.world_size = get_decode_parallel_world_size()
 
@@ -657,10 +665,19 @@ class SpatialParallelCausalConv3d(nn.Conv3d):
                 self.groups,
             )
 
+        # Bind ``super().forward`` lazily: doing it unconditionally costs two
+        # extra Dynamo frames per conv when the ROCm hook is what runs.  The
+        # hook is only ever installed on ROCm, so it doubles as the platform
+        # check.
+        if self._halo_conv_forward is not None:
+            conv_forward = self._halo_conv_forward
+        else:
+            conv_forward = super().forward
+
         return _spatial_parallel_conv_forward(
             self,
             x,
-            super().forward,
+            conv_forward,
             height_pad_mode="zeros",
             match_conv3d_format=True,
         )
@@ -712,6 +729,9 @@ class SpatialParallelConv3d(nn.Conv3d):
         _set_conv_padding(self, (self.padding[0], 0, self.padding[2]))
         self._halo_recv_top_buf: torch.Tensor | None = None
         self._halo_recv_bottom_buf: torch.Tensor | None = None
+        # Set only by the ROCm Conv3D->Conv2D fast path, to swap the inner
+        # conv without displacing the halo exchange and output trim.
+        self._halo_conv_forward = None
         self.rank = get_decode_parallel_rank()
         self.world_size = get_decode_parallel_world_size()
 
@@ -722,10 +742,19 @@ class SpatialParallelConv3d(nn.Conv3d):
         if any(self._padding):
             x = _pad_with_mode(x, self._padding, self.padding_mode)
 
+        # Bind ``super().forward`` lazily: doing it unconditionally costs two
+        # extra Dynamo frames per conv when the ROCm hook is what runs.  The
+        # hook is only ever installed on ROCm, so it doubles as the platform
+        # check.
+        if self._halo_conv_forward is not None:
+            conv_forward = self._halo_conv_forward
+        else:
+            conv_forward = super().forward
+
         return _spatial_parallel_conv_forward(
             self,
             x,
-            super().forward,
+            conv_forward,
             height_pad_mode=self.padding_mode,
             match_conv3d_format=True,
         )

python/sglang/multimodal_gen/runtime/platforms/rocm.py

@@ -332,8 +332,16 @@ class RocmPlatform(Platform):
         Kw>1) are replaced; pointwise or 1-D-temporal convolutions are left
         untouched.  Modules with non-default ``groups`` or ``dilation`` are
         skipped as the 2-D decomposition assumes groups=1 and dilation=1.
+
+        Spatial-parallel convs shard the height dimension across ranks and get
+        their missing rows from a halo exchange, so their ``_padding`` carries
+        no height padding.  Replacing their ``forward`` would drop the halo
+        exchange and silently shrink the output height, so for those the 2-D
+        decomposition is installed as the inner ``_halo_conv_forward`` kernel
+        instead, leaving the halo exchange and output trim intact.
         """
         patched = 0
+        patched_halo = 0
         skipped = 0
         for _name, child in module.named_modules():
             if not isinstance(child, nn.Conv3d):
@@ -348,6 +356,15 @@ class RocmPlatform(Platform):
                 skipped += 1
                 continue
 
+            is_spatial_parallel = hasattr(child, "height_halo_size")
+            if is_spatial_parallel and (
+                not hasattr(child, "_halo_conv_forward")
+                or child.padding_mode != "zeros"
+            ):
+                # No safe hook to patch without breaking the halo exchange.
+                skipped += 1
+                continue
+
             padding = child._padding
             stride = child.stride
 
@@ -386,17 +403,48 @@ class RocmPlatform(Platform):
                     compute_bf16=_bf16,
                 )
 
+            def _patched_halo_conv_forward(
+                self,
+                x,
+                *,
+                _stride=stride,
+                _kt=kt,
+                _bf16=use_bf16,
+            ):
+                # ``x`` is already halo-exchanged and causally padded; only the
+                # conv's own ``padding`` is still outstanding.
+                pad_t, pad_h, pad_w = self.padding
+                if pad_t or pad_h or pad_w:
+                    x = F.pad(x, (pad_w, pad_w, pad_h, pad_h, pad_t, pad_t))
+                x = x.to(self.weight.dtype)
+                return RocmPlatform._conv3d_as_batched_conv2d(
+                    x,
+                    self._weight_2d,
+                    self.bias,
+                    _stride,
+                    _kt,
+                    compute_bf16=_bf16,
+                )
+
+            if is_spatial_parallel:
+                child._halo_conv_forward = types.MethodType(
+                    _patched_halo_conv_forward, child
+                )
+                patched_halo += 1
+            else:
                 child.forward = types.MethodType(_patched_forward, child)
                 patched += 1
 
         logger.info(
-            "Conv3D→Conv2D: patched %d CausalConv3d (3D kernel, compute=%s), "
+            "Conv3D→Conv2D: patched %d CausalConv3d + %d spatial-parallel halo "
-            "skipped %d (1D/pointwise/grouped)",
+            "kernels (3D kernel, compute=%s), skipped %d (1D/pointwise/grouped/"
+            "unsupported spatial-parallel)",
             patched,
+            patched_halo,
             "BF16" if use_bf16 else "same dtype",
             skipped,
         )
-        return patched
+        return patched + patched_halo
 
     @classmethod
     def enable_dit_layerwise_offload_by_default(cls) -> bool: