diff --git a/benchmarks/diffusion/benchmark_qwen_image_qk_norm_rope.py b/benchmarks/diffusion/benchmark_qwen_image_qk_norm_rope.py new file mode 100644 index 000000000..ee69f7e59 --- /dev/null +++ b/benchmarks/diffusion/benchmark_qwen_image_qk_norm_rope.py @@ -0,0 +1,224 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Benchmark Qwen-Image Q/K RMSNorm + interleaved RoPE. + +Example: + + CUDA_VISIBLE_DEVICES=0 python benchmarks/diffusion/benchmark_qwen_image_qk_norm_rope.py \ + --seq-len 4096 --warmup 20 --iters 100 --include-compiled --profile +""" + +from __future__ import annotations + +import argparse +import json +from collections.abc import Callable + +import torch +import torch.nn.functional as F + +from vllm_omni.diffusion.layers.fused_qk_norm_rope import ( + _fused_cuda_supported, + fused_qk_norm_rope, +) +from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( + _apply_qwen_image_rotary_emb, +) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--seq-len", type=int, default=4096) + parser.add_argument("--heads", type=int, default=24) + parser.add_argument("--kv-heads", type=int, default=24) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--iters", type=int, default=100) + parser.add_argument("--dtype", choices=("bf16", "fp16", "fp32"), default="bf16") + parser.add_argument("--packed-qkv-view", action="store_true") + parser.add_argument("--include-compiled", action="store_true") + parser.add_argument("--profile", action="store_true") + return parser.parse_args() + + +def _dtype(name: str) -> torch.dtype: + return { + "bf16": torch.bfloat16, + "fp16": torch.float16, + "fp32": torch.float32, + }[name] + + +def _make_inputs(args: argparse.Namespace, dtype: torch.dtype, device: torch.device): + torch.manual_seed(2026) + if args.packed_qkv_view: + qkv_dim = (args.heads + args.kv_heads + args.kv_heads) * args.head_dim + qkv = torch.randn(args.batch, args.seq_len, qkv_dim, device=device, dtype=dtype) + q, k, _ = qkv.split( + [ + args.heads * args.head_dim, + args.kv_heads * args.head_dim, + args.kv_heads * args.head_dim, + ], + dim=-1, + ) + q = q.unflatten(-1, (args.heads, args.head_dim)) + k = k.unflatten(-1, (args.kv_heads, args.head_dim)) + else: + q = torch.randn(args.batch, args.seq_len, args.heads, args.head_dim, device=device, dtype=dtype) + k = torch.randn(args.batch, args.seq_len, args.kv_heads, args.head_dim, device=device, dtype=dtype) + + q_weight = torch.randn(args.head_dim, device=device, dtype=dtype) + k_weight = torch.randn(args.head_dim, device=device, dtype=dtype) + angles = torch.randn(args.seq_len, args.head_dim // 2, device=device, dtype=torch.float32) + freqs = torch.polar(torch.ones_like(angles), angles) + return q, k, q_weight, k_weight, freqs + + +def _native( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + freqs: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + q = F.rms_norm(q, (q.shape[-1],), q_weight, eps) + k = F.rms_norm(k, (k.shape[-1],), k_weight, eps) + return ( + _apply_qwen_image_rotary_emb(q, freqs), + _apply_qwen_image_rotary_emb(k, freqs), + ) + + +def _fused( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + freqs: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + batch, seq_len, num_heads, head_dim = q.shape + num_kv_heads = k.shape[2] + rope_table = torch.cat((freqs.real, freqs.imag), dim=-1) + rope_table = rope_table.unsqueeze(0).expand(batch, -1, -1).reshape(batch * seq_len, head_dim) + out_q, out_k = fused_qk_norm_rope( + q.reshape(batch * seq_len, num_heads, head_dim), + k.reshape(batch * seq_len, num_kv_heads, head_dim), + q_weight, + k_weight, + rope_table, + eps, + interleaved=True, + ) + return ( + out_q.reshape(batch, seq_len, num_heads, head_dim), + out_k.reshape(batch, seq_len, num_kv_heads, head_dim), + ) + + +def _measure(fn: Callable[[], tuple[torch.Tensor, torch.Tensor]], warmup: int, iters: int) -> list[float]: + for _ in range(warmup): + fn() + torch.accelerator.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + samples = [] + for _ in range(iters): + start.record() + fn() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end)) + return samples + + +def _stats(samples: list[float]) -> dict[str, float]: + samples = sorted(samples) + return { + "median_ms": samples[len(samples) // 2], + "p90_ms": samples[int(len(samples) * 0.9)], + "mean_ms": sum(samples) / len(samples), + } + + +def _profile(fn: Callable[[], tuple[torch.Tensor, torch.Tensor]]) -> str: + from torch.profiler import ProfilerActivity, profile + + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + fn() + torch.accelerator.synchronize() + return prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=15) + + +def main() -> None: + args = _parse_args() + if not torch.cuda.is_available(): + raise RuntimeError("This benchmark requires CUDA") + device = torch.device("cuda") + dtype = _dtype(args.dtype) + eps = 1e-6 + q, k, q_weight, k_weight, freqs = _make_inputs(args, dtype, device) + + def native(): + return _native(q, k, q_weight, k_weight, freqs, eps) + + def fused(): + return _fused(q, k, q_weight, k_weight, freqs, eps) + + results: dict[str, object] = { + "device": torch.cuda.get_device_name(), + "batch": args.batch, + "seq_len": args.seq_len, + "heads": args.heads, + "kv_heads": args.kv_heads, + "head_dim": args.head_dim, + "dtype": args.dtype, + "packed_qkv_view": args.packed_qkv_view, + "fused_fast_path_supported": _fused_cuda_supported( + q.reshape(args.batch * args.seq_len, args.heads, args.head_dim), + k.reshape(args.batch * args.seq_len, args.kv_heads, args.head_dim), + args.head_dim, + args.head_dim, + interleaved=True, + ), + } + + native_samples = _measure(native, args.warmup, args.iters) + fused_samples = _measure(fused, args.warmup, args.iters) + expected_q, expected_k = native() + actual_q, actual_k = fused() + torch.accelerator.synchronize() + + native_stats = _stats(native_samples) + fused_stats = _stats(fused_samples) + results["native"] = native_stats + results["fused"] = fused_stats + results["speedup_vs_native"] = native_stats["median_ms"] / fused_stats["median_ms"] + results["max_abs_q"] = (actual_q.float() - expected_q.float()).abs().max().item() + results["max_abs_k"] = (actual_k.float() - expected_k.float()).abs().max().item() + results["mean_abs_q"] = (actual_q.float() - expected_q.float()).abs().mean().item() + results["mean_abs_k"] = (actual_k.float() - expected_k.float()).abs().mean().item() + + if args.include_compiled: + compiled_native = torch.compile(native, dynamic=True, fullgraph=True) + compiled_fused = torch.compile(fused, dynamic=True, fullgraph=True) + compiled_native_stats = _stats(_measure(compiled_native, args.warmup, args.iters)) + compiled_fused_stats = _stats(_measure(compiled_fused, args.warmup, args.iters)) + results["compiled_native"] = compiled_native_stats + results["compiled_fused"] = compiled_fused_stats + results["compiled_fused_speedup_vs_compiled_native"] = ( + compiled_native_stats["median_ms"] / compiled_fused_stats["median_ms"] + ) + + print(json.dumps(results, indent=2)) + if args.profile: + print("\n# fused profile") + print(_profile(fused)) + + +if __name__ == "__main__": + main() diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_fused_qk_norm_rope.py b/tests/diffusion/models/qwen_image/test_qwen_image_fused_qk_norm_rope.py new file mode 100644 index 000000000..2f10ea33e --- /dev/null +++ b/tests/diffusion/models/qwen_image/test_qwen_image_fused_qk_norm_rope.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from dataclasses import dataclass + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.diffusion.layers.rope import RotaryEmbedding +from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( + _apply_qwen_image_rotary_emb, + _qwen_image_qk_norm_rope, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] + +BATCH = 2 +Q_HEADS = 24 +K_HEADS = 24 +HEAD_DIM = 128 +EPS = 1e-6 +SEED = 2026 + + +@dataclass(frozen=True) +class QwenImageQKInput: + q: torch.Tensor + k: torch.Tensor + norm_q: nn.RMSNorm + norm_k: nn.RMSNorm + freqs: torch.Tensor + + +def _make_input( + *, + seq_len: int, + dtype: torch.dtype, + device: torch.device, + packed_qkv_view: bool, +) -> QwenImageQKInput: + gen = torch.Generator(device=device) + gen.manual_seed(SEED + seq_len) + + if packed_qkv_view: + qkv_dim = (Q_HEADS + K_HEADS + K_HEADS) * HEAD_DIM + qkv = torch.randn(BATCH, seq_len, qkv_dim, device=device, dtype=dtype, generator=gen) + q, k, _ = qkv.split( + [ + Q_HEADS * HEAD_DIM, + K_HEADS * HEAD_DIM, + K_HEADS * HEAD_DIM, + ], + dim=-1, + ) + q = q.unflatten(-1, (Q_HEADS, HEAD_DIM)) + k = k.unflatten(-1, (K_HEADS, HEAD_DIM)) + else: + q = torch.randn(BATCH, seq_len, Q_HEADS, HEAD_DIM, device=device, dtype=dtype, generator=gen) + k = torch.randn(BATCH, seq_len, K_HEADS, HEAD_DIM, device=device, dtype=dtype, generator=gen) + + norm_q = nn.RMSNorm(HEAD_DIM, eps=EPS, device=device, dtype=dtype) + norm_k = nn.RMSNorm(HEAD_DIM, eps=EPS, device=device, dtype=dtype) + norm_q.requires_grad_(False) + norm_k.requires_grad_(False) + with torch.no_grad(): + norm_q.weight.copy_(torch.randn(HEAD_DIM, device=device, dtype=dtype, generator=gen)) + norm_k.weight.copy_(torch.randn(HEAD_DIM, device=device, dtype=dtype, generator=gen)) + angles = torch.randn(seq_len, HEAD_DIM // 2, device=device, dtype=torch.float32, generator=gen) + freqs = torch.polar(torch.ones_like(angles), angles) + return QwenImageQKInput(q=q, k=k, norm_q=norm_q, norm_k=norm_k, freqs=freqs) + + +def _reference(data: QwenImageQKInput) -> tuple[torch.Tensor, torch.Tensor]: + q = data.norm_q(data.q) + k = data.norm_k(data.k) + if data.q.device.type == "cuda": + return ( + _apply_qwen_image_rotary_emb(q, data.freqs), + _apply_qwen_image_rotary_emb(k, data.freqs), + ) + + rope = RotaryEmbedding(is_neox_style=False) + cos = data.freqs.real.to(data.q.dtype) + sin = data.freqs.imag.to(data.q.dtype) + return rope(q, cos, sin), rope(k, cos, sin) + + +def _run(data: QwenImageQKInput) -> tuple[torch.Tensor, torch.Tensor]: + return _qwen_image_qk_norm_rope( + data.q, + data.k, + data.norm_q, + data.norm_k, + data.freqs, + RotaryEmbedding(is_neox_style=False), + EPS, + ) + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_qwen_image_qk_norm_rope_cuda_fp32_fallback_matches_reference(): + data = _make_input( + seq_len=7, + dtype=torch.float32, + device=torch.device("cuda:0"), + packed_qkv_view=True, + ) + + actual_q, actual_k = _run(data) + expected_q, expected_k = _reference(data) + + torch.testing.assert_close(actual_q, expected_q, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(actual_k, expected_k, atol=1e-5, rtol=1e-5) + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_qwen_image_qk_norm_rope_cuda_fp16_fallback_matches_reference(): + data = _make_input( + seq_len=7, + dtype=torch.float16, + device=torch.device("cuda:0"), + packed_qkv_view=True, + ) + + actual_q, actual_k = _run(data) + expected_q, expected_k = _reference(data) + + torch.testing.assert_close(actual_q, expected_q, atol=1e-3, rtol=1e-3) + torch.testing.assert_close(actual_k, expected_k, atol=1e-3, rtol=1e-3) + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("seq_len", [1, 7, 257, 1024]) +@pytest.mark.parametrize("packed_qkv_view", [False, True]) +def test_qwen_image_fused_qk_norm_rope_cuda_matches_fp32_rope_reference( + seq_len: int, + packed_qkv_view: bool, +): + data = _make_input( + seq_len=seq_len, + dtype=torch.bfloat16, + device=torch.device("cuda:0"), + packed_qkv_view=packed_qkv_view, + ) + + actual_q, actual_k = _run(data) + expected_q, expected_k = _reference(data) + + torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02) + torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02) + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_qwen_image_fused_qk_norm_rope_torch_compile_fullgraph_capture(): + data = _make_input( + seq_len=257, + dtype=torch.bfloat16, + device=torch.device("cuda:0"), + packed_qkv_view=True, + ) + rope = RotaryEmbedding(is_neox_style=False) + + def fn( + q: torch.Tensor, + k: torch.Tensor, + freqs: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + return _qwen_image_qk_norm_rope( + q, + k, + data.norm_q, + data.norm_k, + freqs, + rope, + EPS, + ) + + compiled_fn = torch.compile(fn, dynamic=True, fullgraph=True) + expected_q, expected_k = fn(data.q, data.k, data.freqs) + actual_q, actual_k = compiled_fn(data.q, data.k, data.freqs) + + torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02) + torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02) 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 696728b90..8e4dac3a3 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py @@ -45,6 +45,10 @@ from vllm_omni.diffusion.distributed.sp_plan import ( ) from vllm_omni.diffusion.forward_context import get_forward_context from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm +from vllm_omni.diffusion.layers.fused_qk_norm_rope import ( + _fused_cuda_supported, + fused_qk_norm_rope, +) from vllm_omni.diffusion.layers.rope import RotaryEmbedding logger = init_logger(__name__) @@ -61,6 +65,53 @@ def _apply_qwen_image_rotary_emb(x: torch.Tensor, freqs: torch.Tensor) -> torch. return torch.view_as_real(paired * freqs.unsqueeze(1)).flatten(3).to(x.dtype) +def _qwen_image_qk_norm_rope( + q: torch.Tensor, + k: torch.Tensor, + norm_q: nn.Module, + norm_k: nn.Module, + freqs: torch.Tensor, + rope: RotaryEmbedding, + eps: float, + *, + use_fused: bool = True, +) -> tuple[torch.Tensor, torch.Tensor]: + head_dim = q.shape[-1] + rotary_dim = freqs.shape[-1] * 2 + if use_fused and _fused_cuda_supported(q, k, head_dim, rotary_dim, interleaved=True): + batch, seq_len, num_heads, _ = q.shape + num_kv_heads = k.shape[2] + rope_table = torch.cat((freqs.real, freqs.imag), dim=-1) + rope_table = rope_table.unsqueeze(0).expand(batch, -1, -1).reshape(batch * seq_len, rotary_dim) + fused_q, fused_k = fused_qk_norm_rope( + q.reshape(batch * seq_len, num_heads, head_dim), + k.reshape(batch * seq_len, num_kv_heads, head_dim), + norm_q.weight, + norm_k.weight, + rope_table, + eps, + interleaved=True, + ) + return ( + fused_q.reshape(batch, seq_len, num_heads, head_dim), + fused_k.reshape(batch, seq_len, num_kv_heads, head_dim), + ) + + q = norm_q(q) + k = norm_k(k) + if q.device.type == "cuda": + return ( + _apply_qwen_image_rotary_emb(q, freqs), + _apply_qwen_image_rotary_emb(k, freqs), + ) + + # Retain the platform-specific kernels on other accelerators, which may + # not support complex tensors. + cos = torch.real(freqs).to(q.dtype) + sin = torch.imag(freqs).to(q.dtype) + return rope(q, cos, sin), rope(k, cos, sin) + + def _normalize_qwen_image_weight_name(name: str) -> str: name = name.removeprefix("transformer.") if ".to_out.0." in name: @@ -648,28 +699,25 @@ class QwenImageCrossAttention(nn.Module): txt_key = txt_key.unflatten(-1, (self.add_kv_num_heads, self.head_dim)) txt_value = txt_value.unflatten(-1, (self.add_kv_num_heads, self.head_dim)) - img_query = self.norm_q(img_query) - img_key = self.norm_k(img_key) - txt_query = self.norm_added_q(txt_query) - txt_key = self.norm_added_k(txt_key) - - if img_query.device.type == "cuda": - img_query = _apply_qwen_image_rotary_emb(img_query, vid_freqs) - img_key = _apply_qwen_image_rotary_emb(img_key, vid_freqs) - txt_query = _apply_qwen_image_rotary_emb(txt_query, txt_freqs) - txt_key = _apply_qwen_image_rotary_emb(txt_key, txt_freqs) - else: - # Retain the platform-specific kernels on other accelerators, - # which may not support complex tensors. - img_cos = torch.real(vid_freqs).to(img_query.dtype) - img_sin = torch.imag(vid_freqs).to(img_query.dtype) - txt_cos = torch.real(txt_freqs).to(txt_query.dtype) - txt_sin = torch.imag(txt_freqs).to(txt_query.dtype) - - img_query = self.rope(img_query, img_cos, img_sin) - img_key = self.rope(img_key, img_cos, img_sin) - txt_query = self.rope(txt_query, txt_cos, txt_sin) - txt_key = self.rope(txt_key, txt_cos, txt_sin) + img_query, img_key = _qwen_image_qk_norm_rope( + img_query, + img_key, + self.norm_q, + self.norm_k, + vid_freqs, + self.rope, + self.eps, + use_fused=self.qk_norm, + ) + txt_query, txt_key = _qwen_image_qk_norm_rope( + txt_query, + txt_key, + self.norm_added_q, + self.norm_added_k, + txt_freqs, + self.rope, + self.eps, + ) seq_len_txt = encoder_hidden_states.shape[1] joint_query = torch.cat([txt_query, img_query], dim=1)