334 lines
12 KiB
Python
334 lines
12 KiB
Python
# 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())
|