88 lines
3.7 KiB
Python
88 lines
3.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Compare native and fused MAGI-2 SwiGLU7 operators on CUDA or MUSA.
|
|
|
|
Run from the checkout root in an environment with PyTorch and the selected
|
|
backend dependencies. torchada is loaded automatically only on MUSA::
|
|
|
|
python -m benchmarks.kernels.benchmark_magi2_swiglu7 \
|
|
--tokens 1 4 16 64 129 4096 --intermediate-size 1280 --num-tests 30
|
|
|
|
The released MAGI-2 config supplies the 1280 intermediate width. Token counts
|
|
are a synthetic sweep; pass per-rank sizes from the target trace for production
|
|
comparisons. Input is contiguous interleaved gate/up [tokens, 2 * intermediate].
|
|
Both paths evaluate FP32 arithmetic with one cast to the requested output dtype.
|
|
This measures eager operator latency, not compiled serving performance.
|
|
No CUDA/MUSA device (or no PyTorch) produces a JSON skip record and exit status zero.
|
|
"""
|
|
|
|
import argparse
|
|
from functools import partial
|
|
|
|
from benchmarks.kernels._magi2_musa_bench import (
|
|
add_arguments,
|
|
check_pair,
|
|
emit,
|
|
load_musa,
|
|
measure_pair,
|
|
report_environment,
|
|
)
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
add_arguments(parser)
|
|
parser.add_argument("--intermediate-size", type=int, default=1280)
|
|
parser.add_argument("--out-dtype", choices=["bfloat16", "float16", "float32"], default=None)
|
|
args = parser.parse_args()
|
|
if args.intermediate_size <= 0:
|
|
parser.error("intermediate-size must be positive")
|
|
torch = load_musa(args, parser)
|
|
if torch is None:
|
|
return
|
|
cuda = getattr(torch, "cuda")
|
|
from vllm_omni.diffusion.layers import swiglu7
|
|
|
|
report_environment(torch, args, swiglu7)
|
|
dtype = getattr(torch, args.dtype)
|
|
out_dtype = getattr(torch, args.out_dtype or args.dtype)
|
|
device = torch.device("cuda")
|
|
rtol = {torch.bfloat16: 1.6e-2, torch.float16: 2e-3, torch.float32: 1e-5}[out_dtype]
|
|
atol = 1e-5 if out_dtype == torch.float32 else 2e-3
|
|
op = swiglu7.SwiGLU7()
|
|
with torch.inference_mode():
|
|
for tokens in args.tokens:
|
|
# Values outside both clamp boundaries exercise the released formula.
|
|
x = (torch.randn(tokens, 2 * args.intermediate_size, device=device) * 10).to(dtype)
|
|
reference = partial(op.forward_native, x, out_dtype=out_dtype)
|
|
fused = partial(op, x, out_dtype=out_dtype)
|
|
parity = check_pair(torch, reference, fused, swiglu7, ("_swiglu7_kernel",), rtol=rtol, atol=atol)
|
|
timing = measure_pair(torch, reference, fused, args)
|
|
output_elements = tokens * args.intermediate_size
|
|
minimum_io_bytes = (
|
|
x.numel() * x.element_size() + output_elements * torch.empty((), dtype=out_dtype).element_size()
|
|
)
|
|
emit(
|
|
status="pass",
|
|
op="swiglu7",
|
|
input_shape=list(x.shape),
|
|
input_stride=list(x.stride()),
|
|
dtype=args.dtype,
|
|
out_dtype=args.out_dtype or args.dtype,
|
|
alpha=1.702,
|
|
limit=7.0,
|
|
minimum_io_bytes=minimum_io_bytes,
|
|
# Six scalar arithmetic operations and one exp per output;
|
|
# clamps and transcendental costs are not modeled as FLOPs.
|
|
scalar_flops_excluding_exp_and_clamps=6 * output_elements,
|
|
effective_fused_gbps=minimum_io_bytes / timing["fused"]["p50_us"] / 1000,
|
|
**parity,
|
|
**timing,
|
|
)
|
|
del x, reference, fused
|
|
cuda.empty_cache()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|