[Refactor][Diffusion] Move single-model layers into their model directories (#5908)

Signed-off-by: congw729 <115451386+congw729@users.noreply.github.com>
Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

benchmarks/kernels/mot_linear_benchmarks.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # ruff: noqa: N803
 
 """MoT (Mixture-of-Tokens) GEMM kernel benchmark and auto-tuning.
@@ -10,7 +13,7 @@ Usage:
     python benchmarks/kernels/mot_linear_benchmarks.py \
         --model ByteDance-Seed/BAGEL-7B-MoT \
         --tp-size 1 --dtype w16a16 --tune \
-        --save-dir vllm_omni/diffusion/layers/mot/configs/
+        --save-dir vllm_omni/diffusion/models/bagel/mot/configs/
 
     # Auto-tune with local model path (offline clusters):
     python benchmarks/kernels/mot_linear_benchmarks.py \
@@ -43,7 +46,7 @@ from vllm.triton_utils import triton
 from vllm.utils.torch_utils import set_random_seed
 
 # NOTE: you should use the same naming syetem for the kernel to load properly
-from vllm_omni.diffusion.layers.mot.ops.mot_gemm import build_config_filename, get_device_name
+from vllm_omni.diffusion.models.bagel.mot.ops.mot_gemm import build_config_filename, get_device_name
 
 # clear the triton cache from time to time, usaully no need to change
 _CACHE_CLEAR_INTERVAL_ENV = "VLLM_MOT_TUNE_CACHE_CLEAR_INTERVAL"
@@ -428,7 +431,7 @@ def benchmark_config(
     cache_flusher: torch.Tensor | None = None,
 ) -> float:
     """Run a MoT GEMM with the given tile config and return avg latency (us)."""
-    from vllm_omni.diffusion.layers.mot.ops.mot_gemm import invoke_mot_gemm
+    from vllm_omni.diffusion.models.bagel.mot.ops.mot_gemm import invoke_mot_gemm
 
     text_indices, vae_indices, M = build_regular_indices(
         image_num=image_num,
@@ -553,7 +556,7 @@ class BenchmarkWorker:
         use_int8_w8a16: bool,
     ) -> tuple[dict[str, int], float]:
         set_random_seed(self.seed)
-        from vllm_omni.diffusion.layers.mot.ops.mot_gemm import (
+        from vllm_omni.diffusion.models.bagel.mot.ops.mot_gemm import (
             get_best_mot_config,
         )
 

docs/user_guide/diffusion/mot_config.md

@@ -28,7 +28,7 @@ and their optimal Triton tile configurations:
 
 Config loading order (3-tier, see ops/mot_gemm.py):
     1. $VLLM_TUNED_CONFIG_FOLDER/<filename>  (env override)
-    2. This directory: vllm_omni/diffusion/layers/mot/configs/<filename>
+    2. This directory: vllm_omni/diffusion/models/bagel/mot/configs/<filename>
     3. Conservative default config (compiles everywhere, sub-optimal perf)
 
 If no config file matches the current device, a warning is printed with
@@ -39,7 +39,7 @@ To generate configs for your hardware:
     python benchmarks/kernels/mot_linear_benchmarks.py \
         --model ByteDance-Seed/BAGEL-7B-MoT \
         --tp-size 1 --dtype w16a16 --tune \
-        --save-dir vllm_omni/diffusion/layers/mot/configs/
+        --save-dir vllm_omni/diffusion/models/bagel/mot/configs/
 
 For multi-GPU tuning (uses Ray for parallel search):
 

tests/diffusion/kernels/mot/test_mot_linear.py

@@ -41,13 +41,13 @@ from vllm_omni.diffusion.distributed.parallel_state import (
     initialize_model_parallel,
     model_parallel_is_initialized,
 )
-from vllm_omni.diffusion.layers.mot.mot_qkv_parallel_linear import (
+from vllm_omni.diffusion.models.bagel.mot.mot_qkv_parallel_linear import (
     MoTQKVParallelLinear,
 )
-from vllm_omni.diffusion.layers.mot.mot_row_parallel_linear import (
+from vllm_omni.diffusion.models.bagel.mot.mot_row_parallel_linear import (
     MoTRowParallelLinear,
 )
-from vllm_omni.diffusion.layers.mot.ops.mot_gemm import (
+from vllm_omni.diffusion.models.bagel.mot.ops.mot_gemm import (
     get_best_mot_config,
 )
 

tests/diffusion/kernels/mot/test_mot_norm.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # ruff: noqa: N803, E741
 """Layer-level correctness & performance test for MoTRMSNorm.
 
@@ -20,7 +23,7 @@ import torch
 from vllm import ir
 from vllm.config import VllmConfig, set_current_vllm_config
 
-from vllm_omni.diffusion.layers.mot.mot_layernorm import MoTRMSNorm
+from vllm_omni.diffusion.models.bagel.mot.mot_layernorm import MoTRMSNorm
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
 

tests/diffusion/kernels/mot/test_mot_norm_npu.py

@@ -1,8 +1,11 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import pytest
 import torch
 
 import vllm_omni.diffusion.layers.custom_op as custom_op
-from vllm_omni.diffusion.layers.mot.mot_layernorm import MoTRMSNorm
+from vllm_omni.diffusion.models.bagel.mot.mot_layernorm import MoTRMSNorm
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 

tests/diffusion/layers/test_custom_op.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for the CustomOp platform-dispatch base class."""
 
 import pytest
@@ -39,7 +39,7 @@ def test_forward_cuda_still_requires_an_implementation():
 
 def test_mot_rmsnorm_dispatches_on_xpu_without_an_override():
     """MoTRMSNorm relies on the inherited fallback (regression test)."""
-    from vllm_omni.diffusion.layers.mot.mot_layernorm import MoTRMSNorm
+    from vllm_omni.diffusion.models.bagel.mot.mot_layernorm import MoTRMSNorm
 
     assert "forward_xpu" not in MoTRMSNorm.__dict__
 

tests/diffusion/models/wan2_2/test_wan22_quant_config_propagation.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for Wan2.2 quant_config propagation through transformer creation.
 
 Tests cover:
@@ -196,7 +196,7 @@ class TestPatchWanRmsNorm:
     """Test that patch_wan_rms_norm doesn't raise on concurrent module registration."""
 
     def test_patches_modules_with_wan_rms_norm(self):
-        from vllm_omni.diffusion.layers.norm import RMSNormVAE
+        from vllm_omni.diffusion.models.wan2_2.norm import RMSNormVAE
         from vllm_omni.diffusion.models.wan2_2.patch_diffusers import patch_wan_rms_norm
 
         # Create a fake module that has WanRMS_norm

tools/pre_commit/check_forbidden_imports.py

@@ -108,7 +108,7 @@ CHECK_IMPORTS = {
             "setup.py",
             "vllm_omni/benchmarks/data_modules/daily_omni_eval.py",
             "vllm_omni/config/stage_config.py",
-            "vllm_omni/diffusion/layers/mot/ops/mot_gemm.py",
+            "vllm_omni/diffusion/models/bagel/mot/ops/mot_gemm.py",
             "vllm_omni/diffusion/model_loader/diffusers_loader.py",
             "vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py",
             "vllm_omni/diffusion/models/dreamid_omni/fusion.py",

vllm_omni/diffusion/layers/norm.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from importlib.util import find_spec
 
 import torch
@@ -163,75 +166,3 @@ class RMSNorm(CustomOp):
         out = self.weight.to(torch.float32) * out
         out = out.to(input_dtype)
         return (out, residual) if residual is not None else out
-
-
-class RMSNormVAE(CustomOp):
-    """Root Mean Square Layer Normalization for Channel-First or Last"""
-
-    def __init__(
-        self,
-        dim: int,
-        channel_first: bool = True,
-        images: bool = True,
-        bias: bool = False,
-        epsilon: float = 1e-6,
-    ) -> None:
-        super().__init__()
-        broadcastable_dims = (1, 1, 1) if not images else (1, 1)
-        shape = (dim, *broadcastable_dims) if channel_first else (dim,)
-
-        self.channel_first = channel_first
-        self.scale = dim**0.5
-        self.gamma = nn.Parameter(torch.ones(shape))
-        self.bias = nn.Parameter(torch.zeros(shape)) if bias else None
-        self.epsilon = epsilon
-
-        self.gamma_rmsnorm = None
-
-    def forward_cuda(
-        self,
-        x: torch.Tensor,
-    ) -> torch.Tensor:
-        return self.forward_native(x)
-
-    def forward_hip(
-        self,
-        x: torch.Tensor,
-    ) -> torch.Tensor:
-        return self.forward_native(x)
-
-    def forward_npu(
-        self,
-        x: torch.Tensor,
-    ) -> torch.Tensor:
-        import torch_npu
-
-        if self.gamma_rmsnorm is None:
-            self.gamma_rmsnorm = self.gamma.reshape(-1)
-
-        if self.channel_first:
-            x = x.transpose(1, -1)
-            out = torch_npu.npu_rms_norm(x, self.gamma_rmsnorm, epsilon=self.epsilon)[0].transpose(1, -1)
-        else:
-            out = torch_npu.npu_rms_norm(x, self.gamma_rmsnorm, epsilon=self.epsilon)[0]
-
-        if self.bias is not None:
-            out = out + self.bias
-        return out
-
-    def forward_native(
-        self,
-        x: torch.Tensor,
-    ) -> torch.Tensor:
-        out = (
-            F.normalize(
-                x,
-                dim=(1 if self.channel_first else -1),
-                eps=self.epsilon,
-            )
-            * self.scale
-            * self.gamma
-        )
-        if self.bias is not None:
-            out = out + self.bias
-        return out

vllm_omni/diffusion/layers/rope.py

@@ -344,140 +344,6 @@ def apply_rope_to_qk(
     return query, key
 
 
-class WanS2VRotaryPosEmbed(torch.nn.Module):
-    """Precompute complex-valued RoPE embeddings for S2V multi-grid positions.
-
-    Owns the base frequency buffer and provides forward() to compute position
-    embeddings given hidden_states (for shape) and grid_sizes.
-    """
-
-    def __init__(self, num_heads: int, head_dim: int, max_seq_len: int = 1024):
-        super().__init__()
-        self.num_heads = num_heads
-        self.head_dim = head_dim
-        d = head_dim
-        freqs = torch.cat(
-            [
-                self._rope_params(max_seq_len, d - 4 * (d // 6)),
-                self._rope_params(max_seq_len, 2 * (d // 6)),
-                self._rope_params(max_seq_len, 2 * (d // 6)),
-            ],
-            dim=1,
-        )
-        self.register_buffer("freqs", freqs.to(torch.complex64), persistent=False)
-
-    @staticmethod
-    @torch.amp.autocast(current_omni_platform.device_type, enabled=False)
-    def _rope_params(max_seq_len, dim, theta=10000):
-        if dim % 2 != 0:
-            raise ValueError(f"dim ({dim}) must be even")
-        freqs = torch.outer(
-            torch.arange(max_seq_len), 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float64).div(dim))
-        )
-        return torch.polar(torch.ones_like(freqs), freqs)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        grid_sizes: list,
-        trainable_freqs: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        """Precompute RoPE embeddings for the given grid layout.
-
-        Args:
-            hidden_states: Tensor [B, S, ...] (used for batch/seq shape and device)
-            grid_sizes: Grid specification (list of [offsets, sizes, totals])
-            trainable_freqs: Optional trainable frequency overrides for t_f < 0
-
-        Returns:
-            Complex tensor [B, S, 1, head_dim//2] of precomputed position embeddings
-        """
-        b, s = hidden_states.shape[0], hidden_states.shape[1]
-        c = self.head_dim // 2
-        device = hidden_states.device
-
-        freqs = self.freqs.to(device)
-        if trainable_freqs is not None:
-            freqs_input = [freqs, trainable_freqs]
-        else:
-            freqs_input = freqs
-
-        if isinstance(freqs_input, list):
-            trainable_f = freqs_input[1]
-            freqs_split = freqs_input[0]
-        else:
-            trainable_f = None
-            freqs_split = freqs_input
-        freqs_split = freqs_split.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
-
-        output = torch.empty((b, s, 1, c), device=device, dtype=torch.complex64)
-        seq_bucket = [0]
-        if not isinstance(grid_sizes, list):
-            grid_sizes = [grid_sizes]
-        for g in grid_sizes:
-            if not isinstance(g, list):
-                g = [torch.zeros_like(g), g]
-            batch_size = g[0].shape[0]
-            for i in range(batch_size):
-                f_o, h_o, w_o = g[0][i]
-                f, h, w = g[1][i]
-                t_f, t_h, t_w = g[2][i]
-                seq_f, seq_h, seq_w = f - f_o, h - h_o, w - w_o
-                seq_len = int(seq_f * seq_h * seq_w)
-                if seq_len > 0:
-                    if t_f > 0:
-                        assert f_o * f >= 0 and h_o * h >= 0 and w_o * w >= 0
-                        seq_f_int = int(seq_f)
-                        seq_h_int = int(seq_h)
-                        seq_w_int = int(seq_w)
-
-                        if f_o >= 0:
-                            f_sam = torch.linspace(int(f_o), int(t_f + f_o) - 1, seq_f_int, device=device).long()
-                        else:
-                            f_sam = torch.linspace(int(-f_o), int(-t_f - f_o) + 1, seq_f_int, device=device).long()
-                        h_sam = torch.linspace(int(h_o), int(t_h + h_o) - 1, seq_h_int, device=device).long()
-                        w_sam = torch.linspace(int(w_o), int(t_w + w_o) - 1, seq_w_int, device=device).long()
-
-                        freqs_0 = torch.index_select(freqs_split[0] if f_o >= 0 else freqs_split[0].conj(), 0, f_sam)
-                        freqs_0 = freqs_0.view(seq_f_int, 1, 1, -1)
-
-                        freqs_i = torch.cat(
-                            [
-                                freqs_0.expand(seq_f_int, seq_h_int, seq_w_int, -1),
-                                torch.index_select(freqs_split[1], 0, h_sam)
-                                .view(1, seq_h_int, 1, -1)
-                                .expand(seq_f_int, seq_h_int, seq_w_int, -1),
-                                torch.index_select(freqs_split[2], 0, w_sam)
-                                .view(1, 1, seq_w_int, -1)
-                                .expand(seq_f_int, seq_h_int, seq_w_int, -1),
-                            ],
-                            dim=-1,
-                        ).reshape(seq_len, 1, -1)
-                    elif t_f < 0:
-                        freqs_i = trainable_f.unsqueeze(1)
-                    output[i, seq_bucket[-1] : seq_bucket[-1] + seq_len] = freqs_i
-            seq_bucket.append(seq_bucket[-1] + seq_len)
-        return output
-
-
-class RotaryEmbeddingWanS2V(RotaryEmbeddingWan):
-    """Apply RoPE using precomputed complex freqs for Wan S2V main transformer.
-
-    Converts complex freqs (from WanS2VRotaryPosEmbed) to cos/sin and delegates
-    to RotaryEmbeddingWan for platform-optimized application (float32 kernel).
-    Under TP, freqs has 1 head — broadcasts automatically via cos/sin.
-    """
-
-    def __init__(self) -> None:
-        super().__init__(is_neox_style=False, half_head_dim=True)
-
-    def forward(self, x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
-        freqs_sliced = freqs[:, : x.size(1)]
-        cos = freqs_sliced.real.to(x.dtype)
-        sin = freqs_sliced.imag.to(x.dtype)
-        return super().forward(x, cos, sin)
-
-
 class RotaryEmbeddingS2VGrid(torch.nn.Module):
     """Grid-based RoPE for S2V motioner/init attention.
 

vllm_omni/diffusion/models/bagel/bagel_transformer.py

@@ -1,6 +1,7 @@
 # Copyright 2025 Bytedance Ltd. and/or its affiliates.
 # Copyright (c) 2024 The Qwen Team and The HuggingFace Inc. team.
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 #
 # This file has been modified by ByteDance Ltd. and/or its affiliates.
 #
@@ -48,10 +49,10 @@ from vllm_omni.diffusion.distributed.parallel_state import (
     get_sp_group,
 )
 from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
-from vllm_omni.diffusion.layers.mot.mot_layernorm import MoTRMSNorm
-from vllm_omni.diffusion.layers.mot.mot_qkv_parallel_linear import MoTQKVParallelLinear
-from vllm_omni.diffusion.layers.mot.mot_row_parallel_linear import MoTRowParallelLinear
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding
+from vllm_omni.diffusion.models.bagel.mot.mot_layernorm import MoTRMSNorm
+from vllm_omni.diffusion.models.bagel.mot.mot_qkv_parallel_linear import MoTQKVParallelLinear
+from vllm_omni.diffusion.models.bagel.mot.mot_row_parallel_linear import MoTRowParallelLinear
 from vllm_omni.diffusion.utils.kv_utils import left_pad_stack
 from vllm_omni.model_executor.layers.timestep_embedding import timestep_embedding
 

vllm_omni/diffusion/models/bagel/mot/mot_layernorm.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # ruff: noqa: N803, E741
 """Mixture-of-Tokens (MoT) RMS Normalization layer.
 
@@ -68,7 +71,7 @@ class MoTRMSNorm(CustomOp):
             return ir.ops.rms_norm(x, self.weight.data, self.variance_epsilon)
 
         # gen mode – fused MoT Triton kernel
-        from vllm_omni.diffusion.layers.mot.ops.mot_rms_norm import (
+        from vllm_omni.diffusion.models.bagel.mot.ops.mot_rms_norm import (
             mot_rms_norm,
         )
 

vllm_omni/diffusion/models/bagel/mot/mot_qkv_parallel_linear.py

@@ -19,7 +19,7 @@ from vllm.model_executor.layers.quantization.base_config import (
 from vllm.model_executor.utils import set_weight_attrs
 from vllm.platforms import current_platform
 
-from vllm_omni.diffusion.layers.mot.ops.mot_gemm import invoke_mot_gemm
+from vllm_omni.diffusion.models.bagel.mot.ops.mot_gemm import invoke_mot_gemm
 
 logger = init_logger(__name__)
 

vllm_omni/diffusion/models/bagel/mot/mot_row_parallel_linear.py

@@ -19,7 +19,7 @@ from vllm.model_executor.layers.quantization.base_config import (
 from vllm.model_executor.utils import set_weight_attrs
 from vllm.platforms import current_platform
 
-from vllm_omni.diffusion.layers.mot.ops.mot_gemm import invoke_mot_gemm
+from vllm_omni.diffusion.models.bagel.mot.ops.mot_gemm import invoke_mot_gemm
 
 
 class MoTRowParallelLinear(RowParallelLinear):

vllm_omni/diffusion/models/bagel/mot/ops/mot_gemm.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # ruff: noqa: N803, E741
 import functools
 import json

vllm_omni/diffusion/models/bagel/mot/ops/mot_rms_norm.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # ruff: noqa: N803, E741
 import torch
 from vllm.triton_utils import tl, triton

vllm_omni/diffusion/models/wan2_2/norm.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Normalization layer used by the Wan VAE patch."""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from vllm_omni.diffusion.layers.custom_op import CustomOp
+
+
+class RMSNormVAE(CustomOp):
+    """Root Mean Square Layer Normalization for Channel-First or Last"""
+
+    def __init__(
+        self,
+        dim: int,
+        channel_first: bool = True,
+        images: bool = True,
+        bias: bool = False,
+        epsilon: float = 1e-6,
+    ) -> None:
+        super().__init__()
+        broadcastable_dims = (1, 1, 1) if not images else (1, 1)
+        shape = (dim, *broadcastable_dims) if channel_first else (dim,)
+
+        self.channel_first = channel_first
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(shape))
+        self.bias = nn.Parameter(torch.zeros(shape)) if bias else None
+        self.epsilon = epsilon
+
+        self.gamma_rmsnorm = None
+
+    def forward_cuda(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_hip(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_npu(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        import torch_npu
+
+        if self.gamma_rmsnorm is None:
+            self.gamma_rmsnorm = self.gamma.reshape(-1)
+
+        if self.channel_first:
+            x = x.transpose(1, -1)
+            out = torch_npu.npu_rms_norm(x, self.gamma_rmsnorm, epsilon=self.epsilon)[0].transpose(1, -1)
+        else:
+            out = torch_npu.npu_rms_norm(x, self.gamma_rmsnorm, epsilon=self.epsilon)[0]
+
+        if self.bias is not None:
+            out = out + self.bias
+        return out
+
+    def forward_native(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        out = (
+            F.normalize(
+                x,
+                dim=(1 if self.channel_first else -1),
+                eps=self.epsilon,
+            )
+            * self.scale
+            * self.gamma
+        )
+        if self.bias is not None:
+            out = out + self.bias
+        return out

vllm_omni/diffusion/models/wan2_2/patch_diffusers.py

@@ -1,9 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import sys
 
-from vllm_omni.diffusion.layers.norm import RMSNormVAE
+from vllm_omni.diffusion.models.wan2_2.norm import RMSNormVAE
 
 
 def patch_wan_rms_norm():

vllm_omni/diffusion/models/wan2_2/rope.py

@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Wan S2V rotary positional embedding layers."""
+
+import torch
+
+from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan
+from vllm_omni.platforms import current_omni_platform
+
+
+class WanS2VRotaryPosEmbed(torch.nn.Module):
+    """Precompute complex-valued RoPE embeddings for S2V multi-grid positions.
+
+    Owns the base frequency buffer and provides forward() to compute position
+    embeddings given hidden_states (for shape) and grid_sizes.
+    """
+
+    def __init__(self, num_heads: int, head_dim: int, max_seq_len: int = 1024):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        d = head_dim
+        freqs = torch.cat(
+            [
+                self._rope_params(max_seq_len, d - 4 * (d // 6)),
+                self._rope_params(max_seq_len, 2 * (d // 6)),
+                self._rope_params(max_seq_len, 2 * (d // 6)),
+            ],
+            dim=1,
+        )
+        self.register_buffer("freqs", freqs.to(torch.complex64), persistent=False)
+
+    @staticmethod
+    @torch.amp.autocast(current_omni_platform.device_type, enabled=False)
+    def _rope_params(max_seq_len, dim, theta=10000):
+        if dim % 2 != 0:
+            raise ValueError(f"dim ({dim}) must be even")
+        freqs = torch.outer(
+            torch.arange(max_seq_len), 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float64).div(dim))
+        )
+        return torch.polar(torch.ones_like(freqs), freqs)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        grid_sizes: list,
+        trainable_freqs: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Precompute RoPE embeddings for the given grid layout.
+
+        Args:
+            hidden_states: Tensor [B, S, ...] (used for batch/seq shape and device)
+            grid_sizes: Grid specification (list of [offsets, sizes, totals])
+            trainable_freqs: Optional trainable frequency overrides for t_f < 0
+
+        Returns:
+            Complex tensor [B, S, 1, head_dim//2] of precomputed position embeddings
+        """
+        b, s = hidden_states.shape[0], hidden_states.shape[1]
+        c = self.head_dim // 2
+        device = hidden_states.device
+
+        freqs = self.freqs.to(device)
+        if trainable_freqs is not None:
+            freqs_input = [freqs, trainable_freqs]
+        else:
+            freqs_input = freqs
+
+        if isinstance(freqs_input, list):
+            trainable_f = freqs_input[1]
+            freqs_split = freqs_input[0]
+        else:
+            trainable_f = None
+            freqs_split = freqs_input
+        freqs_split = freqs_split.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+        output = torch.empty((b, s, 1, c), device=device, dtype=torch.complex64)
+        seq_bucket = [0]
+        if not isinstance(grid_sizes, list):
+            grid_sizes = [grid_sizes]
+        for g in grid_sizes:
+            if not isinstance(g, list):
+                g = [torch.zeros_like(g), g]
+            batch_size = g[0].shape[0]
+            for i in range(batch_size):
+                f_o, h_o, w_o = g[0][i]
+                f, h, w = g[1][i]
+                t_f, t_h, t_w = g[2][i]
+                seq_f, seq_h, seq_w = f - f_o, h - h_o, w - w_o
+                seq_len = int(seq_f * seq_h * seq_w)
+                if seq_len > 0:
+                    if t_f > 0:
+                        assert f_o * f >= 0 and h_o * h >= 0 and w_o * w >= 0
+                        seq_f_int = int(seq_f)
+                        seq_h_int = int(seq_h)
+                        seq_w_int = int(seq_w)
+
+                        if f_o >= 0:
+                            f_sam = torch.linspace(int(f_o), int(t_f + f_o) - 1, seq_f_int, device=device).long()
+                        else:
+                            f_sam = torch.linspace(int(-f_o), int(-t_f - f_o) + 1, seq_f_int, device=device).long()
+                        h_sam = torch.linspace(int(h_o), int(t_h + h_o) - 1, seq_h_int, device=device).long()
+                        w_sam = torch.linspace(int(w_o), int(t_w + w_o) - 1, seq_w_int, device=device).long()
+
+                        freqs_0 = torch.index_select(freqs_split[0] if f_o >= 0 else freqs_split[0].conj(), 0, f_sam)
+                        freqs_0 = freqs_0.view(seq_f_int, 1, 1, -1)
+
+                        freqs_i = torch.cat(
+                            [
+                                freqs_0.expand(seq_f_int, seq_h_int, seq_w_int, -1),
+                                torch.index_select(freqs_split[1], 0, h_sam)
+                                .view(1, seq_h_int, 1, -1)
+                                .expand(seq_f_int, seq_h_int, seq_w_int, -1),
+                                torch.index_select(freqs_split[2], 0, w_sam)
+                                .view(1, 1, seq_w_int, -1)
+                                .expand(seq_f_int, seq_h_int, seq_w_int, -1),
+                            ],
+                            dim=-1,
+                        ).reshape(seq_len, 1, -1)
+                    elif t_f < 0:
+                        freqs_i = trainable_f.unsqueeze(1)
+                    output[i, seq_bucket[-1] : seq_bucket[-1] + seq_len] = freqs_i
+            seq_bucket.append(seq_bucket[-1] + seq_len)
+        return output
+
+
+class RotaryEmbeddingWanS2V(RotaryEmbeddingWan):
+    """Apply RoPE using precomputed complex freqs for Wan S2V main transformer.
+
+    Converts complex freqs (from WanS2VRotaryPosEmbed) to cos/sin and delegates
+    to RotaryEmbeddingWan for platform-optimized application (float32 kernel).
+    Under TP, freqs has 1 head — broadcasts automatically via cos/sin.
+    """
+
+    def __init__(self) -> None:
+        super().__init__(is_neox_style=False, half_head_dim=True)
+
+    def forward(self, x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+        freqs_sliced = freqs[:, : x.size(1)]
+        cos = freqs_sliced.real.to(x.dtype)
+        sin = freqs_sliced.imag.to(x.dtype)
+        return super().forward(x, cos, sin)

vllm_omni/diffusion/models/wan2_2/wan2_2_s2v_transformer.py

@@ -32,7 +32,8 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput
-from vllm_omni.diffusion.layers.rope import RotaryEmbeddingS2VGrid, RotaryEmbeddingWanS2V, WanS2VRotaryPosEmbed
+from vllm_omni.diffusion.layers.rope import RotaryEmbeddingS2VGrid
+from vllm_omni.diffusion.models.wan2_2.rope import RotaryEmbeddingWanS2V, WanS2VRotaryPosEmbed
 from vllm_omni.platforms import current_omni_platform
 
 from .wan2_2_transformer import DistributedRMSNorm, WanFeedForward