91 lines
3.0 KiB
Python
91 lines
3.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""CPU contracts for the benchmark, independent of GPU availability."""
|
|
|
|
import os
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from benchmarks.kernels import benchmark_magi2_bf16_moe as bench
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def test_shape_sweep_and_cli_validation():
|
|
assert bench.parse_args(["--tokens", "2", "10000"]).tokens == [2, 10000]
|
|
for args in (["--iters", "0"], ["--warmup", "-1"], ["--top-k", "257"], ["--tokens", "0"]):
|
|
with pytest.raises(SystemExit):
|
|
bench.parse_args(args)
|
|
|
|
|
|
@pytest.mark.parametrize("values", [[], [0], [-1], [float("nan")], [float("inf")]])
|
|
def test_invalid_timings_are_not_reported(values):
|
|
with pytest.raises(ValueError):
|
|
bench.stats(values)
|
|
|
|
|
|
def test_stats_preserve_all_samples():
|
|
result = bench.stats([1.0, 4.0, 2.0, 3.0])
|
|
assert result["p50_us"] == 2.5
|
|
assert result["p90_us"] == 4.0
|
|
assert result["samples_us"] == [1.0, 4.0, 2.0, 3.0]
|
|
|
|
|
|
def test_switch_restores_environment_on_error(monkeypatch):
|
|
monkeypatch.setenv("MAGI2_DETERMINISTIC", "1")
|
|
with pytest.raises(RuntimeError), bench.env_switch(True):
|
|
assert os.environ["MAGI2_DETERMINISTIC"] == "0"
|
|
raise RuntimeError("compile failure")
|
|
assert os.environ["MAGI2_DETERMINISTIC"] == "1"
|
|
|
|
|
|
def test_routed_reference_sort_is_inside_every_timed_call(monkeypatch):
|
|
data = (SimpleNamespace(device="device"), *(object() for _ in range(5)))
|
|
events = []
|
|
monkeypatch.setattr(bench, "make_inputs", lambda *args: data)
|
|
monkeypatch.setattr(bench, "check_and_capture", lambda *args: ([], []))
|
|
|
|
def sort(*args):
|
|
assert args == (data[1], data[2], 256)
|
|
events.append("sort")
|
|
return (1, 2, 3)
|
|
|
|
def legacy(*args, **kwargs):
|
|
assert kwargs == {"deterministic": True}
|
|
events.append("reference")
|
|
|
|
def candidate(*args):
|
|
assert args == (data[0], data[1], data[2], "packed-w13", data[5], "route-buffers")
|
|
events.append("candidate")
|
|
|
|
def measure(torch, calls, args):
|
|
assert events == ["pack", "allocate"]
|
|
events.clear()
|
|
for _ in range(2):
|
|
for call in calls.values():
|
|
call()
|
|
return {name: {"synchronized_wall": {"p50_us": 1}} for name in calls}
|
|
|
|
monkeypatch.setattr(bench, "measure", measure)
|
|
|
|
def pack_w13(*args):
|
|
events.append("pack")
|
|
return "packed-w13"
|
|
|
|
def allocate_route_buffers(*args):
|
|
events.append("allocate")
|
|
return "route-buffers"
|
|
|
|
moe = SimpleNamespace(
|
|
global_sort_routes=sort,
|
|
triton_mh_moe_forward=legacy,
|
|
_bf16_fused_moe_forward=candidate,
|
|
_pack_bf16_w13=pack_w13,
|
|
_allocate_bf16_route_buffers=allocate_route_buffers,
|
|
)
|
|
result = bench.benchmark(None, moe, None, 2, bench.parse_args(["--mode", "routed"]))
|
|
assert events == ["sort", "reference", "candidate"] * 2
|
|
assert result["mode"] == "routed"
|
|
assert "env_off" not in str(result["timings"])
|