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

[Feat] support HSDP for Qwen-image series, Z-Image, GLM-Image (#2029)

Browse Source 
Signed-off-by: Lancer <maruixiang6688@gmail.com>
This commit is contained in:
Lancer authored and GitHub committed 2026-04-01 20:01:09 +08:00
1 parent 4e4bbc42a6
commit 70a62651b9
15 files changed
+83 -17
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/user_guide/diffusion_features.md
+7 -7
View File
@@ -105,7 +105,7 @@ The following tables show which models support each feature:
| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **FLUX.2-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
| **LongCat-Image** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **LongCat-Image-Edit** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
@@ -113,13 +113,13 @@ The following tables show which models support each feature:
| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **OmniGen2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ |
| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ |
| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ |
| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ |
| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ |
| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ❌ | ✅ | ✅ | ❌ |
| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ✅ | ❌ | ✅ | ✅ | ❌ |
> Notes:
> 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
tests/e2e/online_serving/test_qwen_image_edit_expansion.py
+12
View File
@@ -98,6 +98,18 @@ def _get_diffusion_feature_cases(model: str):
id="parallel_004",
marks=PARALLEL_FEATURE_MARKS,
),
pytest.param(
OmniServerParams(
model=model,
server_args=[
"--use-hsdp",
"--hsdp-shard-size",
"2",
],
),
id="parallel_005",
marks=PARALLEL_FEATURE_MARKS,
),
]
tests/e2e/online_serving/test_qwen_image_expansion.py
+12
View File
@@ -107,6 +107,18 @@ def _get_diffusion_feature_cases(model: str):
id="vae_patch_parallel_2",
marks=PARALLEL_FEATURE_MARKS,
),
pytest.param(
OmniServerParams(
model=model,
server_args=[
"--use-hsdp",
"--hsdp-shard-size",
"2",
],
),
id="parallel_hsdp",
marks=PARALLEL_FEATURE_MARKS,
),
]
tests/e2e/online_serving/test_qwen_image_layered_expansion.py
+12
View File
@@ -62,6 +62,18 @@ FEATURE_CASES = [
id="cfg_parallel_001",
marks=PARALLEL_FEATURE_MARKS,
),
pytest.param(
OmniServerParams(
model=MODEL,
server_args=[
"--use-hsdp",
"--hsdp-shard-size",
"2",
],
),
id="parallel_hsdp",
marks=PARALLEL_FEATURE_MARKS,
),
]
tests/e2e/online_serving/test_zimage_expansion.py
+12
View File
@@ -60,6 +60,18 @@ def _get_diffusion_feature_cases():
id="parallel_teacache_fp8_ulysses2_ring2",
marks=FOUR_CARD_MARKS,
),
pytest.param(
OmniServerParams(
model=MODEL,
server_args=[
"--use-hsdp",
"--hsdp-shard-size",
"2",
],
),
id="parallel_hsdp",
marks=FOUR_CARD_MARKS,
),
]
vllm_omni/diffusion/distributed/hsdp_utils.py
+9
View File
@@ -0,0 +1,9 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import Any
def is_transformer_block_module(name: str, module: Any) -> bool:
"""Return True for numbered modules under `transformer_blocks`."""
return "transformer_blocks" in name and name.split(".")[-1].isdigit()
vllm_omni/diffusion/models/glm_image/glm_image_transformer.py
+3
View File
@@ -22,6 +22,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm_omni.diffusion.attention.layer import Attention
from vllm_omni.diffusion.cache.base import CachedTransformer
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module
logger = init_logger(__name__)
@@ -724,6 +725,8 @@ class GlmImageTransformer2DModel(CachedTransformer):
_repeated_blocks = ["GlmImageTransformerBlock"]
_hsdp_shard_conditions = [is_transformer_block_module]
def __init__(
self,
od_config: OmniDiffusionConfig,
vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py
+2 -5
View File
@@ -23,6 +23,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
from vllm_omni.diffusion.attention.layer import Attention
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module
from vllm_omni.diffusion.layers.rope import RotaryEmbedding
from vllm_omni.diffusion.models.flux.flux_transformer import FeedForward
@@ -544,11 +545,7 @@ class HunyuanVideo15Transformer3DModel(nn.Module):
"add_kv_proj": ["add_q_proj", "add_k_proj", "add_v_proj"],
}
@staticmethod
def _is_transformer_block(name: str, module) -> bool:
return "transformer_blocks" in name and name.split(".")[-1].isdigit()
_hsdp_shard_conditions = [_is_transformer_block]
_hsdp_shard_conditions = [is_transformer_block_module]
def __init__(
self,
vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py
+1 -1
View File
@@ -273,7 +273,7 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
)
self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
model, subfolder="text_encoder", local_files_only=local_files_only
)
).to(self.device)
self.vae = DistributedAutoencoderKLQwenImage.from_pretrained(
model, subfolder="vae", local_files_only=local_files_only
).to(self.device)
vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py
+1 -1
View File
@@ -245,7 +245,7 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
)
self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
model, subfolder="text_encoder", local_files_only=local_files_only
)
).to(self.device)
self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
self.device
vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py
+1 -1
View File
@@ -198,7 +198,7 @@ class QwenImageEditPlusPipeline(
)
self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
model, subfolder="text_encoder", local_files_only=local_files_only
)
).to(self.device)
self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
self.device
vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py
+1 -1
View File
@@ -219,7 +219,7 @@ class QwenImageLayeredPipeline(nn.Module, SupportImageInput, QwenImageCFGParalle
)
self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
model, subfolder="text_encoder", local_files_only=local_files_only
)
).to(self.device)
self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
self.device
)
vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py
+3
View File
@@ -37,6 +37,7 @@ from vllm_omni.diffusion.attention.backends.abstract import (
from vllm_omni.diffusion.attention.layer import Attention
from vllm_omni.diffusion.cache.base import CachedTransformer
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module
from vllm_omni.diffusion.distributed.sp_plan import (
SequenceParallelInput,
SequenceParallelOutput,
@@ -887,6 +888,8 @@ class QwenImageTransformer2DModel(CachedTransformer):
"add_kv_proj": ["add_q_proj", "add_k_proj", "add_v_proj"],
}
_hsdp_shard_conditions = [is_transformer_block_module]
# Sequence Parallelism plan (following diffusers' _cp_plan pattern)
# Similar to Z-Image's UnifiedPrepare, we use ImageRopePrepare to create
# a module boundary where _sp_plan can shard hidden_states and vid_freqs together.
vllm_omni/diffusion/models/z_image/pipeline_z_image.py
+1 -1
View File
@@ -170,7 +170,7 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
self.text_encoder = AutoModel.from_pretrained(
model, subfolder="text_encoder", local_files_only=local_files_only
)
).to(self._execution_device)
self.vae = DistributedAutoencoderKL.from_pretrained(
model, subfolder="vae", local_files_only=local_files_only
).to(self._execution_device)
vllm_omni/diffusion/models/z_image/z_image_transformer.py
+6
View File
@@ -580,6 +580,12 @@ class ZImageTransformer2DModel(CachedTransformer):
_repeated_blocks = ["ZImageTransformerBlock"]
@staticmethod
def _is_transformer_block(name: str, module) -> bool:
return "layers" in name and name.split(".")[-1].isdigit()
_hsdp_shard_conditions = [_is_transformer_block]
# Sequence Parallelism for Z-Image (following diffusers' _cp_plan pattern)
# Similar to how Wan uses `rope` module's split_output to shard rotary embeddings,
# Z-Image uses `unified_prepare` module's split_output to shard unified tensors.