vllm-omni/benchmarks/diffusion/benchmark_qwen_image_qk_norm_rope.py

# 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.layers.rope import RotaryEmbedding


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)
    rope = RotaryEmbedding(is_neox_style=False)
    cos = freqs.real.to(q.dtype)
    sin = freqs.imag.to(q.dtype)
    return rope(q, cos, sin), rope(k, cos, sin)


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