(func, *args, **kwargs)
| 313 | ###################################################################### |
| 314 | # Check correctness and performance |
| 315 | def benchmark(func, *args, **kwargs): |
| 316 | torch.cuda.synchronize() |
| 317 | begin = timeit.default_timer() |
| 318 | output = func(*args, **kwargs) |
| 319 | torch.cuda.synchronize() |
| 320 | end = timeit.default_timer() |
| 321 | return output, (end - begin) |
| 322 | |
| 323 | output_nested, time_nested = benchmark(mha, query, key, value) |
| 324 | output_padded, time_padded = benchmark(mha, padded_query, padded_key, padded_value) |