[Diffusion] Fuse Qwen-Image-Edit select01 modulation for CUDA (#5921)

Signed-off-by: dongbo910220 <1275604947@qq.com>

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())

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)

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

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

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