diff --git a/benchmarks/kernels/bench_qwen_select01_modulation.py b/benchmarks/kernels/bench_qwen_select01_modulation.py new file mode 100644 index 000000000..9a97c8de3 --- /dev/null +++ b/benchmarks/kernels/bench_qwen_select01_modulation.py @@ -0,0 +1,333 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Benchmark the Qwen-Image-Edit select01 modulation hot path. + +This script intentionally starts with the current PyTorch implementation. It is +used as the before/after baseline for SGLang PR #20395/#21318 style fusion: + +* layernorm + scale/shift + binary select01 gate +* residual + gated update + layernorm + scale/shift + binary select01 gate + +Run on a CUDA machine for actionable numbers, for example: + + python benchmarks/kernels/bench_qwen_select01_modulation.py \ + --device cuda --dtype bf16 --seq-lens 4096 8192 16384 \ + --output-json /tmp/qwen_select01_baseline.json +""" + +from __future__ import annotations + +import argparse +import json +import statistics +import sys +import time +from collections.abc import Callable +from pathlib import Path + +import torch +import torch.nn.functional as F + +for candidate in Path(__file__).resolve().parents: + if (candidate / "vllm_omni").is_dir(): + if str(candidate) not in sys.path: + sys.path.insert(0, str(candidate)) + break + +from vllm_omni.diffusion.layers.qwen_select01_modulation import ( # noqa: E402 + fused_layernorm_select01, + fused_residual_layernorm_select01, +) + + +def _dtype(name: str) -> torch.dtype: + return { + "fp32": torch.float32, + "float32": torch.float32, + "fp16": torch.float16, + "float16": torch.float16, + "bf16": torch.bfloat16, + "bfloat16": torch.bfloat16, + }[name] + + +def _select01_modulation( + mod_params: torch.Tensor, + index: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + shift, scale, gate = mod_params.chunk(3, dim=-1) + + if index is None: + return scale.unsqueeze(1), shift.unsqueeze(1), gate.unsqueeze(1) + + actual_batch = shift.size(0) // 2 + shift_0, shift_1 = shift[:actual_batch], shift[actual_batch:] + scale_0, scale_1 = scale[:actual_batch], scale[actual_batch:] + gate_0, gate_1 = gate[:actual_batch], gate[actual_batch:] + + idx = index.unsqueeze(-1) == 0 + shift = torch.where(idx, shift_0.unsqueeze(1), shift_1.unsqueeze(1)) + scale = torch.where(idx, scale_0.unsqueeze(1), scale_1.unsqueeze(1)) + gate = torch.where(idx, gate_0.unsqueeze(1), gate_1.unsqueeze(1)) + return scale, shift, gate + + +def _native_layernorm_select01( + x: torch.Tensor, + mod_params: torch.Tensor, + index: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + scale, shift, gate = _select01_modulation(mod_params, index) + out = F.layer_norm(x.float(), (x.shape[-1],), eps=eps).to(x.dtype) + return out * (1 + scale) + shift, gate + + +def _native_residual_layernorm_select01( + x: torch.Tensor, + residual: torch.Tensor, + residual_gate: torch.Tensor, + mod_params: torch.Tensor, + index: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + scale, shift, gate = _select01_modulation(mod_params, index) + residual_out = residual + residual_gate * x + out = F.layer_norm(residual_out.float(), (residual_out.shape[-1],), eps=eps).to(residual_out.dtype) + return out * (1 + scale) + shift, residual_out, gate + + +def _time_cuda(fn: Callable[[], object], warmup: int, iters: int) -> list[float]: + for _ in range(warmup): + fn() + torch.accelerator.synchronize() + + times: list[float] = [] + for _ in range(iters): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + fn() + end.record() + torch.accelerator.synchronize() + times.append(float(start.elapsed_time(end))) + return times + + +def _reset_peak_memory_stats(device: torch.device) -> None: + reset_peak_memory_stats = getattr(torch.accelerator, "reset_peak_memory_stats", None) + if reset_peak_memory_stats is None: + reset_peak_memory_stats = getattr(getattr(torch, device.type), "reset_peak_memory_stats") + reset_peak_memory_stats(device) + + +def _max_memory_allocated(device: torch.device) -> int: + max_memory_allocated = getattr(torch.accelerator, "max_memory_allocated", None) + if max_memory_allocated is None: + max_memory_allocated = getattr(getattr(torch, device.type), "max_memory_allocated") + return int(max_memory_allocated(device)) + + +def _time_cpu(fn: Callable[[], object], warmup: int, iters: int) -> list[float]: + for _ in range(warmup): + fn() + + times: list[float] = [] + for _ in range(iters): + t0 = time.perf_counter() + fn() + times.append((time.perf_counter() - t0) * 1000.0) + return times + + +def _summarize(times: list[float]) -> dict[str, float]: + ordered = sorted(times) + return { + "median_ms": statistics.median(ordered), + "min_ms": ordered[0], + "p20_ms": ordered[max(0, int(len(ordered) * 0.20) - 1)], + "p80_ms": ordered[min(len(ordered) - 1, int(len(ordered) * 0.80))], + } + + +def _make_inputs( + batch: int, + seq_len: int, + hidden_size: int, + dtype: torch.dtype, + device: torch.device, +) -> dict[str, torch.Tensor]: + torch.manual_seed(0) + return { + "x": torch.randn(batch, seq_len, hidden_size, device=device, dtype=dtype), + "residual": torch.randn(batch, seq_len, hidden_size, device=device, dtype=dtype), + "residual_gate": torch.randn(batch, 1, hidden_size, device=device, dtype=dtype), + "mod_params": torch.randn(batch * 2, hidden_size * 3, device=device, dtype=dtype), + "index": torch.randint(0, 2, (batch, seq_len), device=device, dtype=torch.int64), + } + + +def _bench_one( + name: str, + fn: Callable[[], object], + device: torch.device, + warmup: int, + iters: int, +) -> dict[str, object]: + if device.type == "cuda": + _reset_peak_memory_stats(device) + times = _time_cuda(fn, warmup, iters) + peak_memory_mb = _max_memory_allocated(device) / 1024**2 + else: + times = _time_cpu(fn, warmup, iters) + peak_memory_mb = None + + return { + "name": name, + **_summarize(times), + "peak_memory_mb": peak_memory_mb, + } + + +def _run(args: argparse.Namespace) -> dict[str, object]: + device = torch.device(args.device) + if device.type == "cuda" and not torch.cuda.is_available(): + raise SystemExit("CUDA is not available on this machine.") + + dtype = _dtype(args.dtype) + rows: list[dict[str, object]] = [] + + for seq_len in args.seq_lens: + tensors = _make_inputs(args.batch_size, seq_len, args.hidden_size, dtype, device) + impls = ["native"] if args.impl == "native" else ["fused"] if args.impl == "fused" else ["native", "fused"] + + def first_norm() -> object: + return _native_layernorm_select01( + tensors["x"], + tensors["mod_params"], + tensors["index"], + args.eps, + ) + + def residual_norm() -> object: + return _native_residual_layernorm_select01( + tensors["x"], + tensors["residual"], + tensors["residual_gate"], + tensors["mod_params"], + tensors["index"], + args.eps, + ) + + callables: list[tuple[str, Callable[[], object]]] = [] + if "native" in impls: + callables.extend( + [ + ("native_layernorm_select01", first_norm), + ("native_residual_layernorm_select01", residual_norm), + ] + ) + if "fused" in impls: + + def fused_first_norm() -> object: + return fused_layernorm_select01( + tensors["x"], + tensors["mod_params"], + tensors["index"], + args.eps, + ) + + def fused_residual_norm() -> object: + return fused_residual_layernorm_select01( + tensors["x"], + tensors["residual"], + tensors["residual_gate"].expand_as(tensors["residual"]), + tensors["mod_params"], + tensors["index"], + args.eps, + ) + + callables.extend( + [ + ("fused_layernorm_select01", fused_first_norm), + ("fused_residual_layernorm_select01", fused_residual_norm), + ] + ) + + if args.check and "fused" in impls: + ref_norm, ref_gate = first_norm() + actual_norm, actual_gate = fused_layernorm_select01( + tensors["x"], + tensors["mod_params"], + tensors["index"], + args.eps, + ) + torch.testing.assert_close(actual_norm, ref_norm, atol=args.check_atol, rtol=args.check_rtol) + torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0) + ref_norm, ref_residual, ref_gate = residual_norm() + actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01( + tensors["x"], + tensors["residual"], + tensors["residual_gate"].expand_as(tensors["residual"]), + tensors["mod_params"], + tensors["index"], + args.eps, + ) + torch.testing.assert_close(actual_norm, ref_norm, atol=args.check_atol, rtol=args.check_rtol) + torch.testing.assert_close(actual_residual, ref_residual, atol=args.check_atol, rtol=args.check_rtol) + torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0) + + for name, fn in callables: + row = _bench_one(name, fn, device, args.warmup, args.iters) + row.update( + { + "batch_size": args.batch_size, + "seq_len": seq_len, + "hidden_size": args.hidden_size, + "dtype": args.dtype, + "device": str(device), + } + ) + rows.append(row) + print( + f"{name:<38} B={args.batch_size:<2} L={seq_len:<6} C={args.hidden_size:<5} " + f"median={row['median_ms']:.4f} ms p20={row['p20_ms']:.4f} p80={row['p80_ms']:.4f}" + ) + + result = { + "benchmark": "qwen_select01_modulation", + "torch": torch.__version__, + "cuda_available": torch.cuda.is_available(), + "gpu": torch.cuda.get_device_name(device) if device.type == "cuda" else None, + "rows": rows, + } + + if args.output_json: + output_path = Path(args.output_json) + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(json.dumps(result, indent=2), encoding="utf-8") + + return result + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + default_device = "cuda" if torch.cuda.is_available() else "cpu" + parser.add_argument("--device", default=default_device) + parser.add_argument("--dtype", default="bf16", choices=["fp32", "float32", "fp16", "float16", "bf16", "bfloat16"]) + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument("--seq-lens", type=int, nargs="+", default=[4096, 8192, 16384]) + parser.add_argument("--hidden-size", type=int, default=3072) + parser.add_argument("--eps", type=float, default=1e-6) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--iters", type=int, default=100) + parser.add_argument("--impl", choices=["native", "fused", "all"], default="native") + parser.add_argument("--check", action="store_true") + parser.add_argument("--check-atol", type=float, default=5e-2) + parser.add_argument("--check-rtol", type=float, default=5e-2) + parser.add_argument("--output-json") + return parser.parse_args() + + +if __name__ == "__main__": + _run(_parse_args()) diff --git a/tests/diffusion/layers/test_qwen_select01_modulation.py b/tests/diffusion/layers/test_qwen_select01_modulation.py new file mode 100644 index 000000000..050da4717 --- /dev/null +++ b/tests/diffusion/layers/test_qwen_select01_modulation.py @@ -0,0 +1,299 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from vllm_omni.diffusion.layers import qwen_select01_modulation +from vllm_omni.diffusion.layers.qwen_select01_modulation import ( + can_use_qwen_select01_triton, + fused_layernorm_select01, + fused_residual_layernorm_select01, + select01_modulation_native, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] + + +@pytest.mark.cpu +@pytest.mark.parametrize( + ("platform_is_cuda", "tensor_is_cuda", "has_triton", "expected"), + [ + (True, True, True, True), + (False, True, True, False), + (True, False, True, False), + (True, True, False, False), + ], +) +def test_qwen_select01_triton_requires_cuda_platform( + monkeypatch, + platform_is_cuda: bool, + tensor_is_cuda: bool, + has_triton: bool, + expected: bool, +): + class FakePlatform: + @staticmethod + def is_cuda() -> bool: + return platform_is_cuda + + class FakeTensor: + is_cuda = tensor_is_cuda + + monkeypatch.setattr(qwen_select01_modulation, "current_platform", FakePlatform()) + monkeypatch.setattr(qwen_select01_modulation, "HAS_TRITON", has_triton) + + assert can_use_qwen_select01_triton(FakeTensor()) is expected + + +def _make_inputs( + batch_size: int, + seq_len: int, + hidden_size: int, + dtype: torch.dtype, + device: torch.device, +) -> tuple[torch.Tensor, ...]: + generator = torch.Generator(device=device) + generator.manual_seed(202817) + x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device=device, generator=generator) + residual = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device=device, generator=generator) + residual_gate = torch.randn(batch_size, 1, hidden_size, dtype=dtype, device=device, generator=generator) + mod_params = torch.randn(batch_size * 2, hidden_size * 3, dtype=dtype, device=device, generator=generator) + index = torch.randint(0, 2, (batch_size, seq_len), dtype=torch.int64, device=device, generator=generator) + return x, residual, residual_gate, mod_params, index + + +def _reference_layernorm_select01( + x: torch.Tensor, + mod_params: torch.Tensor, + index: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + scale, shift, gate = select01_modulation_native(mod_params, index) + out = F.layer_norm(x.float(), (x.shape[-1],), eps=eps).to(x.dtype) + return out * (1 + scale) + shift, gate + + +def _reference_residual_layernorm_select01( + x: torch.Tensor, + residual: torch.Tensor, + residual_gate: torch.Tensor, + mod_params: torch.Tensor, + index: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + scale, shift, gate = select01_modulation_native(mod_params, index) + residual_out = residual + residual_gate * x + out = F.layer_norm(residual_out.float(), (residual_out.shape[-1],), eps=eps).to(residual_out.dtype) + return out * (1 + scale) + shift, residual_out, gate + + +@pytest.mark.cpu +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("broadcast_index", [False, True]) +def test_qwen_select01_native_fallback_matches_reference(dtype: torch.dtype, broadcast_index: bool): + device = torch.device("cpu") + x, residual, residual_gate, mod_params, index = _make_inputs(2, 17, 64, dtype, device) + if broadcast_index: + index = index[:1] + eps = 1e-6 + + actual_norm, actual_gate = fused_layernorm_select01(x, mod_params, index, eps) + ref_norm, ref_gate = _reference_layernorm_select01(x, mod_params, index, eps) + torch.testing.assert_close(actual_norm, ref_norm, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0) + + actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01( + x, + residual, + residual_gate, + mod_params, + index, + eps, + ) + ref_norm, ref_residual, ref_gate = _reference_residual_layernorm_select01( + x, + residual, + residual_gate, + mod_params, + index, + eps, + ) + torch.testing.assert_close(actual_norm, ref_norm, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(actual_residual, ref_residual, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize( + ("dtype", "atol", "rtol"), + [ + (torch.float32, 1e-5, 1e-5), + (torch.bfloat16, 5e-2, 5e-2), + ], +) +@pytest.mark.parametrize("seq_len", [1, 7, 128, 1024]) +def test_qwen_select01_triton_matches_reference(dtype: torch.dtype, atol: float, rtol: float, seq_len: int): + device = torch.device("cuda:0") + x, residual, residual_gate, mod_params, index = _make_inputs(2, seq_len, 256, dtype, device) + eps = 1e-6 + + actual_norm, actual_gate = fused_layernorm_select01(x, mod_params, index, eps) + ref_norm, ref_gate = _reference_layernorm_select01(x, mod_params, index, eps) + torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol) + torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0) + + actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01( + x, + residual, + residual_gate.expand_as(residual), + mod_params, + index, + eps, + ) + ref_norm, ref_residual, ref_gate = _reference_residual_layernorm_select01( + x, + residual, + residual_gate.expand_as(residual), + mod_params, + index, + eps, + ) + torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol) + torch.testing.assert_close(actual_residual, ref_residual, atol=atol, rtol=rtol) + torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize( + ("dtype", "atol", "rtol"), + [ + (torch.float32, 1e-5, 1e-5), + (torch.bfloat16, 5e-2, 5e-2), + ], +) +def test_qwen_select01_triton_matches_reference_for_structured_multi_batch_index( + dtype: torch.dtype, + atol: float, + rtol: float, +): + device = torch.device("cuda:0") + x, residual, residual_gate, mod_params, _ = _make_inputs(3, 16, 256, dtype, device) + index = torch.tensor( + [ + [0] * 16, + [1] * 16, + [0, 1] * 8, + ], + dtype=torch.int64, + device=device, + ) + eps = 1e-6 + + actual_norm, actual_gate = fused_layernorm_select01(x, mod_params, index, eps) + ref_norm, ref_gate = _reference_layernorm_select01(x, mod_params, index, eps) + torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol) + torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0) + + actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01( + x, + residual, + residual_gate.expand_as(residual), + mod_params, + index, + eps, + ) + ref_norm, ref_residual, ref_gate = _reference_residual_layernorm_select01( + x, + residual, + residual_gate.expand_as(residual), + mod_params, + index, + eps, + ) + torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol) + torch.testing.assert_close(actual_residual, ref_residual, atol=atol, rtol=rtol) + torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize( + ("dtype", "atol", "rtol"), + [ + (torch.float32, 1e-5, 1e-5), + (torch.bfloat16, 5e-2, 5e-2), + ], +) +def test_qwen_select01_triton_broadcasts_single_index_batch(dtype: torch.dtype, atol: float, rtol: float): + device = torch.device("cuda:0") + x, residual, residual_gate, mod_params, index = _make_inputs(2, 17, 256, dtype, device) + index = index[:1] + eps = 1e-6 + + actual_norm, actual_gate = fused_layernorm_select01(x, mod_params, index, eps) + ref_norm, ref_gate = _reference_layernorm_select01(x, mod_params, index, eps) + torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol) + torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0) + + actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01( + x, + residual, + residual_gate.expand_as(residual), + mod_params, + index, + eps, + ) + ref_norm, ref_residual, ref_gate = _reference_residual_layernorm_select01( + x, + residual, + residual_gate.expand_as(residual), + mod_params, + index, + eps, + ) + torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol) + torch.testing.assert_close(actual_residual, ref_residual, atol=atol, rtol=rtol) + torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_qwen_select01_cuda_falls_back_without_triton(monkeypatch): + device = torch.device("cuda:0") + x, residual, residual_gate, mod_params, index = _make_inputs(2, 7, 256, torch.bfloat16, device) + residual_gate = residual_gate.expand_as(residual) + eps = 1e-6 + + def fail_if_launched(*args, **kwargs): + pytest.fail("Triton launcher must not run when HAS_TRITON is false") + + monkeypatch.setattr(qwen_select01_modulation, "HAS_TRITON", False) + monkeypatch.setattr(qwen_select01_modulation, "_launch_layernorm_select01", fail_if_launched) + monkeypatch.setattr(qwen_select01_modulation, "_launch_residual_layernorm_select01", fail_if_launched) + + actual_norm, actual_gate = fused_layernorm_select01(x, mod_params, index, eps) + ref_norm, ref_gate = _reference_layernorm_select01(x, mod_params, index, eps) + torch.testing.assert_close(actual_norm, ref_norm, atol=0, rtol=0) + torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0) + + actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01( + x, + residual, + residual_gate, + mod_params, + index, + eps, + ) + ref_norm, ref_residual, ref_gate = _reference_residual_layernorm_select01( + x, + residual, + residual_gate, + mod_params, + index, + eps, + ) + torch.testing.assert_close(actual_norm, ref_norm, atol=0, rtol=0) + torch.testing.assert_close(actual_residual, ref_residual, atol=0, rtol=0) + torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0) diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_select01_routing.py b/tests/diffusion/models/qwen_image/test_qwen_image_select01_routing.py new file mode 100644 index 000000000..1a3acbb48 --- /dev/null +++ b/tests/diffusion/models/qwen_image/test_qwen_image_select01_routing.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.diffusion.models.qwen_image import qwen_image_transformer +from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import QwenImageTransformerBlock + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +class _StaticOutput(nn.Module): + def __init__(self, output: torch.Tensor): + super().__init__() + self.output = output + + def forward(self, _: torch.Tensor) -> torch.Tensor: + return self.output + + +class _RecordingNorm(nn.Module): + def __init__(self): + super().__init__() + self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = [] + + def forward(self, x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor) -> torch.Tensor: + self.calls.append((x, scale, shift)) + return x + + +class _ZeroAttention(nn.Module): + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + **_: object, + ) -> tuple[torch.Tensor, torch.Tensor]: + return torch.zeros_like(hidden_states), torch.zeros_like(encoder_hidden_states) + + +class _ZeroMLP(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.zeros_like(x) + + +def test_non_cuda_select01_uses_adalayernorm_dispatch(monkeypatch): + hidden_size = 4 + block = QwenImageTransformerBlock.__new__(QwenImageTransformerBlock) + nn.Module.__init__(block) + block.zero_cond_t = True + block.img_mod = _StaticOutput(torch.zeros(2, 6 * hidden_size)) + block.txt_mod = _StaticOutput(torch.zeros(1, 6 * hidden_size)) + block.img_norm1 = _RecordingNorm() + block.img_norm2 = _RecordingNorm() + block.txt_norm1 = _RecordingNorm() + block.txt_norm2 = _RecordingNorm() + block.attn = _ZeroAttention() + block.img_mlp = _ZeroMLP() + block.txt_mlp = _ZeroMLP() + + monkeypatch.setattr(qwen_image_transformer, "can_use_qwen_select01_triton", lambda _: False) + + def fail_if_called(*args, **kwargs): + pytest.fail("non-CUDA select01 must preserve the AdaLayerNorm dispatch") + + monkeypatch.setattr(qwen_image_transformer, "fused_layernorm_select01", fail_if_called) + monkeypatch.setattr(qwen_image_transformer, "fused_residual_layernorm_select01", fail_if_called) + + hidden_states = torch.randn(1, 3, hidden_size) + encoder_hidden_states = torch.randn(1, 2, hidden_size) + modulate_index = torch.tensor([[0, 1, 0]]) + block( + hidden_states, + encoder_hidden_states, + torch.ones(1, 2), + torch.zeros(2, hidden_size), + (torch.empty(0), torch.empty(0)), + modulate_index=modulate_index, + ) + + assert len(block.img_norm1.calls) == 1 + assert len(block.img_norm2.calls) == 1 + for _, scale, shift in (*block.img_norm1.calls, *block.img_norm2.calls): + assert scale.shape == hidden_states.shape + assert shift.shape == hidden_states.shape diff --git a/vllm_omni/diffusion/layers/qwen_select01_modulation.py b/vllm_omni/diffusion/layers/qwen_select01_modulation.py new file mode 100644 index 000000000..4fe1e3e0e --- /dev/null +++ b/vllm_omni/diffusion/layers/qwen_select01_modulation.py @@ -0,0 +1,491 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import torch +import torch.nn.functional as F +from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON, tl, triton + + +@triton.jit +def _layernorm_select01_kernel( + output_ptr, + gate_out_ptr, + x_ptr, + weight_ptr, + bias_ptr, + scale0_ptr, + shift0_ptr, + gate0_ptr, + scale1_ptr, + shift1_ptr, + gate1_ptr, + index_ptr, + inner_dim, + seq_len, + stride_x_row, + stride_out_row, + stride_gate_out_row, + stride_w, + stride_b, + stride_s0_b, + stride_s0_c, + stride_sh0_b, + stride_sh0_c, + stride_g0_b, + stride_g0_c, + stride_s1_b, + stride_s1_c, + stride_sh1_b, + stride_sh1_c, + stride_g1_b, + stride_g1_c, + stride_i_b, + stride_i_l, + eps, + has_weight: tl.constexpr, + has_bias: tl.constexpr, + block_n: tl.constexpr, +): + row = tl.program_id(0) + cols = tl.arange(0, block_n) + mask = cols < inner_dim + + x_row_ptr = x_ptr + row * stride_x_row + x = tl.load(x_row_ptr + cols, mask=mask, other=0.0).to(tl.float32) + + mean = tl.sum(x, axis=0) / inner_dim + xbar = tl.where(mask, x - mean, 0.0) + var = tl.sum(xbar * xbar, axis=0) / inner_dim + x_hat = (x - mean) * tl.rsqrt(var + eps) + + if has_weight: + weight = tl.load(weight_ptr + cols * stride_w, mask=mask, other=1.0).to(tl.float32) + x_hat = x_hat * weight + if has_bias: + bias = tl.load(bias_ptr + cols * stride_b, mask=mask, other=0.0).to(tl.float32) + x_hat = x_hat + bias + + batch_idx = row // seq_len + seq_idx = row % seq_len + idx = tl.load(index_ptr + batch_idx * stride_i_b + seq_idx * stride_i_l).to(tl.int1) + + scale0_ptrs = scale0_ptr + batch_idx * stride_s0_b + cols * stride_s0_c + shift0_ptrs = shift0_ptr + batch_idx * stride_sh0_b + cols * stride_sh0_c + gate0_ptrs = gate0_ptr + batch_idx * stride_g0_b + cols * stride_g0_c + scale1_ptrs = scale1_ptr + batch_idx * stride_s1_b + cols * stride_s1_c + shift1_ptrs = shift1_ptr + batch_idx * stride_sh1_b + cols * stride_sh1_c + gate1_ptrs = gate1_ptr + batch_idx * stride_g1_b + cols * stride_g1_c + + scale_ptrs = tl.where(idx, scale1_ptrs, scale0_ptrs) + shift_ptrs = tl.where(idx, shift1_ptrs, shift0_ptrs) + gate_ptrs = tl.where(idx, gate1_ptrs, gate0_ptrs) + + scale = tl.load(scale_ptrs, mask=mask, other=0.0).to(tl.float32) + shift = tl.load(shift_ptrs, mask=mask, other=0.0).to(tl.float32) + gate = tl.load(gate_ptrs, mask=mask, other=0.0) + + y = x_hat * (1.0 + scale) + shift + tl.store(output_ptr + row * stride_out_row + cols, y, mask=mask) + tl.store(gate_out_ptr + row * stride_gate_out_row + cols, gate, mask=mask) + + +@triton.jit +def _residual_layernorm_select01_kernel( + output_ptr, + residual_out_ptr, + gate_out_ptr, + x_ptr, + residual_ptr, + residual_gate_ptr, + weight_ptr, + bias_ptr, + scale0_ptr, + shift0_ptr, + gate0_ptr, + scale1_ptr, + shift1_ptr, + gate1_ptr, + index_ptr, + inner_dim, + seq_len, + stride_x_row, + stride_res_row, + stride_res_gate_row, + stride_out_row, + stride_res_out_row, + stride_gate_out_row, + stride_w, + stride_b, + stride_s0_b, + stride_s0_c, + stride_sh0_b, + stride_sh0_c, + stride_g0_b, + stride_g0_c, + stride_s1_b, + stride_s1_c, + stride_sh1_b, + stride_sh1_c, + stride_g1_b, + stride_g1_c, + stride_i_b, + stride_i_l, + eps, + has_weight: tl.constexpr, + has_bias: tl.constexpr, + block_n: tl.constexpr, +): + row = tl.program_id(0) + cols = tl.arange(0, block_n) + mask = cols < inner_dim + + x = tl.load(x_ptr + row * stride_x_row + cols, mask=mask, other=0.0).to(tl.float32) + residual = tl.load(residual_ptr + row * stride_res_row + cols, mask=mask, other=0.0).to(tl.float32) + residual_gate = tl.load( + residual_gate_ptr + row * stride_res_gate_row + cols, + mask=mask, + other=0.0, + ).to(tl.float32) + + residual_out = residual + residual_gate * x + tl.store(residual_out_ptr + row * stride_res_out_row + cols, residual_out, mask=mask) + + mean = tl.sum(residual_out, axis=0) / inner_dim + xbar = tl.where(mask, residual_out - mean, 0.0) + var = tl.sum(xbar * xbar, axis=0) / inner_dim + x_hat = (residual_out - mean) * tl.rsqrt(var + eps) + + if has_weight: + weight = tl.load(weight_ptr + cols * stride_w, mask=mask, other=1.0).to(tl.float32) + x_hat = x_hat * weight + if has_bias: + bias = tl.load(bias_ptr + cols * stride_b, mask=mask, other=0.0).to(tl.float32) + x_hat = x_hat + bias + + batch_idx = row // seq_len + seq_idx = row % seq_len + idx = tl.load(index_ptr + batch_idx * stride_i_b + seq_idx * stride_i_l).to(tl.int1) + + scale0_ptrs = scale0_ptr + batch_idx * stride_s0_b + cols * stride_s0_c + shift0_ptrs = shift0_ptr + batch_idx * stride_sh0_b + cols * stride_sh0_c + gate0_ptrs = gate0_ptr + batch_idx * stride_g0_b + cols * stride_g0_c + scale1_ptrs = scale1_ptr + batch_idx * stride_s1_b + cols * stride_s1_c + shift1_ptrs = shift1_ptr + batch_idx * stride_sh1_b + cols * stride_sh1_c + gate1_ptrs = gate1_ptr + batch_idx * stride_g1_b + cols * stride_g1_c + + scale_ptrs = tl.where(idx, scale1_ptrs, scale0_ptrs) + shift_ptrs = tl.where(idx, shift1_ptrs, shift0_ptrs) + gate_ptrs = tl.where(idx, gate1_ptrs, gate0_ptrs) + + scale = tl.load(scale_ptrs, mask=mask, other=0.0).to(tl.float32) + shift = tl.load(shift_ptrs, mask=mask, other=0.0).to(tl.float32) + gate = tl.load(gate_ptrs, mask=mask, other=0.0) + + y = x_hat * (1.0 + scale) + shift + tl.store(output_ptr + row * stride_out_row + cols, y, mask=mask) + tl.store(gate_out_ptr + row * stride_gate_out_row + cols, gate, mask=mask) + + +def split_select01_mod_params(mod_params: torch.Tensor) -> tuple[torch.Tensor, ...]: + shift, scale, gate = mod_params.chunk(3, dim=-1) + actual_batch = shift.size(0) // 2 + if actual_batch == 0 or shift.size(0) != actual_batch * 2: + raise ValueError("select01 modulation expects mod_params batch dim to be 2 * batch") + return ( + scale[:actual_batch], + shift[:actual_batch], + gate[:actual_batch], + scale[actual_batch:], + shift[actual_batch:], + gate[actual_batch:], + ) + + +def select01_modulation_native( + mod_params: torch.Tensor, + index: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + scale0, shift0, gate0, scale1, shift1, gate1 = split_select01_mod_params(mod_params) + idx = index.to(dtype=torch.bool).unsqueeze(-1) + scale = torch.where(idx, scale1.unsqueeze(1), scale0.unsqueeze(1)) + shift = torch.where(idx, shift1.unsqueeze(1), shift0.unsqueeze(1)) + gate = torch.where(idx, gate1.unsqueeze(1), gate0.unsqueeze(1)) + return scale, shift, gate + + +def _validate_select01_inputs( + x: torch.Tensor, + index: torch.Tensor, + tensors: tuple[torch.Tensor, ...], +) -> None: + if x.dim() != 3: + raise ValueError("x must be 3D [B, L, C]") + if index.dim() != 2 or index.shape[1] != x.shape[1] or index.shape[0] not in (1, x.shape[0]): + raise ValueError("index must be 2D [B or 1, L] and match x batch/sequence") + batch_size, _, hidden_size = x.shape + for tensor in tensors: + if tensor.dim() != 2 or tensor.shape != (batch_size, hidden_size): + raise ValueError("scale/shift/gate tensors must be 2D [B, C]") + + +def _maybe_contiguous(x: torch.Tensor) -> torch.Tensor: + return x if x.is_contiguous() else x.contiguous() + + +def _launch_layernorm_select01( + x: torch.Tensor, + weight: torch.Tensor | None, + bias: torch.Tensor | None, + scale0: torch.Tensor, + shift0: torch.Tensor, + gate0: torch.Tensor, + scale1: torch.Tensor, + shift1: torch.Tensor, + gate1: torch.Tensor, + index: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + _validate_select01_inputs(x, index, (scale0, shift0, gate0, scale1, shift1, gate1)) + batch_size, seq_len, hidden_size = x.shape + output = torch.empty_like(x) + gate_out = torch.empty_like(x) + + x_2d = _maybe_contiguous(x).view(batch_size * seq_len, hidden_size) + output_2d = output.view(batch_size * seq_len, hidden_size) + gate_out_2d = gate_out.view(batch_size * seq_len, hidden_size) + index = _maybe_contiguous(index) + scale0 = _maybe_contiguous(scale0) + shift0 = _maybe_contiguous(shift0) + gate0 = _maybe_contiguous(gate0) + scale1 = _maybe_contiguous(scale1) + shift1 = _maybe_contiguous(shift1) + gate1 = _maybe_contiguous(gate1) + + weight_arg = _maybe_contiguous(weight) if weight is not None else x_2d + bias_arg = _maybe_contiguous(bias) if bias is not None else x_2d + block_n = min(65536 // x_2d.element_size(), triton.next_power_of_2(hidden_size)) + if hidden_size > block_n: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + + _layernorm_select01_kernel[(batch_size * seq_len,)]( + output_2d, + gate_out_2d, + x_2d, + weight_arg, + bias_arg, + scale0, + shift0, + gate0, + scale1, + shift1, + gate1, + index, + hidden_size, + seq_len, + x_2d.stride(0), + output_2d.stride(0), + gate_out_2d.stride(0), + weight_arg.stride(0) if weight is not None else 0, + bias_arg.stride(0) if bias is not None else 0, + scale0.stride(0), + scale0.stride(1), + shift0.stride(0), + shift0.stride(1), + gate0.stride(0), + gate0.stride(1), + scale1.stride(0), + scale1.stride(1), + shift1.stride(0), + shift1.stride(1), + gate1.stride(0), + gate1.stride(1), + index.stride(0) if index.shape[0] > 1 else 0, + index.stride(1), + eps, + has_weight=weight is not None, + has_bias=bias is not None, + block_n=block_n, + num_warps=4, + num_stages=4, + ) + return output, gate_out + + +def _launch_residual_layernorm_select01( + x: torch.Tensor, + residual: torch.Tensor, + residual_gate: torch.Tensor, + weight: torch.Tensor | None, + bias: torch.Tensor | None, + scale0: torch.Tensor, + shift0: torch.Tensor, + gate0: torch.Tensor, + scale1: torch.Tensor, + shift1: torch.Tensor, + gate1: torch.Tensor, + index: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + _validate_select01_inputs(x, index, (scale0, shift0, gate0, scale1, shift1, gate1)) + if residual.shape != x.shape: + raise ValueError("residual must have the same shape as x") + if residual_gate.shape != x.shape: + raise ValueError("residual_gate must have the same shape as x") + + batch_size, seq_len, hidden_size = x.shape + output = torch.empty_like(x) + residual_out = torch.empty_like(x) + gate_out = torch.empty_like(x) + + x_2d = _maybe_contiguous(x).view(batch_size * seq_len, hidden_size) + residual_2d = _maybe_contiguous(residual).view(batch_size * seq_len, hidden_size) + residual_gate_2d = _maybe_contiguous(residual_gate).view(batch_size * seq_len, hidden_size) + output_2d = output.view(batch_size * seq_len, hidden_size) + residual_out_2d = residual_out.view(batch_size * seq_len, hidden_size) + gate_out_2d = gate_out.view(batch_size * seq_len, hidden_size) + index = _maybe_contiguous(index) + scale0 = _maybe_contiguous(scale0) + shift0 = _maybe_contiguous(shift0) + gate0 = _maybe_contiguous(gate0) + scale1 = _maybe_contiguous(scale1) + shift1 = _maybe_contiguous(shift1) + gate1 = _maybe_contiguous(gate1) + + weight_arg = _maybe_contiguous(weight) if weight is not None else x_2d + bias_arg = _maybe_contiguous(bias) if bias is not None else x_2d + block_n = min(65536 // x_2d.element_size(), triton.next_power_of_2(hidden_size)) + if hidden_size > block_n: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + + _residual_layernorm_select01_kernel[(batch_size * seq_len,)]( + output_2d, + residual_out_2d, + gate_out_2d, + x_2d, + residual_2d, + residual_gate_2d, + weight_arg, + bias_arg, + scale0, + shift0, + gate0, + scale1, + shift1, + gate1, + index, + hidden_size, + seq_len, + x_2d.stride(0), + residual_2d.stride(0), + residual_gate_2d.stride(0), + output_2d.stride(0), + residual_out_2d.stride(0), + gate_out_2d.stride(0), + weight_arg.stride(0) if weight is not None else 0, + bias_arg.stride(0) if bias is not None else 0, + scale0.stride(0), + scale0.stride(1), + shift0.stride(0), + shift0.stride(1), + gate0.stride(0), + gate0.stride(1), + scale1.stride(0), + scale1.stride(1), + shift1.stride(0), + shift1.stride(1), + gate1.stride(0), + gate1.stride(1), + index.stride(0) if index.shape[0] > 1 else 0, + index.stride(1), + eps, + has_weight=weight is not None, + has_bias=bias is not None, + block_n=block_n, + num_warps=4, + num_stages=4, + ) + return output, residual_out, gate_out + + +def can_use_qwen_select01_triton(x: torch.Tensor) -> bool: + return HAS_TRITON and current_platform.is_cuda() and x.is_cuda + + +def fused_layernorm_select01( + x: torch.Tensor, + mod_params: torch.Tensor, + index: torch.Tensor, + eps: float, + weight: torch.Tensor | None = None, + bias: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + scale0, shift0, gate0, scale1, shift1, gate1 = split_select01_mod_params(mod_params) + is_compiling = torch.compiler.is_compiling() + if can_use_qwen_select01_triton(x) and not is_compiling: + return _launch_layernorm_select01( + x, + weight, + bias, + scale0, + shift0, + gate0, + scale1, + shift1, + gate1, + index, + eps, + ) + + scale, shift, gate = select01_modulation_native(mod_params, index) + out = F.layer_norm( + x.float(), + (x.shape[-1],), + weight=weight.float() if weight is not None else None, + bias=bias.float() if bias is not None else None, + eps=eps, + ).to(x.dtype) + return out * (1 + scale) + shift, gate + + +def fused_residual_layernorm_select01( + x: torch.Tensor, + residual: torch.Tensor, + residual_gate: torch.Tensor, + mod_params: torch.Tensor, + index: torch.Tensor, + eps: float, + weight: torch.Tensor | None = None, + bias: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + scale0, shift0, gate0, scale1, shift1, gate1 = split_select01_mod_params(mod_params) + is_compiling = torch.compiler.is_compiling() + if can_use_qwen_select01_triton(x) and not is_compiling: + return _launch_residual_layernorm_select01( + x, + residual, + residual_gate, + weight, + bias, + scale0, + shift0, + gate0, + scale1, + shift1, + gate1, + index, + eps, + ) + + scale, shift, gate = select01_modulation_native(mod_params, index) + residual_out = residual + residual_gate * x + out = F.layer_norm( + residual_out.float(), + (residual_out.shape[-1],), + weight=weight.float() if weight is not None else None, + bias=bias.float() if bias is not None else None, + eps=eps, + ).to(residual_out.dtype) + return out * (1 + scale) + shift, residual_out, gate 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 8e4dac3a3..11568f62c 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.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 from __future__ import annotations @@ -49,6 +49,12 @@ from vllm_omni.diffusion.layers.fused_qk_norm_rope import ( _fused_cuda_supported, fused_qk_norm_rope, ) +from vllm_omni.diffusion.layers.qwen_select01_modulation import ( + can_use_qwen_select01_triton, + fused_layernorm_select01, + fused_residual_layernorm_select01, + select01_modulation_native, +) from vllm_omni.diffusion.layers.rope import RotaryEmbedding logger = init_logger(__name__) @@ -855,41 +861,11 @@ class QwenImageTransformerBlock(nn.Module): self.zero_cond_t = zero_cond_t - def _modulate(self, mod_params, index=None): + def _modulate(self, mod_params): """Apply modulation to input 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 scale_result, shift_result, gate_result + return scale.unsqueeze(1), shift.unsqueeze(1), gate.unsqueeze(1) def forward( self, @@ -915,8 +891,22 @@ 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_scale1, img_shift1, img_gate1 = self._modulate(img_mod1, modulate_index) - img_modulated = self.img_norm1(hidden_states, img_scale1, img_shift1) + use_fused_select01 = modulate_index is not None and can_use_qwen_select01_triton(hidden_states) + if use_fused_select01: + img_modulated, img_gate1 = fused_layernorm_select01( + hidden_states, + img_mod1, + modulate_index, + self.img_norm1.eps, + self.img_norm1.layernorm.weight, + self.img_norm1.layernorm.bias, + ) + elif modulate_index is not None: + img_scale1, img_shift1, img_gate1 = select01_modulation_native(img_mod1, modulate_index) + img_modulated = self.img_norm1(hidden_states, img_scale1, img_shift1) + else: + img_scale1, img_shift1, img_gate1 = self._modulate(img_mod1) + img_modulated = self.img_norm1(hidden_states, img_scale1, img_shift1) # Process text stream - norm1 + modulation txt_scale1, txt_shift1, txt_gate1 = self._modulate(txt_mod1) @@ -941,12 +931,29 @@ class QwenImageTransformerBlock(nn.Module): img_attn_output, txt_attn_output = attn_output # Apply attention gates and add residual (like in Megatron) - hidden_states = hidden_states + img_gate1 * img_attn_output + hidden_states_before_attn = hidden_states + if not use_fused_select01: + hidden_states = hidden_states + img_gate1 * img_attn_output encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn_output # Process image stream - norm2 + MLP - img_scale2, img_shift2, img_gate2 = self._modulate(img_mod2, modulate_index) - img_modulated2 = self.img_norm2(hidden_states, img_scale2, img_shift2) + if use_fused_select01: + img_modulated2, hidden_states, img_gate2 = fused_residual_layernorm_select01( + img_attn_output, + hidden_states_before_attn, + img_gate1, + img_mod2, + modulate_index, + self.img_norm2.eps, + self.img_norm2.layernorm.weight, + self.img_norm2.layernorm.bias, + ) + elif modulate_index is not None: + img_scale2, img_shift2, img_gate2 = select01_modulation_native(img_mod2, modulate_index) + img_modulated2 = self.img_norm2(hidden_states, img_scale2, img_shift2) + else: + img_scale2, img_shift2, img_gate2 = self._modulate(img_mod2) + img_modulated2 = self.img_norm2(hidden_states, img_scale2, img_shift2) img_mlp_output = self.img_mlp(img_modulated2) hidden_states = hidden_states + img_gate2 * img_mlp_output