195 lines
7.7 KiB
Python
195 lines
7.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Cold-cache torch.cuda timing and correctness gates for MAGI-2 microbenchmarks."""
|
|
|
|
import argparse
|
|
import importlib
|
|
import importlib.metadata
|
|
import json
|
|
import logging
|
|
import math
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from collections.abc import Callable
|
|
from pathlib import Path
|
|
from types import ModuleType
|
|
from typing import Any
|
|
|
|
LOGGER = logging.getLogger(__name__)
|
|
|
|
|
|
def emit(**fields: Any) -> None:
|
|
LOGGER.info(json.dumps(fields, sort_keys=True))
|
|
|
|
|
|
def add_arguments(parser: argparse.ArgumentParser) -> None:
|
|
parser.add_argument("--tokens", type=int, nargs="+", default=[1, 4, 16, 64, 129, 4096])
|
|
parser.add_argument("--dtype", choices=["bfloat16", "float16", "float32"], default="bfloat16")
|
|
parser.add_argument("--num-tests", type=int, default=30, help="Measured samples per implementation.")
|
|
parser.add_argument("--warmup", type=int, default=10, help="Interleaved warmup calls before sampling.")
|
|
parser.add_argument("--seed", type=int, default=42)
|
|
|
|
|
|
def load_musa(args: argparse.Namespace, parser: argparse.ArgumentParser) -> Any:
|
|
if min(*args.tokens, args.num_tests) <= 0 or args.warmup < 0:
|
|
parser.error("tokens and num-tests must be positive; warmup must be nonnegative")
|
|
logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stdout)
|
|
torch = importlib.import_module("torch")
|
|
cuda = getattr(torch, "cuda")
|
|
# torchada is needed only on MUSA; after it is loaded, use torch.cuda for
|
|
# both CUDA and MUSA as documented by torchada.
|
|
if hasattr(torch.version, "musa") and torch.version.musa is not None:
|
|
try:
|
|
importlib.import_module("torchada")
|
|
except ModuleNotFoundError as error:
|
|
if error.name != "torchada":
|
|
raise
|
|
has_gpu = cuda.is_available() or (
|
|
hasattr(torch.version, "musa") and torch.version.musa is not None and torch.musa.is_available()
|
|
)
|
|
if not has_gpu:
|
|
emit(status="skip", reason="CUDA/MUSA device is unavailable", torch=torch.__version__)
|
|
return None
|
|
torch.manual_seed(args.seed)
|
|
cuda.manual_seed_all(args.seed)
|
|
return torch
|
|
|
|
|
|
def report_environment(torch: Any, args: argparse.Namespace, module: ModuleType) -> None:
|
|
root = Path(__file__).resolve().parents[2]
|
|
packages = {}
|
|
for name in ("torch", "torch_musa", "torchada", "triton", "vllm", "vllm-omni"):
|
|
try:
|
|
packages[name] = importlib.metadata.version(name)
|
|
except importlib.metadata.PackageNotFoundError:
|
|
packages[name] = None
|
|
cuda = getattr(torch, "cuda")
|
|
properties = cuda.get_device_properties(cuda.current_device())
|
|
try:
|
|
source_sha = subprocess.check_output(["git", "-C", str(root), "rev-parse", "HEAD"], text=True).strip()
|
|
source_dirty = bool(subprocess.check_output(["git", "-C", str(root), "status", "--porcelain"], text=True))
|
|
except (subprocess.CalledProcessError, FileNotFoundError):
|
|
source_sha = os.environ.get("MAGI2_BENCH_SOURCE_SHA", "unknown")
|
|
source_dirty = None
|
|
emit(
|
|
status="environment",
|
|
command=sys.argv,
|
|
source_root=str(root),
|
|
source_sha=source_sha,
|
|
source_dirty=source_dirty,
|
|
kernel_source=module.__file__,
|
|
device=cuda.get_device_name(cuda.current_device()),
|
|
sm_count=getattr(properties, "multi_processor_count", None),
|
|
packages=packages,
|
|
env={name: os.environ.get(name) for name in ("MUSA_VISIBLE_DEVICES", "TORCHDYNAMO_DISABLE")},
|
|
config=vars(args),
|
|
mode="eager operator diagnostic; no compile or graph capture",
|
|
comparison="forward_native versus the default CustomOp dispatch",
|
|
timing="torch.cuda.Event; alternating native/fused calls; vanilla 512 MiB flush outside timed regions",
|
|
timing_scope="whole callable, including device launch gaps; not isolated kernel duration",
|
|
l2_flush=True,
|
|
l2_flush_bytes=512 * 1024 * 1024,
|
|
)
|
|
|
|
|
|
class _LaunchProbe:
|
|
def __init__(self, kernel: Any) -> None:
|
|
self.kernel = kernel
|
|
self.count = 0
|
|
|
|
def __getitem__(self, grid: Any) -> Any:
|
|
self.count += 1
|
|
return self.kernel[grid]
|
|
|
|
|
|
def check_pair(
|
|
torch: Any,
|
|
reference: Callable,
|
|
fused: Callable,
|
|
module: ModuleType,
|
|
kernels: tuple[str, ...],
|
|
*,
|
|
rtol: float,
|
|
atol: float,
|
|
) -> dict:
|
|
cuda = getattr(torch, "cuda")
|
|
expected = reference()
|
|
probes = {name: _LaunchProbe(getattr(module, name)) for name in kernels}
|
|
try:
|
|
for name, probe in probes.items():
|
|
setattr(module, name, probe)
|
|
actual = fused()
|
|
cuda.synchronize()
|
|
finally:
|
|
for name, probe in probes.items():
|
|
setattr(module, name, probe.kernel)
|
|
if any(probe.count != 1 for probe in probes.values()):
|
|
raise RuntimeError(f"Expected one launch per fused kernel; got {[(k, p.count) for k, p in probes.items()]}")
|
|
expected = expected if isinstance(expected, tuple) else (expected,)
|
|
actual = actual if isinstance(actual, tuple) else (actual,)
|
|
errors = []
|
|
for result, baseline in zip(actual, expected, strict=True):
|
|
for value in (baseline, result):
|
|
if not torch.isfinite(value).all().item():
|
|
raise AssertionError("Nonfinite benchmark output")
|
|
if value.numel() > 1 and (value == value.flatten()[0]).all().item():
|
|
raise AssertionError("Constant benchmark output")
|
|
torch.testing.assert_close(result, baseline, rtol=rtol, atol=atol)
|
|
difference = (result.float() - baseline.float()).abs()
|
|
errors.append(
|
|
{
|
|
"max_abs": difference.max().item(),
|
|
"max_rel": (difference / baseline.float().abs().clamp_min(1e-6)).max().item(),
|
|
}
|
|
)
|
|
return {"parity": "pass", "errors": errors, "rtol": rtol, "atol": atol, "kernels": list(kernels)}
|
|
|
|
|
|
def _percentile(samples: list[float], quantile: float) -> float:
|
|
ordered = sorted(samples)
|
|
position = (len(ordered) - 1) * quantile
|
|
lower, upper = math.floor(position), math.ceil(position)
|
|
return ordered[lower] + (ordered[upper] - ordered[lower]) * (position - lower)
|
|
|
|
|
|
def measure_pair(torch: Any, reference: Callable, fused: Callable, args: argparse.Namespace) -> dict:
|
|
cuda = getattr(torch, "cuda")
|
|
functions = (reference, fused)
|
|
order: list[int] = []
|
|
|
|
def alternate() -> Any:
|
|
mode = len(order) % 2
|
|
order.append(mode)
|
|
return functions[mode]()
|
|
|
|
times_ms = []
|
|
for _ in range(args.warmup):
|
|
alternate()
|
|
cuda.synchronize()
|
|
order.clear()
|
|
flush = torch.empty(512 * 1024 * 1024, dtype=torch.uint8, device="cuda")
|
|
for _ in range(2 * args.num_tests):
|
|
flush.zero_()
|
|
cuda.synchronize()
|
|
start = cuda.Event(enable_timing=True)
|
|
end = cuda.Event(enable_timing=True)
|
|
start.record()
|
|
alternate()
|
|
end.record()
|
|
end.synchronize()
|
|
times_ms.append(start.elapsed_time(end))
|
|
samples: tuple[list[float], list[float]] = ([], [])
|
|
for mode, milliseconds in zip(order[-len(times_ms) :], times_ms, strict=True):
|
|
if not math.isfinite(milliseconds) or milliseconds <= 0:
|
|
raise RuntimeError(f"Invalid GPU event sample: {milliseconds}")
|
|
samples[mode].append(milliseconds * 1000)
|
|
stats = {}
|
|
for label, values in zip(("native", "fused"), samples, strict=True):
|
|
if len(values) != args.num_tests:
|
|
raise RuntimeError("GPU timing harness returned an unexpected sample count")
|
|
p50 = _percentile(values, 0.5)
|
|
stats[label] = {"median_us": p50, "p50_us": p50, "p90_us": _percentile(values, 0.9), "samples": len(values)}
|
|
stats["speedup_p50"] = stats["native"]["p50_us"] / stats["fused"]["p50_us"]
|
|
return stats
|