()
| 23 | |
| 24 | from ..bench.speed_test import run_speed_test |
| 25 | from ..bench.triton_benchmark_lib import Args |
| 26 | |
| 27 | set_volume_caches() |
| 28 | args = Args( |
| 29 | n_hidden_states=n_hidden_states, |
| 30 | n_procs=n_procs, |
| 31 | name=name, |
| 32 | bench_fn=bench_fn, |
| 33 | case=case, |
| 34 | tgt_dir=Path(volume_path) / "speed-test", |
| 35 | ) |
| 36 | run_speed_test(args) |
| 37 |