[NPU][diffusion]: support parallel decoding of qwen-image (#20757)

Co-authored-by: 高鑫 <gaoxin@gaoxindeMacBook-Pro.local>

python/sglang/multimodal_gen/configs/models/vaes/qwenimage.py

@@ -38,6 +38,8 @@ class QwenImageVAEConfig(VAEConfig):
     use_temporal_tiling: bool = False
     use_parallel_tiling: bool = False
 
+    use_parallel_decode: bool = False
+
     def get_vae_scale_factor(self):
         return 2 ** len(self.arch_config.temperal_downsample)
 

python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_qwenimage.py

@@ -3,6 +3,7 @@
 from typing import Optional, Tuple, Union
 
 import torch
+import torch.distributed as dist
 import torch.nn as nn
 import torch.nn.functional as F
 from diffusers.models.activations import get_activation
@@ -13,7 +14,10 @@ from diffusers.models.autoencoders.vae import (
 from diffusers.models.modeling_outputs import AutoencoderKLOutput
 
 from sglang.multimodal_gen.configs.models.vaes.qwenimage import QwenImageVAEConfig
-from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.distributed import (
+    get_local_torch_device,
+    get_sp_world_size,
+)
 from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
@@ -789,6 +793,7 @@ class AutoencoderKLQwenImage(ParallelTiledVAE):
         self.input_channels = config.arch_config.input_channels
         self.latents_mean = config.arch_config.latents_mean
         self.config = config.arch_config
+        self.use_parallel_decode = config.use_parallel_decode
 
         self.encoder = QwenImageEncoder3d(
             base_dim, z_dim * 2, dim_mult, num_res_blocks, attn_scales, self.temperal_downsample, dropout,
@@ -841,6 +846,8 @@ class AutoencoderKLQwenImage(ParallelTiledVAE):
             .to(cuda_device, dtype)
         )
 
+
+
     def enable_tiling(
         self,
         tile_sample_min_height: Optional[int] = None,
@@ -956,30 +963,43 @@ class AutoencoderKLQwenImage(ParallelTiledVAE):
 
         return posterior
 
-    def _decode(self, z: torch.Tensor, return_dict: bool = True):
+    def _decode_with_parallel_dispatch(self, z: torch.Tensor) -> DecoderOutput:
+        if self.use_parallel_decode and get_sp_world_size() > 1:
+            num_frame = z.shape[2]
+            num_sample_frames = (num_frame - 1) * self.temporal_compression_ratio + 1
+            decoded = super().parallel_tiled_decode(z)[:, :, :num_sample_frames]
+            return DecoderOutput(sample=decoded)
+
+        return DecoderOutput(sample=self._decode(z))
+
+    def _decode(self, z: torch.Tensor) -> torch.Tensor:
         _, _, num_frame, height, width = z.shape
         tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
         tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
 
         if self.use_tiling and (width > tile_latent_min_width or height > tile_latent_min_height):
-            return self.tiled_decode(z, return_dict=return_dict)
+            return self.tiled_decode(z).sample
 
         self.clear_cache()
         x = self.post_quant_conv(z)
         for i in range(num_frame):
             self._conv_idx = [0]
             if i == 0:
-                out = self.decoder(x[:, :, i: i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx)
+                out = self.decoder(
+                    x[:, :, i : i + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx,
+                )
             else:
-                out_ = self.decoder(x[:, :, i: i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx)
+                out_ = self.decoder(
+                    x[:, :, i : i + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx,
+                )
                 out = torch.cat([out, out_], 2)
-
         out = torch.clamp(out, min=-1.0, max=1.0)
         self.clear_cache()
-        if not return_dict:
+        return out
-            return (out,)
-
-        return DecoderOutput(sample=out)
 
     def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
         r"""
@@ -996,29 +1016,121 @@ class AutoencoderKLQwenImage(ParallelTiledVAE):
                 returned.
         """
         if self.use_slicing and z.shape[0] > 1:
-            decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
+            decoded_slices = [
+                self._decode_with_parallel_dispatch(z_slice).sample
+                for z_slice in z.split(1)
+            ]
             decoded = torch.cat(decoded_slices)
         else:
-            decoded = self._decode(z).sample
+            decoded = self._decode_with_parallel_dispatch(z).sample
 
         return decoded
 
-    def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+    def blend_v(
+        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+    ) -> torch.Tensor:
         blend_extent = min(a.shape[-2], b.shape[-2], blend_extent)
-        for y in range(blend_extent):
+        if blend_extent <= 0:
-            b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * (
+            return b
-                y / blend_extent
+        weight = (
-            )
+            torch.arange(blend_extent, device=b.device, dtype=b.dtype) / blend_extent
+        ).view(1, 1, 1, blend_extent, 1)
+        b[:, :, :, :blend_extent, :] = (
+            a[:, :, :, -blend_extent:, :] * (1 - weight)
+            + b[:, :, :, :blend_extent, :] * weight
+        )
         return b
 
-    def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
+    def blend_h(
+        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+    ) -> torch.Tensor:
         blend_extent = min(a.shape[-1], b.shape[-1], blend_extent)
-        for x in range(blend_extent):
+        if blend_extent <= 0:
-            b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * (
+            return b
-                x / blend_extent
+        weight = (
-            )
+            torch.arange(blend_extent, device=b.device, dtype=b.dtype) / blend_extent
+        ).view(1, 1, 1, 1, blend_extent)
+        b[:, :, :, :, :blend_extent] = (
+            a[:, :, :, :, -blend_extent:] * (1 - weight)
+            + b[:, :, :, :, :blend_extent] * weight
+        )
         return b
 
+    def _process_parallel_tiled_outputs(
+        self,
+        results: torch.Tensor,
+        local_dim_metadata: list[torch.Size],
+        z: torch.Tensor,
+        world_size: int,
+        rank: int,
+        num_t_tiles: int,
+        num_h_tiles: int,
+        num_w_tiles: int,
+        total_spatial_tiles: int,
+        blend_height: int,
+        blend_width: int,
+    ) -> torch.Tensor:
+        local_size = torch.tensor(
+            [results.size(0)], device=results.device, dtype=torch.int64
+        )
+        if rank == 0:
+            gathered_sizes = [
+                torch.zeros(1, device=results.device, dtype=torch.int64)
+                for _ in range(world_size)
+            ]
+        else:
+            gathered_sizes = None
+        dist.gather(local_size, gather_list=gathered_sizes, dst=0)
+
+        max_size = 0
+        if rank == 0:
+            max_size = max(size.item() for size in gathered_sizes)
+
+        max_size_tensor = torch.tensor(
+            [max_size], device=results.device, dtype=torch.int64
+        )
+        dist.broadcast(max_size_tensor, src=0)
+        max_size = int(max_size_tensor.item())
+
+        padded_results = torch.zeros(
+            max_size, device=results.device, dtype=results.dtype
+        )
+        padded_results[: results.size(0)] = results
+
+        gathered_dim_metadata = [None] * world_size
+        dist.all_gather_object(gathered_dim_metadata, local_dim_metadata)
+
+        if rank == 0:
+            gathered_results = [
+                torch.empty_like(padded_results) for _ in range(world_size)
+            ]
+        else:
+            gathered_results = None
+        dist.gather(padded_results, gather_list=gathered_results, dst=0)
+
+        if rank == 0:
+            gathered_results = torch.stack(gathered_results, dim=0).contiguous()
+            dec = super()._merge_parallel_tiled_results(
+                gathered_results,
+                gathered_dim_metadata,
+                num_t_tiles,
+                num_h_tiles,
+                num_w_tiles,
+                total_spatial_tiles,
+                blend_height,
+                blend_width,
+            )
+            shape_tensor = torch.tensor(dec.shape, device=dec.device, dtype=torch.int64)
+        else:
+            dec = None
+            shape_tensor = torch.zeros(5, device=z.device, dtype=torch.int64)
+
+        dist.broadcast(shape_tensor, src=0)
+        if rank != 0:
+            dec = z.new_empty(tuple(shape_tensor.tolist()))
+        dist.broadcast(dec, src=0)
+        return dec
+
     def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput:
         r"""Encode a batch of images using a tiled encoder.
 

python/sglang/multimodal_gen/runtime/models/vaes/common.py

@@ -220,14 +220,110 @@ class ParallelTiledVAE(ABC, nn.Module):
                 _start_shape += mul_shape
                 global_idx += 1
 
+    def _merge_parallel_tiled_results(
+        self,
+        gathered_results: torch.Tensor,
+        gathered_dim_metadata: list[list[torch.Size]],
+        num_t_tiles: int,
+        num_h_tiles: int,
+        num_w_tiles: int,
+        total_spatial_tiles: int,
+        blend_height: int,
+        blend_width: int,
+    ) -> torch.Tensor:
+        data: list = [
+            [[[] for _ in range(num_w_tiles)] for _ in range(num_h_tiles)]
+            for _ in range(num_t_tiles)
+        ]
+        for current_data, global_idx in self._parallel_data_generator(
+            gathered_results, gathered_dim_metadata
+        ):
+            t_idx = global_idx // total_spatial_tiles
+            spatial_idx = global_idx % total_spatial_tiles
+            h_idx = spatial_idx // num_w_tiles
+            w_idx = spatial_idx % num_w_tiles
+            data[t_idx][h_idx][w_idx] = current_data
+
+        result_slices = []
+        last_slice_data = None
+        for i, tem_data in enumerate(data):
+            slice_data = self._merge_spatial_tiles(
+                tem_data,
+                blend_height,
+                blend_width,
+                self.tile_sample_stride_height,
+                self.tile_sample_stride_width,
+            )
+            if i > 0:
+                slice_data = self.blend_t(
+                    last_slice_data, slice_data, self.blend_num_frames
+                )
+                result_slices.append(
+                    slice_data[:, :, : self.tile_sample_stride_num_frames, :, :]
+                )
+            else:
+                result_slices.append(
+                    slice_data[:, :, : self.tile_sample_stride_num_frames + 1, :, :]
+                )
+            last_slice_data = slice_data
+        return torch.cat(result_slices, dim=2)
+
+    def _process_parallel_tiled_outputs(
+        self,
+        results: torch.Tensor,
+        local_dim_metadata: list[torch.Size],
+        z: torch.Tensor,
+        world_size: int,
+        rank: int,
+        num_t_tiles: int,
+        num_h_tiles: int,
+        num_w_tiles: int,
+        total_spatial_tiles: int,
+        blend_height: int,
+        blend_width: int,
+    ) -> torch.Tensor:
+        local_size = torch.tensor(
+            [results.size(0)], device=results.device, dtype=torch.int64
+        )
+        all_sizes = [
+            torch.zeros(1, device=results.device, dtype=torch.int64)
+            for _ in range(world_size)
+        ]
+        dist.all_gather(all_sizes, local_size)
+        max_size = max(size.item() for size in all_sizes)
+
+        padded_results = torch.zeros(
+            max_size, device=results.device, dtype=results.dtype
+        )
+        padded_results[: results.size(0)] = results
+
+        gathered_dim_metadata = [None] * world_size
+        gathered_results = (
+            torch.zeros_like(padded_results)
+            .repeat(world_size, *[1] * len(padded_results.shape))
+            .contiguous()
+        )
+        dist.all_gather_into_tensor(gathered_results, padded_results)
+        dist.all_gather_object(gathered_dim_metadata, local_dim_metadata)
+        gathered_dim_metadata = cast(list[list[torch.Size]], gathered_dim_metadata)
+        return self._merge_parallel_tiled_results(
+            gathered_results,
+            gathered_dim_metadata,
+            num_t_tiles,
+            num_h_tiles,
+            num_w_tiles,
+            total_spatial_tiles,
+            blend_height,
+            blend_width,
+        )
+
     def parallel_tiled_decode(self, z: torch.FloatTensor) -> torch.FloatTensor:
         """
         Parallel version of tiled_decode that distributes both temporal and spatial computation across GPUs
         """
         world_size, rank = get_sp_world_size(), get_sp_parallel_rank()
-        B, C, T, H, W = z.shape
+        _, _, T, H, W = z.shape
 
-        # Calculate parameters
         tile_latent_min_height = (
             self.tile_sample_min_height // self.spatial_compression_ratio
         )
@@ -259,26 +355,22 @@ class ParallelTiledVAE(ABC, nn.Module):
         total_spatial_tiles = num_h_tiles * num_w_tiles
         total_tiles = num_t_tiles * total_spatial_tiles
 
-        # Calculate tiles per rank and padding
         tiles_per_rank = (total_tiles + world_size - 1) // world_size
         start_tile_idx = rank * tiles_per_rank
         end_tile_idx = min((rank + 1) * tiles_per_rank, total_tiles)
 
         local_results = []
         local_dim_metadata = []
-        # Process assigned tiles
+        for global_idx in range(start_tile_idx, end_tile_idx):
-        for local_idx, global_idx in enumerate(range(start_tile_idx, end_tile_idx)):
             t_idx = global_idx // total_spatial_tiles
             spatial_idx = global_idx % total_spatial_tiles
             h_idx = spatial_idx // num_w_tiles
             w_idx = spatial_idx % num_w_tiles
 
-            # Calculate positions
             t_start = t_idx * tile_latent_stride_num_frames
             h_start = h_idx * tile_latent_stride_height
             w_start = w_idx * tile_latent_stride_width
 
-            # Extract and process tile
             tile = z[
                 :,
                 :,
@@ -286,84 +378,31 @@ class ParallelTiledVAE(ABC, nn.Module):
                 h_start : h_start + tile_latent_min_height,
                 w_start : w_start + tile_latent_min_width,
             ]
-
+            decoded_tile = self._decode(tile)
-            # Process tile
-            tile = self._decode(tile)
-
             if t_start > 0:
-                tile = tile[:, :, 1:, :, :]
+                decoded_tile = decoded_tile[:, :, 1:, :, :]
+            local_results.append(decoded_tile.reshape(-1))
+            local_dim_metadata.append(decoded_tile.shape)
 
-            # Store metadata
+        if local_results:
-            shape = tile.shape
+            results = torch.cat(local_results, dim=0).contiguous()
-            # Store decoded data (flattened)
+        else:
-            decoded_flat = tile.reshape(-1)
+            results = z.new_empty((0,), dtype=z.dtype)
-            local_results.append(decoded_flat)
-            local_dim_metadata.append(shape)
-
-        results = torch.cat(local_results, dim=0).contiguous()
         del local_results
-        # first gather size to pad the results
+
-        local_size = torch.tensor(
+        dec = self._process_parallel_tiled_outputs(
-            [results.size(0)], device=results.device, dtype=torch.int64
+            results,
+            local_dim_metadata,
+            z,
+            world_size,
+            rank,
+            num_t_tiles,
+            num_h_tiles,
+            num_w_tiles,
+            total_spatial_tiles,
+            blend_height,
+            blend_width,
         )
-        all_sizes = [
-            torch.zeros(1, device=results.device, dtype=torch.int64)
-            for _ in range(world_size)
-        ]
-        dist.all_gather(all_sizes, local_size)
-        max_size = max(size.item() for size in all_sizes)
-        padded_results = torch.zeros(max_size, device=results.device)
-        padded_results[: results.size(0)] = results
-        del results
-
-        # Gather all results
-        gathered_dim_metadata = [None] * world_size
-        gathered_results = (
-            torch.zeros_like(padded_results)
-            .repeat(world_size, *[1] * len(padded_results.shape))
-            .contiguous()
-        )  # use contiguous to make sure it won't copy data in the following operations
-        # TODO (PY): use sgl_diffusion distributed methods
-        dist.all_gather_into_tensor(gathered_results, padded_results)
-        dist.all_gather_object(gathered_dim_metadata, local_dim_metadata)
-        # Process gathered results
-        data: list = [
-            [[[] for _ in range(num_w_tiles)] for _ in range(num_h_tiles)]
-            for _ in range(num_t_tiles)
-        ]
-        for current_data, global_idx in self._parallel_data_generator(
-            gathered_results, gathered_dim_metadata
-        ):
-            t_idx = global_idx // total_spatial_tiles
-            spatial_idx = global_idx % total_spatial_tiles
-            h_idx = spatial_idx // num_w_tiles
-            w_idx = spatial_idx % num_w_tiles
-            data[t_idx][h_idx][w_idx] = current_data
-        # Merge results
-        result_slices = []
-        last_slice_data = None
-        for i, tem_data in enumerate(data):
-            slice_data = self._merge_spatial_tiles(
-                tem_data,
-                blend_height,
-                blend_width,
-                self.tile_sample_stride_height,
-                self.tile_sample_stride_width,
-            )
-            if i > 0:
-                slice_data = self.blend_t(
-                    last_slice_data, slice_data, self.blend_num_frames
-                )
-                result_slices.append(
-                    slice_data[:, :, : self.tile_sample_stride_num_frames, :, :]
-                )
-            else:
-                result_slices.append(
-                    slice_data[:, :, : self.tile_sample_stride_num_frames + 1, :, :]
-                )
-            last_slice_data = slice_data
-        dec = torch.cat(result_slices, dim=2)
-
         return dec
 
     def _merge_spatial_tiles(