From 0bdc15d20f8100567801c906741708a5d9541509 Mon Sep 17 00:00:00 2001 From: "jacky.cheng" Date: Sat, 5 Sep 2026 16:00:25 +0800 Subject: [PATCH] [AMD] Fix ROCm VAE Conv2D fast path breaking spatial-parallel decode (#34424) --- .../runtime/layers/parallel_conv.py | 33 ++++++++++- .../multimodal_gen/runtime/platforms/rocm.py | 58 +++++++++++++++++-- 2 files changed, 84 insertions(+), 7 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/parallel_conv.py b/python/sglang/multimodal_gen/runtime/layers/parallel_conv.py index baf49b964..04daace52 100644 --- a/python/sglang/multimodal_gen/runtime/layers/parallel_conv.py +++ b/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, ) diff --git a/python/sglang/multimodal_gen/runtime/platforms/rocm.py b/python/sglang/multimodal_gen/runtime/platforms/rocm.py index 2dac954fb..4c87385b8 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/rocm.py +++ b/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, ) - child.forward = types.MethodType(_patched_forward, child) - patched += 1 + 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), " - "skipped %d (1D/pointwise/grouped)", + "Conv3D→Conv2D: patched %d CausalConv3d + %d spatial-parallel halo " + "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: