[Model] Fuse MammothModa2 QK norm and RoPE (#7969)
Signed-off-by: levius <2114377220@qq.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
This commit is contained in:
SunandOpenAI Codex authored and GitHub committed
1 parent
855254d4af
commit
3f20252151
5 files changed
+801
-21
No files matched your search
@@ -0,0 +1,262 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Benchmark MammothModa2's Q/K RMSNorm + adjacent-pair real RoPE.
|
||||
|
||||
Example:
|
||||
|
||||
CUDA_VISIBLE_DEVICES=0 python \
|
||||
benchmarks/diffusion/benchmark_mammoth_moda2_qk_norm_rope.py \
|
||||
--tokens 512 4173 8346 --warmup 50 --iters 200
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import importlib.metadata
|
||||
import json
|
||||
import os
|
||||
import statistics
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.layers.fused_qk_norm_rope import (
|
||||
_fused_cuda_supported,
|
||||
fused_qk_norm_rope,
|
||||
)
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
from vllm_omni.diffusion.models.mammoth_moda2.rope_real import (
|
||||
apply_real_rotary_emb,
|
||||
)
|
||||
|
||||
_HEAD_DIM = 120
|
||||
_Q_HEADS = 21
|
||||
_KV_HEADS = 7
|
||||
_EPS = 1e-5
|
||||
|
||||
|
||||
def _parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--tokens", nargs="+", type=int, default=[512, 4173, 8346])
|
||||
parser.add_argument("--warmup", type=int, default=50)
|
||||
parser.add_argument("--iters", type=int, default=200)
|
||||
parser.add_argument("--rounds", type=int, default=3)
|
||||
parser.add_argument("--attention-sequences", nargs="*", type=int, default=[])
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _version(distribution: str) -> str | None:
|
||||
try:
|
||||
return importlib.metadata.version(distribution)
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
return None
|
||||
|
||||
|
||||
def _measure(
|
||||
fn: Callable[[], object],
|
||||
warmup: int,
|
||||
iters: int,
|
||||
) -> dict[str, float]:
|
||||
for _ in range(warmup):
|
||||
fn()
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
torch.accelerator.reset_peak_memory_stats()
|
||||
baseline_bytes = torch.accelerator.memory_allocated()
|
||||
samples = []
|
||||
for _ in range(iters):
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
fn()
|
||||
end.record()
|
||||
end.synchronize()
|
||||
samples.append(float(start.elapsed_time(end)))
|
||||
|
||||
ordered = sorted(samples)
|
||||
return {
|
||||
"median_ms": statistics.median(ordered),
|
||||
"mean_ms": statistics.mean(ordered),
|
||||
"stddev_ms": statistics.pstdev(ordered),
|
||||
"min_ms": ordered[0],
|
||||
"max_ms": ordered[-1],
|
||||
"p90_ms": ordered[min(len(ordered) - 1, int(len(ordered) * 0.9))],
|
||||
"extra_peak_mib": (torch.accelerator.max_memory_allocated() - baseline_bytes) / 2**20,
|
||||
}
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def _run_shape(tokens: int, warmup: int, iters: int, *, fused_first: bool = False) -> dict[str, object]:
|
||||
dtype = torch.bfloat16
|
||||
device = torch.device("cuda")
|
||||
torch.manual_seed(42)
|
||||
query = torch.randn(tokens, _Q_HEADS, _HEAD_DIM, device=device, dtype=dtype)
|
||||
key = torch.randn(tokens, _KV_HEADS, _HEAD_DIM, device=device, dtype=dtype)
|
||||
norm_q = RMSNorm(_HEAD_DIM, eps=_EPS).to(device=device, dtype=dtype)
|
||||
norm_k = RMSNorm(_HEAD_DIM, eps=_EPS).to(device=device, dtype=dtype)
|
||||
norm_q.weight.normal_(mean=1.0, std=0.2)
|
||||
norm_k.weight.normal_(mean=1.0, std=0.2)
|
||||
angles = torch.rand(tokens, _HEAD_DIM // 2, device=device) * (2 * torch.pi)
|
||||
cos = angles.cos().repeat_interleave(2, dim=-1).to(dtype)
|
||||
sin = angles.sin().repeat_interleave(2, dim=-1).to(dtype)
|
||||
|
||||
def native() -> tuple[torch.Tensor, torch.Tensor]:
|
||||
q_out = apply_real_rotary_emb(norm_q(query).unsqueeze(0), cos.unsqueeze(0), sin.unsqueeze(0)).squeeze(0)
|
||||
k_out = apply_real_rotary_emb(norm_k(key).unsqueeze(0), cos.unsqueeze(0), sin.unsqueeze(0)).squeeze(0)
|
||||
return q_out, k_out
|
||||
|
||||
def fused() -> tuple[torch.Tensor, torch.Tensor]:
|
||||
rope_table = torch.cat((cos[..., 0::2], sin[..., 0::2]), dim=-1)
|
||||
return fused_qk_norm_rope(
|
||||
query,
|
||||
key,
|
||||
norm_q.weight,
|
||||
norm_k.weight,
|
||||
rope_table,
|
||||
_EPS,
|
||||
head_dim=_HEAD_DIM,
|
||||
rotary_dim=_HEAD_DIM,
|
||||
interleaved=True,
|
||||
)
|
||||
|
||||
expected_q, expected_k = native()
|
||||
actual_q, actual_k = fused()
|
||||
errors = {
|
||||
"q_max_abs": (actual_q.float() - expected_q.float()).abs().max().item(),
|
||||
"q_mean_abs": (actual_q.float() - expected_q.float()).abs().mean().item(),
|
||||
"k_max_abs": (actual_k.float() - expected_k.float()).abs().max().item(),
|
||||
"k_mean_abs": (actual_k.float() - expected_k.float()).abs().mean().item(),
|
||||
}
|
||||
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)
|
||||
|
||||
functions = {"native": native, "fused": fused}
|
||||
order = ("fused", "native") if fused_first else ("native", "fused")
|
||||
measured = {name: _measure(functions[name], warmup, iters) for name in order}
|
||||
native_stats, fused_stats = measured["native"], measured["fused"]
|
||||
return {
|
||||
"tokens": tokens,
|
||||
"native": native_stats,
|
||||
"fused_including_rope_pack": fused_stats,
|
||||
"median_speedup": native_stats["median_ms"] / fused_stats["median_ms"],
|
||||
"errors": errors,
|
||||
}
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def _run_attention(seq: int, warmup: int, iters: int, *, fused_first: bool = False) -> dict[str, object]:
|
||||
from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import TransformerBlock
|
||||
|
||||
torch.manual_seed(42)
|
||||
block = (
|
||||
TransformerBlock(
|
||||
2520,
|
||||
_Q_HEADS,
|
||||
_KV_HEADS,
|
||||
multiple_of=256,
|
||||
ffn_dim_multiplier=1.0,
|
||||
norm_eps=_EPS,
|
||||
rope_repeats_pairs=True,
|
||||
)
|
||||
.to(device="cuda", dtype=torch.bfloat16)
|
||||
.eval()
|
||||
)
|
||||
block.attn.norm_q.weight.normal_(1.0, 0.2)
|
||||
block.attn.norm_k.weight.normal_(1.0, 0.2)
|
||||
hidden = torch.randn(2, seq, 2520, device="cuda", dtype=torch.bfloat16)
|
||||
mask = torch.ones(2, seq, device="cuda", dtype=torch.bool)
|
||||
mask[0, seq - min(300, seq // 2) :] = False
|
||||
mask[1, seq - min(17, seq // 2) :] = False
|
||||
angles = torch.rand(2, seq, _HEAD_DIM // 2, device="cuda") * (2 * torch.pi)
|
||||
# The production producer supplies FP32 tables; each example has its own positions.
|
||||
rotary = angles.cos().repeat_interleave(2, dim=-1), angles.sin().repeat_interleave(2, dim=-1)
|
||||
env = "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS"
|
||||
previous = os.environ.get(env)
|
||||
|
||||
def run():
|
||||
return block.attn(hidden, hidden, attention_mask=mask, image_rotary_emb=rotary)
|
||||
|
||||
try:
|
||||
outputs = {}
|
||||
timings = {}
|
||||
order = ("fused", "native") if fused_first else ("native", "fused")
|
||||
for name in order:
|
||||
os.environ[env] = "0" if name == "fused" else str(2 * seq + 1)
|
||||
outputs[name] = run()
|
||||
timings[name] = _measure(run, warmup, iters)
|
||||
diff = (outputs["native"].float() - outputs["fused"].float()).abs()[mask]
|
||||
max_diff, mean_diff = diff.max().item(), diff.mean().item()
|
||||
assert max_diff < 2e-2 and mean_diff < 1e-3, (max_diff, mean_diff)
|
||||
assert torch.count_nonzero(outputs["fused"][~mask]) == 0
|
||||
return {
|
||||
"sequence": seq,
|
||||
"batch": 2,
|
||||
"backend": block.attn.omni_attn.attn_backend.__name__,
|
||||
"native": timings["native"],
|
||||
"fused": timings["fused"],
|
||||
"median_speedup": timings["native"]["median_ms"] / timings["fused"]["median_ms"],
|
||||
"max_abs_diff": max_diff,
|
||||
"mean_abs_diff": mean_diff,
|
||||
}
|
||||
finally:
|
||||
if previous is None:
|
||||
os.environ.pop(env, None)
|
||||
else:
|
||||
os.environ[env] = previous
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = _parse_args()
|
||||
if (
|
||||
args.rounds < 1
|
||||
or args.iters < 1
|
||||
or args.warmup < 0
|
||||
or any(value < 1 for value in [*args.tokens, *args.attention_sequences])
|
||||
):
|
||||
raise ValueError("rounds, iters and sequence lengths must be positive; warmup must be nonnegative")
|
||||
if not torch.cuda.is_available():
|
||||
raise RuntimeError("This benchmark requires CUDA")
|
||||
|
||||
probe_q = torch.empty(1, _Q_HEADS, _HEAD_DIM, device="cuda", dtype=torch.bfloat16)
|
||||
probe_k = torch.empty(1, _KV_HEADS, _HEAD_DIM, device="cuda", dtype=torch.bfloat16)
|
||||
if not _fused_cuda_supported(probe_q, probe_k, _HEAD_DIM, _HEAD_DIM, interleaved=True):
|
||||
raise RuntimeError("The fused CUDA QK norm/RoPE path is unavailable")
|
||||
|
||||
result = {
|
||||
"environment": {
|
||||
"gpu": torch.cuda.get_device_name(),
|
||||
"torch": torch.__version__,
|
||||
"cuda": torch.version.cuda,
|
||||
"triton": _version("triton"),
|
||||
"vllm": _version("vllm"),
|
||||
"vllm_omni": _version("vllm-omni"),
|
||||
},
|
||||
"config": {
|
||||
"dtype": "bfloat16",
|
||||
"q_heads": _Q_HEADS,
|
||||
"kv_heads": _KV_HEADS,
|
||||
"head_dim": _HEAD_DIM,
|
||||
"warmup": args.warmup,
|
||||
"iters": args.iters,
|
||||
"rounds": args.rounds,
|
||||
"inference_mode": True,
|
||||
"alternating_order": True,
|
||||
"norm_weights": "normal(mean=1.0, std=0.2)",
|
||||
},
|
||||
"results": [
|
||||
{"round": round_id, **_run_shape(tokens, args.warmup, args.iters, fused_first=bool(round_id % 2))}
|
||||
for round_id in range(args.rounds)
|
||||
for tokens in args.tokens
|
||||
],
|
||||
"attention_results": [
|
||||
{"round": round_id, **_run_attention(seq, args.warmup, args.iters, fused_first=bool(round_id % 2))}
|
||||
for round_id in range(args.rounds)
|
||||
for seq in args.attention_sequences
|
||||
],
|
||||
}
|
||||
print(json.dumps(result, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in new issue
Block a user