311 lines
13 KiB
Python
311 lines
13 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Compare MAGI-2 routed MoE with equal preprocessing boundaries.
|
|
|
|
python -m benchmarks.kernels.benchmark_magi2_bf16_moe --tokens 2 4096 10000
|
|
|
|
Default (`--mode dispatch`): the same MAGI-2 layer is timed with its native
|
|
reference and default BF16 fast path, including routing and all preprocessing.
|
|
With `--mode routed`, routing probabilities/IDs are inputs to BOTH timed paths. Sorting,
|
|
padding, packing, allocations and reduction remain inside their respective
|
|
calls. The reference uses deterministic routed storage in diagnostic mode.
|
|
Use --mode dispatch to measure the
|
|
actual layer with the native and default fast paths, including the router. Unsupported
|
|
dispatch fails explicitly; it is never silently replaced by another baseline.
|
|
|
|
Stage timings are standalone diagnostics, not additive or env speedup claims.
|
|
These synthetic local shapes are not full-model performance measurements.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import importlib
|
|
import json
|
|
import math
|
|
import os
|
|
import statistics
|
|
import subprocess
|
|
import time
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
|
|
def parse_args(argv=None):
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--tokens", type=int, nargs="+", default=[2, 4096])
|
|
parser.add_argument("--heads", type=int, default=3, help="EP4 local head count for MAGI-2 Preview.")
|
|
parser.add_argument("--experts", type=int, default=256, help="Released experts per head.")
|
|
parser.add_argument("--top-k", type=int, default=6)
|
|
parser.add_argument("--warmup", type=int, default=10)
|
|
parser.add_argument("--iters", type=int, default=50)
|
|
parser.add_argument("--seed", type=int, default=419)
|
|
parser.add_argument("--mode", choices=["routed", "dispatch"], default="dispatch")
|
|
parser.add_argument(
|
|
"--native-deterministic",
|
|
action="store_true",
|
|
help="Use the portable deterministic native reference (needed by Triton-MUSA 3.2).",
|
|
)
|
|
parser.add_argument("--stages", action="store_true")
|
|
parser.add_argument("--output", type=Path)
|
|
args = parser.parse_args(argv)
|
|
if min(*args.tokens, args.heads, args.experts, args.top_k, args.iters) < 1 or args.warmup < 0:
|
|
parser.error("Shapes/iters must be positive; warmup must be nonnegative")
|
|
if args.top_k > args.experts:
|
|
parser.error("top-k must not exceed experts per head")
|
|
return args
|
|
|
|
|
|
@contextmanager
|
|
def env_switch(enabled):
|
|
# Keep the legacy benchmark context API; dispatch is no longer opt-in.
|
|
with patch.dict(os.environ, MAGI2_DETERMINISTIC="0"):
|
|
yield
|
|
|
|
|
|
def stats(values):
|
|
if not values or any(not math.isfinite(x) or x <= 0 for x in values):
|
|
raise ValueError("Missing/non-positive/non-finite timing sample")
|
|
ordered = sorted(values)
|
|
return {
|
|
"mean_us": statistics.mean(values),
|
|
"p50_us": statistics.median(values),
|
|
"p90_us": ordered[math.ceil(0.9 * len(ordered)) - 1],
|
|
"samples_us": values,
|
|
}
|
|
|
|
|
|
def measure(torch, calls, args):
|
|
"""Alternating AB/BA; cold-cache flush is outside each timed region."""
|
|
cuda = getattr(torch, "cuda")
|
|
names = list(calls)
|
|
sequence = names + names[::-1]
|
|
wall_samples = []
|
|
labels = []
|
|
counter = 0
|
|
|
|
def invoke():
|
|
nonlocal counter
|
|
name = sequence[counter % len(sequence)]
|
|
counter += 1
|
|
cuda.synchronize()
|
|
start = time.perf_counter_ns()
|
|
output = calls[name]()
|
|
cuda.synchronize()
|
|
elapsed = (time.perf_counter_ns() - start) / 1000
|
|
wall_samples.append(elapsed)
|
|
labels.append(name)
|
|
return output
|
|
|
|
count = len(names) * args.iters
|
|
cache = torch.empty(512 * 1024 * 1024, device="cuda", dtype=torch.uint8)
|
|
gpu_ms = []
|
|
for index in range(len(names) * args.warmup + count):
|
|
cache.zero_()
|
|
cuda.synchronize()
|
|
begin, end = cuda.Event(enable_timing=True), cuda.Event(enable_timing=True)
|
|
begin.record()
|
|
invoke()
|
|
end.record()
|
|
end.synchronize()
|
|
if index >= len(names) * args.warmup:
|
|
gpu_ms.append(begin.elapsed_time(end))
|
|
labels, wall_samples = labels[-count:], wall_samples[-count:]
|
|
result = {}
|
|
for name in names:
|
|
wall = [x for label, x in zip(labels, wall_samples) if label == name]
|
|
event = [float(x) * 1000 for label, x in zip(labels, gpu_ms) if label == name]
|
|
if len(wall) != args.iters or len(event) != args.iters:
|
|
raise RuntimeError("Timing harness did not return the requested sample counts")
|
|
result[name] = {"synchronized_wall": stats(wall), "gpu_event": stats(event)}
|
|
return result
|
|
|
|
|
|
def make_inputs(torch, tokens, args):
|
|
# Fan-in-scaled synthetic weights. Unlike the old script, W_up and W_gate
|
|
# have the same distribution and top-k IDs are unique within each token.
|
|
torch.manual_seed(args.seed)
|
|
x = torch.randn(tokens, args.heads, 256, device="cuda").to(torch.bfloat16)
|
|
shape = (args.heads * args.experts, 256, 1280)
|
|
gate, up = [(torch.randn(shape, device="cuda") / 16).to(torch.bfloat16) for _ in range(2)]
|
|
down = (torch.randn(shape[0], 1280, 256, device="cuda") / math.sqrt(1280)).to(torch.bfloat16)
|
|
logits = torch.randn(args.heads, tokens, args.experts, device="cuda")
|
|
selected = logits.topk(args.top_k, dim=-1)
|
|
probabilities = selected.values.softmax(-1)
|
|
return x, probabilities, selected.indices, gate, up, down
|
|
|
|
|
|
def check_and_capture(torch, moe, kernels, data, args):
|
|
x, probabilities, indices, gate, up, down = data
|
|
gather, weights, offsets = moe.global_sort_routes(probabilities, indices, args.experts)
|
|
sorted_routes, launches = [], []
|
|
scatter, invoke = kernels._deterministic_scatter, moe.invoke_fused_moe_bf16
|
|
|
|
def capture_scatter(routes, *params):
|
|
sorted_routes.append(routes.detach().clone())
|
|
return scatter(routes, *params)
|
|
|
|
def capture_launch(*params, **kwargs):
|
|
launches.append((params, kwargs))
|
|
return invoke(*params, **kwargs)
|
|
|
|
with patch.object(kernels, "_deterministic_scatter", capture_scatter):
|
|
moe.triton_mh_moe_forward(x, gather, weights, offsets, gate, up, down, deterministic=True)
|
|
packed_w13 = moe._pack_bf16_w13(gate, up)
|
|
route_buffers = moe._allocate_bf16_route_buffers(
|
|
x.shape[0] * args.heads * args.top_k,
|
|
args.heads * args.experts,
|
|
128,
|
|
x.device,
|
|
)
|
|
with patch.object(moe, "invoke_fused_moe_bf16", capture_launch):
|
|
actual = moe._bf16_fused_moe_forward(x, probabilities, indices, packed_w13, down, route_buffers)
|
|
if len(sorted_routes) != 1 or len(launches) != 2:
|
|
raise RuntimeError("Unexpected kernel dispatch: benchmark contract needs updating")
|
|
flat_experts = (indices + torch.arange(args.heads, device=x.device)[:, None, None] * args.experts).flatten()
|
|
order = flat_experts.argsort(stable=True)
|
|
new_routes = launches[1][0][2].reshape(-1, 256).index_select(0, order)
|
|
scatter_ids = gather.long() * args.heads + flat_experts[order] // args.experts
|
|
expected = torch.zeros(x.shape[0] * args.heads, 256, device=x.device, dtype=torch.float32)
|
|
expected.index_add_(0, scatter_ids, sorted_routes[0].float())
|
|
expected = expected.to(x.dtype).view_as(x)
|
|
errors = []
|
|
# Same expert/FP32 reduction contract as test_bf16_moe_wiring.py, rather
|
|
# than relaxing tolerances on the inherited BF16 scatter reduction.
|
|
for result, reference in ((new_routes, sorted_routes[0]), (actual, expected)):
|
|
if not (torch.isfinite(result).all() and torch.isfinite(reference).all()):
|
|
raise RuntimeError("Non-finite MoE output")
|
|
torch.testing.assert_close(result, reference, rtol=2e-2, atol=0.5)
|
|
relative_l2 = (result.float() - reference.float()).norm() / reference.float().norm().clamp_min(1e-6)
|
|
if relative_l2.item() >= 1e-3:
|
|
raise AssertionError(f"MoE relative L2 {relative_l2.item()} >= 1e-3")
|
|
errors.append(
|
|
{"max_abs": (result.float() - reference.float()).abs().max().item(), "relative_l2": relative_l2.item()}
|
|
)
|
|
return launches, errors
|
|
|
|
|
|
def benchmark(torch, moe, kernels, tokens, args):
|
|
data = make_inputs(torch, tokens, args)
|
|
x, probabilities, indices, gate, up, down = data
|
|
launches, errors = check_and_capture(torch, moe, kernels, data, args)
|
|
packed_w13 = moe._pack_bf16_w13(gate, up)
|
|
route_count = tokens * args.heads * args.top_k
|
|
route_buffers = moe._allocate_bf16_route_buffers(route_count, args.heads * args.experts, 128, x.device)
|
|
|
|
def reference():
|
|
gather, weights, offsets = moe.global_sort_routes(probabilities, indices, args.experts)
|
|
return moe.triton_mh_moe_forward(x, gather, weights, offsets, gate, up, down, deterministic=True)
|
|
|
|
calls = {
|
|
"reference_deterministic_routed": reference,
|
|
"candidate_routed": lambda: moe._bf16_fused_moe_forward(
|
|
x, probabilities, indices, packed_w13, down, route_buffers
|
|
),
|
|
}
|
|
if args.mode == "dispatch":
|
|
config = moe.Magi2MultiHeadMoEConfig(
|
|
args.heads * 256, args.heads, args.experts, args.top_k, 1280, torch.bfloat16
|
|
)
|
|
layer = moe.Magi2MultiHeadMoE(config).to(x.device)
|
|
for parameter, value in ((layer.W_gate, gate), (layer.W_up, up), (layer.W_down, down)):
|
|
parameter.copy_(value)
|
|
layer.gate.normal_(std=1 / 16)
|
|
|
|
flat_x = x.reshape(tokens, -1)
|
|
|
|
def native_dispatch():
|
|
x_heads = flat_x.view(tokens, args.heads, 256)
|
|
probs, idx = layer._route(x_heads)
|
|
gather, weights, offsets = moe.global_sort_routes(probs, idx, args.experts)
|
|
return moe.triton_mh_moe_forward(
|
|
x_heads,
|
|
gather,
|
|
weights,
|
|
offsets,
|
|
layer.W_gate,
|
|
layer.W_up,
|
|
layer.W_down,
|
|
deterministic=args.native_deterministic,
|
|
).reshape(tokens, -1)
|
|
|
|
calls = {"native_dispatch": native_dispatch, "fused_dispatch": lambda: layer(flat_x)}
|
|
torch.testing.assert_close(calls["fused_dispatch"](), calls["native_dispatch"](), rtol=2e-2, atol=0.5)
|
|
timings = measure(torch, calls, args)
|
|
names = list(calls)
|
|
result = {"status": "pass", "tokens": tokens, "mode": args.mode, "correctness": errors, "timings": timings}
|
|
result["speedup_p50_wall"] = (
|
|
timings[names[0]]["synchronized_wall"]["p50_us"] / timings[names[1]]["synchronized_wall"]["p50_us"]
|
|
)
|
|
if args.stages:
|
|
route_ids = (indices.int() + torch.arange(args.heads, device=x.device)[:, None, None] * args.experts).reshape(
|
|
-1, args.top_k
|
|
)
|
|
|
|
def gems():
|
|
for params, kwargs in launches:
|
|
moe.invoke_fused_moe_bf16(*params, **kwargs)
|
|
|
|
result["standalone_stages_not_additive"] = measure(
|
|
torch,
|
|
{
|
|
"reference_sort": lambda: moe.global_sort_routes(probabilities, indices, args.experts),
|
|
"candidate_align": lambda: moe._align_bf16_routes(route_ids, args.heads * args.experts, 128),
|
|
"candidate_pack_w13": lambda: torch.stack((gate.transpose(1, 2), up.transpose(1, 2)), dim=2).reshape(
|
|
args.heads * args.experts, 2560, 256
|
|
),
|
|
"candidate_two_gemms_prepared": gems,
|
|
},
|
|
args,
|
|
)
|
|
return result
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
import torch
|
|
|
|
if hasattr(torch.version, "musa") and torch.version.musa is not None:
|
|
importlib.import_module("torchada") # before importing vLLM / capturing CUDA APIs
|
|
cuda = getattr(torch, "cuda")
|
|
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:
|
|
print(json.dumps({"status": "skip", "reason": "No visible CUDA/MUSA device"}))
|
|
return
|
|
import vllm_omni.diffusion.models.magi2.fused_moe_kernels as kernels
|
|
import vllm_omni.diffusion.models.magi2.mh_moe as moe
|
|
|
|
root = Path(__file__).resolve().parents[2]
|
|
try:
|
|
sha = subprocess.check_output(
|
|
["git", "-C", str(root), "rev-parse", "HEAD"], text=True, stderr=subprocess.DEVNULL
|
|
).strip()
|
|
except (OSError, subprocess.CalledProcessError):
|
|
sha = os.environ.get("MAGI2_BENCH_SOURCE_SHA", "unknown")
|
|
report = {
|
|
"device": cuda.get_device_name(),
|
|
"sm_count": cuda.get_device_properties(0).multi_processor_count,
|
|
"source_sha": sha,
|
|
"module_path": moe.__file__,
|
|
"torch": torch.__version__,
|
|
"args": {**vars(args), "output": str(args.output)},
|
|
"cache": "512 MiB flush outside each region",
|
|
"capture": "eager; JIT warmed; no graph/compile",
|
|
"results": [],
|
|
}
|
|
with torch.inference_mode():
|
|
for tokens in args.tokens:
|
|
result = benchmark(torch, moe, kernels, tokens, args)
|
|
report["results"].append(result)
|
|
print(json.dumps(result), flush=True)
|
|
if args.output:
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output.write_text(json.dumps(report, indent=2) + "\n")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|