diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index a45b5348e..fd9ccaa5c 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -44,4 +44,8 @@ th { |--------------|--------|-------------------| | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B`| | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | +| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` | +| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` | +| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` | +| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | |`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | diff --git a/vllm_omni/diffusion/attention/backends/ascend_attn.py b/vllm_omni/diffusion/attention/backends/ascend_attn.py new file mode 100644 index 000000000..6563385c9 --- /dev/null +++ b/vllm_omni/diffusion/attention/backends/ascend_attn.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from importlib.util import find_spec + +import torch +from vllm.logger import init_logger + +from vllm_omni.diffusion.attention.backends.abstract import ( + AttentionBackend, + AttentionImpl, + AttentionMetadata, +) + +logger = init_logger(__name__) + + +class AscendAttentionBackend(AttentionBackend): + accept_output_buffer: bool = True + + @staticmethod + def get_supported_head_sizes() -> list[int]: + return [x for x in range(1024)] # todo + + @staticmethod + def get_name() -> str: + return "ASCEND" + + @staticmethod + def get_impl_cls() -> type["AscendAttentionBackendImpl"]: + return AscendAttentionBackendImpl + + +class AscendAttentionBackendImpl(AttentionImpl): + def __init__( + self, + num_heads: int, + head_size: int, + softmax_scale: float, + causal: bool = False, + num_kv_heads: int | None = None, + prefix: str = "", + **extra_impl_args, + ) -> None: + self.causal = causal + self.softmax_scale = softmax_scale + + def forward_native( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata = None, + ) -> torch.Tensor: + query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value)) + attention_mask = attn_metadata.attn_mask if attn_metadata else None + + output = torch.nn.functional.scaled_dot_product_attention( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=self.causal, + scale=self.softmax_scale, + ) + out = output.permute(0, 2, 1, 3) + return out + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata = None, + ) -> torch.Tensor: + if find_spec("mindiesd"): + from mindiesd import attention_forward + + attention_mask = attn_metadata.attn_mask if attn_metadata else None + output = attention_forward( + query, + key, + value, + attn_mask=attention_mask, + opt_mode="manual", + op_type="fused_attn_score", + layout="BNSD", + ) + return output + else: + return self.forward_native(query, key, value, attn_metadata) diff --git a/vllm_omni/diffusion/attention/selector.py b/vllm_omni/diffusion/attention/selector.py index fe5a91e36..f2f60243f 100644 --- a/vllm_omni/diffusion/attention/selector.py +++ b/vllm_omni/diffusion/attention/selector.py @@ -28,6 +28,7 @@ _BACKEND_CONFIG = { "module": "vllm_omni.diffusion.attention.backends.sage_attn", "class": "SageAttentionBackend", }, + "ASCEND": {"module": "vllm_omni.diffusion.attention.backends.ascend_attn", "class": "AscendAttentionBackend"}, } diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index 56b4deefb..d3002ba57 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/vllm_omni/diffusion/cache/teacache/extractors.py @@ -208,8 +208,7 @@ def extract_qwen_context( block = module.transformer_blocks[0] img_mod_params = block.img_mod(temb) img_mod1, _ = img_mod_params.chunk(2, dim=-1) - img_normed = block.img_norm1(hidden_states) - img_modulated, _ = block._modulate(img_normed, img_mod1) + img_modulated, _ = block.img_norm1(hidden_states, img_mod1) # ============================================================================ # DEFINE TRANSFORMER EXECUTION (Qwen-specific) diff --git a/vllm_omni/diffusion/layers/adalayernorm.py b/vllm_omni/diffusion/layers/adalayernorm.py new file mode 100644 index 000000000..835ee07c0 --- /dev/null +++ b/vllm_omni/diffusion/layers/adalayernorm.py @@ -0,0 +1,102 @@ +import torch +import torch.nn as nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.layers.custom_op import CustomOp + +logger = init_logger(__name__) + + +class AdaLayerNorm(CustomOp): + """ + AdaLayerNorm: + out = layernorm(x) * (1 + scale) + shift + """ + + def __init__(self, hidden_size: int, elementwise_affine: bool = False, eps: float = 1e-6) -> None: + super().__init__() + self.eps = eps + self.elementwise_affine = elementwise_affine + self.hidden_size = hidden_size + self.layernorm = nn.LayerNorm(self.hidden_size, elementwise_affine=self.elementwise_affine, eps=self.eps) + + def preprocess( + self, + mod_params: torch.Tensor, + index: torch.Tensor = None, + ) -> torch.Tensor: + # shift: b d, scale: b d, gate: b d + shift, scale, gate = mod_params.chunk(3, dim=-1) + + if index is not None: + # Assuming mod_params batch dim is 2*actual_batch (chunked into 2 parts) + # So shift, scale, gate have shape [2*actual_batch, d] + actual_batch = shift.size(0) // 2 + shift_0, shift_1 = shift[:actual_batch], shift[actual_batch:] # each: [actual_batch, d] + scale_0, scale_1 = scale[:actual_batch], scale[actual_batch:] + gate_0, gate_1 = gate[:actual_batch], gate[actual_batch:] + + # index: [b, l] where b is actual batch size + # Expand to [b, l, 1] to match feature dimension + index_expanded = index.unsqueeze(-1) # [b, l, 1] + + # Expand chunks to [b, 1, d] then broadcast to [b, l, d] + shift_0_exp = shift_0.unsqueeze(1) # [b, 1, d] + shift_1_exp = shift_1.unsqueeze(1) # [b, 1, d] + scale_0_exp = scale_0.unsqueeze(1) + scale_1_exp = scale_1.unsqueeze(1) + gate_0_exp = gate_0.unsqueeze(1) + gate_1_exp = gate_1.unsqueeze(1) + + # Use torch.where to select based on index + shift_result = torch.where(index_expanded == 0, shift_0_exp, shift_1_exp) + scale_result = torch.where(index_expanded == 0, scale_0_exp, scale_1_exp) + gate_result = torch.where(index_expanded == 0, gate_0_exp, gate_1_exp) + else: + shift_result = shift.unsqueeze(1) + scale_result = scale.unsqueeze(1) + gate_result = gate.unsqueeze(1) + + return shift_result, scale_result, gate_result + + def forward_cuda( + self, + x: torch.Tensor, + mod_params: torch.Tensor, + index: torch.Tensor = None, + ) -> torch.Tensor: + return self.forward_native(x, mod_params, index) + + def forward_hip( + self, + x: torch.Tensor, + mod_params: torch.Tensor, + index: torch.Tensor = None, + ) -> torch.Tensor: + return self.forward_native(x, mod_params, index) + + def forward_npu( + self, + x: torch.Tensor, + mod_params: torch.Tensor, + index: torch.Tensor = None, + ) -> torch.Tensor: + shift_result, scale_result, gate_result = self.preprocess(mod_params, index) + + import torch_npu + + output = torch_npu.npu_layer_norm_eval( + x, normalized_shape=[self.hidden_size], weight=(1 + scale_result), bias=shift_result, eps=self.eps + ) + + return output, gate_result + + def forward_native( + self, + x: torch.Tensor, + mod_params: torch.Tensor, + index: torch.Tensor = None, + ) -> torch.Tensor: + shift_result, scale_result, gate_result = self.preprocess(mod_params, index) + + return self.layernorm(x) * (1 + scale_result) + shift_result, gate_result diff --git a/vllm_omni/diffusion/layers/custom_op.py b/vllm_omni/diffusion/layers/custom_op.py index 0bf5c4f60..2f50073cc 100644 --- a/vllm_omni/diffusion/layers/custom_op.py +++ b/vllm_omni/diffusion/layers/custom_op.py @@ -3,7 +3,7 @@ from typing import Any import torch.nn as nn -from vllm_omni.utils.platform_utils import detect_device_type, is_rocm +from vllm_omni.utils.platform_utils import detect_device_type, is_npu, is_rocm class CustomOp(nn.Module): @@ -22,6 +22,8 @@ class CustomOp(nn.Module): return self.forward_hip elif self.is_cuda: return self.forward_cuda + elif is_npu(): + return self.forward_npu else: return self.forward_native @@ -39,6 +41,9 @@ class CustomOp(nn.Module): def forward_cuda(self, *args, **kwargs): raise NotImplementedError + def forward_npu(self, *args, **kwargs): + raise NotImplementedError + def forward_hip(self, *args, **kwargs): # By default, we assume that HIP ops are compatible with CUDA ops. return self.forward_cuda(*args, **kwargs) diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py index 528f2425e..81e8c098f 100644 --- a/vllm_omni/diffusion/layers/rope.py +++ b/vllm_omni/diffusion/layers/rope.py @@ -36,6 +36,35 @@ def apply_rotary_emb_torch(x, cos, sin, interleaved=False): ) +def apply_rotary_emb_mindiesd( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + interleaved: bool = False, + half_head_dim: bool = True, # if true, size of sin and cos is (B, S, D/2), otherwise (B, S, D) +) -> torch.Tensor: + from mindiesd import rotary_position_embedding + + if cos.dim() == 3: + # (B, S, D/2) -> (S, D/2) + cos = cos[0] + sin = sin[0] + + if interleaved: + # if last dim of sin and cos is D/2, expand to (S, D) to adapt to mindiesd operators + if half_head_dim: + seqlen = cos.shape[0] + sin = sin.unsqueeze(0).unsqueeze(2).unsqueeze(-1).expand(-1, -1, -1, -1, 2).reshape(1, seqlen, 1, -1) + cos = cos.unsqueeze(0).unsqueeze(2).unsqueeze(-1).expand(-1, -1, -1, -1, 2).reshape(1, seqlen, 1, -1) + return rotary_position_embedding(x, cos, sin, rotated_mode="rotated_interleaved", head_first=False, fused=True) + else: + if half_head_dim: + seqlen = cos.shape[0] + sin = sin.unsqueeze(0).unsqueeze(2).repeat(1, 1, 1, 2) + cos = cos.unsqueeze(0).unsqueeze(2).repeat(1, 1, 1, 2) + return rotary_position_embedding(x, cos, sin, rotated_mode="rotated_half", head_first=False, fused=True) + + class RotaryEmbedding(CustomOp): """ rotary positional embedding. @@ -97,6 +126,17 @@ class RotaryEmbedding(CustomOp): interleaved=self.interleaved, ) + def forward_npu( + self, + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + ) -> torch.Tensor: + if find_spec("mindiesd"): + return apply_rotary_emb_mindiesd(x, cos, sin, self.interleaved) + else: + return self.forward_native(x, cos, sin) + def forward_native( self, x: torch.Tensor, diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py index e6d10b087..1a52769ff 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py @@ -31,6 +31,7 @@ from vllm_omni.diffusion.distributed.parallel_state import ( get_sp_group, ) from vllm_omni.diffusion.forward_context import get_forward_context +from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm from vllm_omni.diffusion.layers.rope import RotaryEmbedding logger = init_logger(__name__) @@ -478,7 +479,7 @@ class QwenImageTransformerBlock(nn.Module): nn.SiLU(), nn.Linear(dim, 6 * dim, bias=True), # For scale, shift, gate for norm1 and norm2 ) - self.img_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.img_norm1 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps) self.attn = QwenImageCrossAttention( dim=dim, num_heads=num_attention_heads, @@ -486,7 +487,7 @@ class QwenImageTransformerBlock(nn.Module): context_pre_only=False, head_dim=attention_head_dim, ) - self.img_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.img_norm2 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps) self.img_mlp = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate") # Text processing modules @@ -494,9 +495,9 @@ class QwenImageTransformerBlock(nn.Module): nn.SiLU(), nn.Linear(dim, 6 * dim, bias=True), # For scale, shift, gate for norm1 and norm2 ) - self.txt_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.txt_norm1 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps) # Text doesn't need separate attention - it's handled by img_attn joint computation - self.txt_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.txt_norm2 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps) self.txt_mlp = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate") self.zero_cond_t = zero_cond_t @@ -560,12 +561,10 @@ class QwenImageTransformerBlock(nn.Module): txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1) # Each [B, 3*dim] # Process image stream - norm1 + modulation - img_normed = self.img_norm1(hidden_states) - img_modulated, img_gate1 = self._modulate(img_normed, img_mod1, modulate_index) + img_modulated, img_gate1 = self.img_norm1(hidden_states, img_mod1) # Process text stream - norm1 + modulation - txt_normed = self.txt_norm1(encoder_hidden_states) - txt_modulated, txt_gate1 = self._modulate(txt_normed, txt_mod1) + txt_modulated, txt_gate1 = self.txt_norm1(encoder_hidden_states, txt_mod1) # Use QwenAttnProcessor2_0 for joint attention computation # This directly implements the DoubleStreamLayerMegatron logic: @@ -588,14 +587,12 @@ class QwenImageTransformerBlock(nn.Module): encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn_output # Process image stream - norm2 + MLP - img_normed2 = self.img_norm2(hidden_states) - img_modulated2, img_gate2 = self._modulate(img_normed2, img_mod2, modulate_index) + img_modulated2, img_gate2 = self.img_norm2(hidden_states, img_mod2) img_mlp_output = self.img_mlp(img_modulated2) hidden_states = hidden_states + img_gate2 * img_mlp_output # Process text stream - norm2 + MLP - txt_normed2 = self.txt_norm2(encoder_hidden_states) - txt_modulated2, txt_gate2 = self._modulate(txt_normed2, txt_mod2) + txt_modulated2, txt_gate2 = self.txt_norm2(encoder_hidden_states, txt_mod2) txt_mlp_output = self.txt_mlp(txt_modulated2) encoder_hidden_states = encoder_hidden_states + txt_gate2 * txt_mlp_output