diff --git a/benchmark/bench_flash_mla.py b/benchmark/bench_flash_mla.py index 2b59f8ed..f9ce5eda 100644 --- a/benchmark/bench_flash_mla.py +++ b/benchmark/bench_flash_mla.py @@ -407,14 +407,18 @@ def flash_mla_triton(): "flash_mla_triton": run_flash_mla_triton, } -def compare_ab(baseline, target, b, s_q, cache_seqlens, h_q, h_kv, d, dv, causal, dtype): +def set_benchmark_seed(seed): + torch.manual_seed(seed) + random.seed(seed) + + +def compare_ab(baseline, target, b, s_q, cache_seqlens, h_q, h_kv, d, dv, causal, dtype, seed=0): print(f"comparing {baseline} vs {target}: {b=}, {s_q=}, mean_seqlens={cache_seqlens.float().mean()}, {h_q=}, {h_kv=}, {d=}, {dv=}, {causal=}, {dtype=}") device = torch.device("cuda:0") torch.set_default_dtype(dtype) torch.set_default_device(device) torch.cuda.set_device(device) - torch.manual_seed(0) - random.seed(0) + set_benchmark_seed(seed) assert baseline in FUNC_TABLE assert target in FUNC_TABLE baseline_func = FUNC_TABLE[baseline] @@ -447,14 +451,13 @@ def compare_ab(baseline, target, b, s_q, cache_seqlens, h_q, h_kv, d, dv, causal return bytes / 10 ** 6 / perf_a, bytes / 10 ** 6 / perf_b -def compare_a(target, b, s_q, cache_seqlens, h_q, h_kv, d, dv, causal, dtype): +def compare_a(target, b, s_q, cache_seqlens, h_q, h_kv, d, dv, causal, dtype, seed=0): print(f"{target}: {b=}, {s_q=}, mean_seqlens={cache_seqlens.float().mean()}, {h_q=}, {h_kv=}, {d=}, {dv=}, {causal=}, {dtype=}") torch.set_default_dtype(dtype) device = torch.device("cuda:0") torch.set_default_device(device) torch.cuda.set_device(device) - torch.manual_seed(0) - random.seed(0) + set_benchmark_seed(seed) assert target in FUNC_TABLE target_func = FUNC_TABLE[target] @@ -497,6 +500,7 @@ def get_args(): parser.add_argument("--all", action="store_true") parser.add_argument("--one", action="store_true") parser.add_argument("--compare", action="store_true") + parser.add_argument("--seed", type=int, default=0) args = parser.parse_args() return args @@ -509,12 +513,12 @@ def get_args(): for shape in shape_configs: if args.all: for target in available_targets: - perf = compare_a(target, shape["b"], shape["s_q"], shape["cache_seqlens"], shape["h_q"], shape["h_kv"], shape["d"], shape["dv"], shape["causal"], shape["dtype"]) + perf = compare_a(target, shape["b"], shape["s_q"], shape["cache_seqlens"], shape["h_q"], shape["h_kv"], shape["d"], shape["dv"], shape["causal"], shape["dtype"], seed=args.seed) fout.write(f'{target},{shape["b"]},{shape["cache_seqlens"].float().mean().cpu().item():.0f},{shape["h_q"]},{perf:.0f}\n') elif args.compare: - perfa, prefb = compare_ab(args.baseline, args.target, shape["b"], shape["s_q"], shape["cache_seqlens"], shape["h_q"], shape["h_kv"], shape["d"], shape["dv"], shape["causal"], shape["dtype"]) + perfa, prefb = compare_ab(args.baseline, args.target, shape["b"], shape["s_q"], shape["cache_seqlens"], shape["h_q"], shape["h_kv"], shape["d"], shape["dv"], shape["causal"], shape["dtype"], seed=args.seed) fout.write(f'{args.baseline},{shape["b"]},{shape["cache_seqlens"].float().mean().cpu().item():.0f},{shape["h_q"]},{perfa:.0f}\n') fout.write(f'{args.target},{shape["b"]},{shape["cache_seqlens"].float().mean().cpu().item():.0f},{shape["h_q"]},{prefb:.0f}\n') elif args.one: - perf = compare_a(args.target, shape["b"], shape["s_q"], shape["cache_seqlens"], shape["h_q"], shape["h_kv"], shape["d"], shape["dv"], shape["causal"], shape["dtype"]) - fout.write(f'{args.target},{shape["b"]},{shape["cache_seqlens"].float().mean().cpu().item():.0f},{shape["h_q"]},{perf:.0f}\n') \ No newline at end of file + perf = compare_a(args.target, shape["b"], shape["s_q"], shape["cache_seqlens"], shape["h_q"], shape["h_kv"], shape["d"], shape["dv"], shape["causal"], shape["dtype"], seed=args.seed) + fout.write(f'{args.target},{shape["b"]},{shape["cache_seqlens"].float().mean().cpu().item():.0f},{shape["h_q"]},{perf:.0f}\n') diff --git a/tests/test_benchmark_seed.py b/tests/test_benchmark_seed.py new file mode 100644 index 00000000..306a3acf --- /dev/null +++ b/tests/test_benchmark_seed.py @@ -0,0 +1,66 @@ +import importlib.util +import random +import sys +import types +from pathlib import Path +from unittest.mock import Mock, patch + + +def _identity_decorator(*args, **kwargs): + def decorate(func): + return func + + return decorate + + +def _load_benchmark_module(): + fake_torch = types.ModuleType("torch") + fake_torch.manual_seed = Mock() + fake_torch.inference_mode = _identity_decorator + fake_torch.tensor = lambda *args, **kwargs: list(args[0]) if args else [] + fake_torch.int32 = object() + fake_torch.bfloat16 = object() + + fake_triton = types.ModuleType("triton") + fake_triton.autotune = _identity_decorator + fake_triton.heuristics = _identity_decorator + fake_triton.jit = _identity_decorator + fake_triton.Config = lambda *args, **kwargs: (args, kwargs) + fake_triton.testing = types.SimpleNamespace(do_bench=lambda func: 0) + + fake_tl = types.ModuleType("triton.language") + fake_tl.__getattr__ = lambda name: object() + fake_flash_mla = types.ModuleType("flash_mla") + fake_flash_mla.flash_mla_with_kvcache = Mock() + fake_flash_mla.get_mla_metadata = Mock() + + module_name = "bench_flash_mla_test_module" + source_path = Path(__file__).resolve().parents[1] / "benchmark" / "bench_flash_mla.py" + spec = importlib.util.spec_from_file_location(module_name, source_path) + module = importlib.util.module_from_spec(spec) + stubbed_modules = { + "flashinfer": types.ModuleType("flashinfer"), + "flash_mla": fake_flash_mla, + "torch": fake_torch, + "triton": fake_triton, + "triton.language": fake_tl, + } + with patch.dict(sys.modules, stubbed_modules, clear=False): + sys.modules.pop(module_name, None) + assert spec.loader is not None + spec.loader.exec_module(module) + return module, fake_torch + + +def test_set_benchmark_seed_updates_torch_and_python_random(): + module, fake_torch = _load_benchmark_module() + module.set_benchmark_seed(123) + + fake_torch.manual_seed.assert_called_once_with(123) + first = random.random() + random.seed(123) + assert first == random.random() + + +if __name__ == "__main__": + test_set_benchmark_seed_updates_torch_and_python_random()