基准测试支持随机种子控制

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')
+                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')

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()