Skip to content

vllm_omni.diffusion.distributed.autoencoders.wan_spatial_shard

Spatially-sharded Wan VAE decode.

The existing distributed Wan VAE path shards tiles. This module adds an opt-in decode backend that shards decoder feature maps along height or width and exchanges boundary rows/columns before spatial convolutions. It is intentionally decode-only and keeps checkpoint loading unchanged by patching the already-loaded decoder.

logger module-attribute 

logger = init_logger(__name__)

SpatialShardContext dataclass

input_extent instance-attribute 

input_extent: int

local_input_extent instance-attribute 

local_input_extent: int

rank instance-attribute 

rank: int

split_dim instance-attribute 

split_dim: str

world_size instance-attribute 

world_size: int

WanDistCausalConv3d

Bases: Conv3d

group instance-attribute 

group = group

halo_size instance-attribute 

halo_size = (self.kernel_extent - 1) // 2

kernel_extent instance-attribute 

kernel_extent = self.kernel_size[-2]

split_dim instance-attribute 

split_dim = split_dim

split_pad_left instance-attribute 

split_pad_left = int(self._source_padding[2])

split_pad_right instance-attribute 

split_pad_right = int(self._source_padding[3])

split_tensor_dim instance-attribute 

split_tensor_dim = _spatial_dim(split_dim)

stride_extent instance-attribute 

stride_extent = self.stride[-2]

forward 

forward(x: Tensor, cache_x: Tensor | None = None) -> Tensor

WanDistConv2d

Bases: Conv2d

group instance-attribute 

group = group

halo_size instance-attribute 

halo_size = (kernel_extent - 1) // 2

kernel_extent instance-attribute 

kernel_extent = self.kernel_size[-2]

split_dim instance-attribute 

split_dim = split_dim

split_tensor_dim instance-attribute 

split_tensor_dim = _spatial_dim(split_dim)

stride_extent instance-attribute 

stride_extent = self.stride[-2]

forward 

forward(x: Tensor) -> Tensor

WanDistZeroPad2d

Bases: Module

Apply ZeroPad2d only at global split-dimension boundaries.

group instance-attribute 

group = group

padding instance-attribute 

padding = padding

split_dim instance-attribute 

split_dim = split_dim

split_padding instance-attribute 

split_padding = split_padding or default_split_padding

forward 

forward(x: Tensor) -> Tensor

all_gather_along_dim 

all_gather_along_dim(
    x: Tensor,
    *,
    group: ProcessGroup,
    dim: int,
    dst: int | None = None,
) -> Tensor

gather_and_trim_extent 

gather_and_trim_extent(
    x: Tensor,
    *,
    expected_extent: int | None,
    split_dim: str,
    group: ProcessGroup,
    dst: int | None = None,
) -> Tensor

halo_exchange 

halo_exchange(
    x: Tensor,
    *,
    group: ProcessGroup,
    halo_size: int,
    split_dim: str = "height",
    recv_top_buf: Tensor | None = None,
    recv_bottom_buf: Tensor | None = None,
) -> tuple[Tensor, Tensor | None, Tensor | None]

install_wan_spatial_shard_decode 

install_wan_spatial_shard_decode(
    vae: Any, group: ProcessGroup, split_dim: str = "height"
) -> None

Patch vae.decoder once for spatially-sharded decode.

This mutates the already-loaded decoder in place by swapping its spatial convolutions/padding for halo-exchanging variants and wrapping decoder.forward. The patch is permanent for the lifetime of the VAE instance and is applied only once (subsequent calls are no-ops). A given instance is bound to a single split_dim; switching between "height" and "width" requires a fresh VAE instance and raises here otherwise.

Only group-relative rank 0 assembles the final decoded frame, mirroring the distributed tiled-decode broadcast_result=False contract; the other ranks take part in the collectives but return an empty placeholder.

reshard_from_trimmed_extent 

reshard_from_trimmed_extent(
    x: Tensor,
    *,
    local_extent: int,
    split_dim: str,
    group: ProcessGroup,
) -> Tensor

spatial_shard_decode 

spatial_shard_decode(
    vae: Any,
    z: Tensor,
    *,
    group: ProcessGroup,
    return_dict: bool = True,
    split_dim: str = "height",
) -> DecoderOutput | tuple[Tensor]

split_for_parallel_decode 

split_for_parallel_decode(
    x: Tensor,
    *,
    upsample_count: int,
    split_dim: str = "height",
    group: ProcessGroup | None = None,
    rank: int | None = None,
    world_size: int | None = None,
) -> tuple[Tensor, int]

Shard latent/feature spatial extent and return expected full output extent.