[Model] Support Qwen-Image series models on NPU (#463)

Signed-off-by: KexiongYu <yukexiong1@huawei.com>
Signed-off-by: betta18 <jiangmengyu1@huawei.com>
Signed-off-by: muziyuhui666 <111362884+muziyuhui666@users.noreply.github.com>
Signed-off-by: jiangmengyu18 <56633611+jiangmengyu18@users.noreply.github.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: jiangmengyu18 <56633611+jiangmengyu18@users.noreply.github.com>
Co-authored-by: KexiongYu <yukexiong1@huawei.com>
Co-authored-by: betta18 <jiangmengyu1@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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` |

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)

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"},
 }
 
 

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)

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

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)

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,

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