anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[New model] Qwen-Image-2.1 support (rebase + review follow-ups) (#8099)

Browse Source 
Signed-off-by: NumberWan <wantszkin2003@gmail.com>
Signed-off-by: SamitHuang <285365963@qq.com>
Signed-off-by: Samit <285365963@qq.com>
Co-authored-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: Tianyu Guo <guoty9@mail2.sysu.edu.cn>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: SamitHuang <285365963@qq.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
This commit is contained in:
NumberWanIsotr0pyTianyu GuoCursorSamitHuangGao Han
 authored and GitHub committed 2026-10-07 21:23:59 +08:00
1 parent 8ef77be66e
commit 3dc35694b3
52 files changed
+8058 -136
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py
+73 -39
View File
@@ -3,6 +3,8 @@
from __future__ import annotations
from collections.abc import Callable
from contextlib import contextmanager
from dataclasses import dataclass, field
import torch
@@ -34,9 +36,9 @@ class TileTask:
@dataclass
class DistributedOperator:
split: callable
exec: callable
merge: callable
split: Callable[[torch.Tensor], tuple[list[TileTask], GridSpec]]
exec: Callable[[TileTask], torch.Tensor]
merge: Callable[[dict[tuple[int, ...], torch.Tensor], GridSpec], torch.Tensor]
class DistributedVaeExecutor:
@@ -46,7 +48,11 @@ class DistributedVaeExecutor:
def __init__(self):
# Use a dedicated process group spanning the complete worker WORLD.
self.group = get_world_group().device_group
world_group = get_world_group()
self.group = world_group.device_group
self._error_group = world_group.cpu_group
# Reporting an OOM must not require another device allocation.
self._error_status = torch.zeros((), dtype=torch.int32, device="cpu")
self.world_size = dist.get_world_size(self.group)
self.rank = dist.get_rank(self.group)
self.parallel_size = 1
@@ -56,8 +62,34 @@ class DistributedVaeExecutor:
self.parallel_size = parallel_size
self.parallel_mode = mode
@contextmanager
def _sync_errors(self, stage: str):
"""Agree on local-stage failures before entering the next collective."""
status = 0
message = ""
try:
yield
except torch.OutOfMemoryError as error:
status = 1
message = str(error)
except Exception as error:
# A non-retryable error takes precedence over OOM on another rank.
status = 2
message = f"{type(error).__name__}: {error}"
self._error_status.fill_(status)
dist.all_reduce(self._error_status, op=dist.ReduceOp.MAX, group=self._error_group)
global_status = self._error_status.item()
if global_status == 2:
detail = message if status == 2 else "another rank failed"
raise RuntimeError(f"Distributed VAE {stage} failed: {detail}")
if global_status == 1:
detail = message if status == 1 else "another rank ran out of memory"
raise torch.OutOfMemoryError(f"Distributed VAE {stage} failed: {detail}")
def gather_tensors(self, tensor: torch.Tensor):
gather_list = [torch.empty_like(tensor) for _ in range(self.world_size)]
with self._sync_errors("gather allocation"):
gather_list = [torch.empty_like(tensor) for _ in range(self.world_size)]
dist.all_gather(gather_list, tensor, group=self.group)
return gather_list if self.rank == 0 else None
@@ -66,11 +98,12 @@ class DistributedVaeExecutor:
return tensor
def _compute_global_padding_shape(self, local_results, output_ndim: int, device):
local_tile_max_dims = [0] * output_ndim
for _, tile_tensor in local_results:
for dim_idx, dim_size in enumerate(tile_tensor.shape):
local_tile_max_dims[dim_idx] = max(local_tile_max_dims[dim_idx], dim_size)
local_shape_stat = torch.tensor([len(local_results), *local_tile_max_dims], device=device)
with self._sync_errors("shape preparation"):
local_tile_max_dims = [0] * output_ndim
for _, tile_tensor in local_results:
for dim_idx, dim_size in enumerate(tile_tensor.shape):
local_tile_max_dims[dim_idx] = max(local_tile_max_dims[dim_idx], dim_size)
local_shape_stat = torch.tensor([len(local_results), *local_tile_max_dims], device=device)
dist.all_reduce(
local_shape_stat,
op=dist.ReduceOp.MAX,
@@ -100,17 +133,16 @@ class DistributedVaeExecutor:
tid = int(meta_src[idx, 0])
if tid < 0:
continue
slices = [slice(None)] * tiles_src[idx].ndim
slices_list = [slice(None)] * tiles_src[idx].ndim
for i, dim in enumerate(grid_spec.split_dims):
slices[dim] = slice(0, int(meta_src[idx, i + 1]))
slices = tuple(slices)
coord_tensor_map[tid_coord_map[tid]] = tiles_src[idx][slices]
slices_list[dim] = slice(0, int(meta_src[idx, i + 1]))
coord_tensor_map[tid_coord_map[tid]] = tiles_src[idx][tuple(slices_list)]
return coord_tensor_map
def _balance_tasks(self, task_list, num_rank):
workloads = [0] * num_rank
assigned = [[] for _ in range(num_rank)]
assigned: list[list[TileTask]] = [[] for _ in range(num_rank)]
for task in sorted(task_list, key=lambda t: t.workload, reverse=True):
r = workloads.index(min(workloads))
@@ -122,49 +154,51 @@ class DistributedVaeExecutor:
def execute(self, z: torch.Tensor, operator: DistributedOperator, broadcast_result: bool = True):
pp_size = min(self.parallel_size, self.world_size)
# 1. Split into tiles
tiletask_list, grid_spec = operator.split(z)
tid_coord_map = {task.tile_id: task.grid_coord for task in tiletask_list}
# 2. local decode
assigned = self._balance_tasks(tiletask_list, pp_size)
local_tasks = assigned[self.rank] if self.rank < pp_size else []
local_results = [(t.tile_id, operator.exec(t)) for t in local_tasks]
with self._sync_errors("tile decode"):
tiletask_list, grid_spec = operator.split(z)
tid_coord_map = {task.tile_id: task.grid_coord for task in tiletask_list}
assigned = self._balance_tasks(tiletask_list, pp_size)
local_tasks = assigned[self.rank] if self.rank < pp_size else []
local_results = [(t.tile_id, operator.exec(t)) for t in local_tasks]
# 3. compute shape per rank
global_padding_shape = self._compute_global_padding_shape(local_results, z.ndim, z.device)
# 5. prepare tile tensors
output_dtype = grid_spec.output_dtype if grid_spec.output_dtype is not None else z.dtype
local_tile_tensor, local_meta_tensor = self._pack_local_tiles(
local_results, global_padding_shape, grid_spec, z.device, output_dtype
)
with self._sync_errors("tile packing"):
local_tile_tensor, local_meta_tensor = self._pack_local_tiles(
local_results, global_padding_shape, grid_spec, z.device, output_dtype
)
# 6. gather tiles & meta
meta_gather = self.gather_tensors(local_meta_tensor)
tile_gather = self.gather_tensors(local_tile_tensor)
if self.rank != 0:
result = torch.empty(0, device=z.device) # Dummy return for non-zero ranks
else:
# 7. reconstruct full tensor (rank 0)
coord_tensor_map = self._unpack_tiles(meta_gather, tile_gather, grid_spec, tid_coord_map)
result = operator.merge(coord_tensor_map, grid_spec)
with self._sync_errors("tile merge"):
if self.rank != 0:
result = torch.empty(0, device=z.device) # Dummy return for non-zero ranks
else:
# 7. reconstruct full tensor (rank 0)
coord_tensor_map = self._unpack_tiles(meta_gather, tile_gather, grid_spec, tid_coord_map)
result = operator.merge(coord_tensor_map, grid_spec)
if broadcast_result:
result = self._sync_final_result(result, z.ndim, z.device, output_dtype)
return result
def _sync_final_result(self, rank0_result, output_ndim, output_device, output_dtype):
shape_tensor = torch.empty((output_ndim,), device=output_device, dtype=torch.int64)
if self.rank == 0:
shape_tensor.copy_(torch.tensor(tuple(rank0_result.shape), device=output_device, dtype=torch.int64))
with self._sync_errors("output shape preparation"):
shape_tensor = torch.empty((output_ndim,), device=output_device, dtype=torch.int64)
if self.rank == 0:
shape_tensor.copy_(torch.tensor(tuple(rank0_result.shape), device=output_device, dtype=torch.int64))
dist.broadcast(shape_tensor, src=0, group=self.group)
if self.rank != 0:
sync_result = torch.empty(tuple(shape_tensor.tolist()), device=output_device, dtype=output_dtype)
else:
sync_result = rank0_result
with self._sync_errors("output allocation"):
if self.rank != 0:
sync_result = torch.empty(tuple(shape_tensor.tolist()), device=output_device, dtype=output_dtype)
else:
sync_result = rank0_result
dist.broadcast(sync_result, src=0, group=self.group)
return sync_result