Merge branch 'main' into feat/diffusion-prefetch-protection

vllm_omni/diffusion/models/bagel/autoencoder.py

@@ -14,6 +14,16 @@ from dataclasses import dataclass
 import torch
 from einops import rearrange
 from torch import Tensor, nn
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+    DistributedOperator,
+    DistributedVaeMixin,
+    GridSpec,
+    TileTask,
+)
+
+logger = init_logger(__name__)
 
 
 @dataclass
@@ -322,3 +332,229 @@ class AutoEncoder(nn.Module):
 
     def forward(self, x: Tensor) -> Tensor:
         return self.decode(self.encode(x))
+
+
+class DistributedAutoEncoder(AutoEncoder, DistributedVaeMixin):
+    def __init__(self, params: AutoEncoderParams):
+        super().__init__(params)
+        self.init_distributed()
+
+        self.use_tiling = False
+        self.tile_sample_min_height = 512
+        self.tile_sample_min_width = 512
+        self.tile_sample_stride_height = 448
+        self.tile_sample_stride_width = 448
+        self.spatial_compression_ratio = params.downsample
+
+    def blend_v(self, above: Tensor, current: Tensor, blend_extent: int) -> Tensor:
+        blend_extent = min(above.shape[-2], current.shape[-2], blend_extent)
+        if blend_extent <= 0:
+            return current
+
+        for y in range(blend_extent):
+            alpha = y / blend_extent
+            current[:, :, y, :] = above[:, :, -blend_extent + y, :] * (1 - alpha) + current[:, :, y, :] * alpha
+        return current
+
+    def blend_h(self, left: Tensor, current: Tensor, blend_extent: int) -> Tensor:
+        blend_extent = min(left.shape[-1], current.shape[-1], blend_extent)
+        if blend_extent <= 0:
+            return current
+
+        for x in range(blend_extent):
+            alpha = x / blend_extent
+            current[:, :, :, x] = left[:, :, :, -blend_extent + x] * (1 - alpha) + current[:, :, :, x] * alpha
+        return current
+
+    def decode_tile_split(self, z: Tensor) -> tuple[list[TileTask], GridSpec]:
+        _, _, latent_height, latent_width = z.shape
+
+        sample_height = latent_height * self.spatial_compression_ratio
+        sample_width = latent_width * self.spatial_compression_ratio
+
+        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
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+        blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
+        blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
+
+        tiletask_list = []
+        for i in range(0, latent_height, tile_latent_stride_height):
+            for j in range(0, latent_width, tile_latent_stride_width):
+                tile = z[
+                    :,
+                    :,
+                    i : i + tile_latent_min_height,
+                    j : j + tile_latent_min_width,
+                ]
+                tiletask_list.append(
+                    TileTask(
+                        tile_id=len(tiletask_list),
+                        grid_coord=(i // tile_latent_stride_height, j // tile_latent_stride_width),
+                        tensor=tile,
+                        workload=tile.shape[-2] * tile.shape[-1],
+                    )
+                )
+        grid_spec = GridSpec(
+            split_dims=(2, 3),
+            grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+            tile_spec={
+                "sample_height": sample_height,
+                "sample_width": sample_width,
+                "blend_height": blend_height,
+                "blend_width": blend_width,
+            },
+            output_dtype=z.dtype,
+        )
+        return tiletask_list, grid_spec
+
+    def decode_tile_exec(self, task: TileTask) -> Tensor:
+        tile = task.tensor
+        tile = tile / self.scale_factor + self.shift_factor
+        return self.decoder(tile)
+
+    def decode_tile_merge(self, coord_tensor_map: dict[tuple[int, ...], Tensor], grid_spec: GridSpec) -> Tensor:
+        grid_h, grid_w = grid_spec.grid_shape
+        result_rows = []
+
+        for i in range(grid_h):
+            result_row = []
+            for j in range(grid_w):
+                tile = coord_tensor_map[(i, j)]
+                if i > 0:
+                    tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
+                if j > 0:
+                    tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
+
+                result_row.append(
+                    tile[
+                        :,
+                        :,
+                        : self.tile_sample_stride_height,
+                        : self.tile_sample_stride_width,
+                    ]
+                )
+            result_rows.append(torch.cat(result_row, dim=-1))
+
+        dec = torch.cat(result_rows, dim=-2)[
+            :,
+            :,
+            : grid_spec.tile_spec["sample_height"],
+            : grid_spec.tile_spec["sample_width"],
+        ]
+        return dec
+
+    def decode(self, z: Tensor) -> Tensor:
+        if not self.is_distributed_enabled():
+            return super().decode(z)
+
+        logger.debug("Bagel VAE decode running with distributed executor")
+        return self.distributed_executor.execute(
+            z,
+            DistributedOperator(
+                split=self.decode_tile_split,
+                exec=self.decode_tile_exec,
+                merge=self.decode_tile_merge,
+            ),
+            broadcast_result=True,
+        )
+
+    def encode_tile_split(self, x: Tensor) -> tuple[list[TileTask], GridSpec]:
+        _, _, height, width = x.shape
+
+        latent_height = height // self.spatial_compression_ratio
+        latent_width = width // self.spatial_compression_ratio
+
+        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
+
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+        blend_height = tile_latent_min_height - tile_latent_stride_height
+        blend_width = tile_latent_min_width - tile_latent_stride_width
+
+        tiletask_list = []
+        for i in range(0, height, self.tile_sample_stride_height):
+            for j in range(0, width, self.tile_sample_stride_width):
+                tile = x[
+                    :,
+                    :,
+                    i : i + self.tile_sample_min_height,
+                    j : j + self.tile_sample_min_width,
+                ]
+                tiletask_list.append(
+                    TileTask(
+                        tile_id=len(tiletask_list),
+                        grid_coord=(i // self.tile_sample_stride_height, j // self.tile_sample_stride_width),
+                        tensor=tile,
+                        workload=tile.shape[-2] * tile.shape[-1],
+                    )
+                )
+        grid_spec = GridSpec(
+            split_dims=(2, 3),
+            grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+            tile_spec={
+                "latent_height": latent_height,
+                "latent_width": latent_width,
+                "blend_height": blend_height,
+                "blend_width": blend_width,
+                "tile_latent_stride_height": tile_latent_stride_height,
+                "tile_latent_stride_width": tile_latent_stride_width,
+            },
+            output_dtype=next(self.parameters()).dtype,
+        )
+        return tiletask_list, grid_spec
+
+    def encode_tile_exec(self, task: TileTask) -> Tensor:
+        tile = task.tensor
+        z = self.reg(self.encoder(tile))
+        z = self.scale_factor * (z - self.shift_factor)
+        return z
+
+    def encode_tile_merge(self, coord_tensor_map: dict[tuple[int, ...], Tensor], grid_spec: GridSpec) -> Tensor:
+        grid_h, grid_w = grid_spec.grid_shape
+        result_rows = []
+
+        for i in range(grid_h):
+            result_row = []
+            for j in range(grid_w):
+                tile = coord_tensor_map[(i, j)]
+                if i > 0:
+                    tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
+                if j > 0:
+                    tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
+
+                result_row.append(
+                    tile[
+                        :,
+                        :,
+                        : grid_spec.tile_spec["tile_latent_stride_height"],
+                        : grid_spec.tile_spec["tile_latent_stride_width"],
+                    ]
+                )
+            result_rows.append(torch.cat(result_row, dim=-1))
+
+        enc = torch.cat(result_rows, dim=-2)[
+            :,
+            :,
+            : grid_spec.tile_spec["latent_height"],
+            : grid_spec.tile_spec["latent_width"],
+        ]
+        return enc
+
+    def encode(self, x: Tensor) -> Tensor:
+        if not self.is_distributed_enabled():
+            return super().encode(x)
+        logger.debug("Bagel VAE encode running with distributed executor")
+        return self.distributed_executor.execute(
+            x,
+            DistributedOperator(
+                split=self.encode_tile_split,
+                exec=self.encode_tile_exec,
+                merge=self.encode_tile_merge,
+            ),
+            broadcast_result=True,
+        )

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -32,7 +32,7 @@ from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPi
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
 
-from .autoencoder import AutoEncoder, AutoEncoderParams
+from .autoencoder import AutoEncoder, AutoEncoderParams, DistributedAutoEncoder
 from .bagel_transformer import Bagel, NaiveCache, Qwen2MoTConfig, Qwen2MoTForCausalLM
 
 logger = init_logger(__name__)
@@ -253,7 +253,7 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
         )
         self.transformer = self.language_model.model
         ae_params: AutoEncoderParams = default_ae_params()
-        self.vae = AutoEncoder(ae_params)
+        self.vae = DistributedAutoEncoder(ae_params)
 
         self.bagel = Bagel(
             language_model=self.language_model,

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

@@ -43,7 +43,6 @@ from vllm_omni.diffusion.distributed.sp_plan import (
 )
 from vllm_omni.diffusion.forward_context import get_forward_context
 from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm
-from vllm_omni.diffusion.layers.norm import RMSNorm
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 
 logger = init_logger(__name__)
@@ -538,8 +537,8 @@ class QwenImageCrossAttention(nn.Module):
         self.query_num_heads = self.to_qkv.num_heads
         self.kv_num_heads = self.to_qkv.num_kv_heads
 
-        self.norm_q = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
-        self.norm_k = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_q = nn.RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = nn.RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
 
         self.inner_dim = out_dim if out_dim is not None else head_dim * self.total_num_heads
 
@@ -576,8 +575,8 @@ class QwenImageCrossAttention(nn.Module):
             prefix=_join_prefix(prefix, "to_out"),
         )
 
-        self.norm_added_q = RMSNorm(head_dim, eps=eps)
-        self.norm_added_k = RMSNorm(head_dim, eps=eps)
+        self.norm_added_q = nn.RMSNorm(head_dim, eps=eps)
+        self.norm_added_k = nn.RMSNorm(head_dim, eps=eps)
 
         self.attn = Attention(
             num_heads=self.query_num_heads,
@@ -992,7 +991,7 @@ class QwenImageTransformer2DModel(CachedTransformer):
             quant_config=quant_config,
         )
 
-        self.txt_norm = RMSNorm(joint_attention_dim, eps=1e-6)
+        self.txt_norm = nn.RMSNorm(joint_attention_dim, eps=1e-6)
 
         # Entry projections (image/text) are kept full precision —
         # small sensitive layers at the network boundary (see #2728).

vllm_omni/diffusion/models/sd3/pipeline_sd3.py

@@ -402,12 +402,12 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
             truncation=True,
             add_special_tokens=True,
             return_tensors="pt",
-        ).to(self.device)
+        )
         text_input_ids = text_inputs.input_ids
         untruncated_ids = self.tokenizer_3(prompt, padding="longest", return_tensors="pt").input_ids
 
         if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
-            removed_text = self.tokenizer_3.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1])
+            removed_text = self.tokenizer_3.batch_decode(untruncated_ids[:, max_sequence_length - 1 : -1])
             logger.warning(
                 "The following part of your input was truncated because `max_sequence_length` is set to "
                 f" {max_sequence_length} tokens: {removed_text}"